RMM/utils.py

from passlib.context import CryptContext
from cryptography.fernet import Fernet, InvalidToken
from pathlib import Path
import os


pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

def hash_password(password: str) -> str:
    return pwd_context.hash(password)

def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)


KEY_FILE = Path(__file__).resolve().parent / "encryption.key"


def load_key() -> bytes:
    env_key = os.getenv("ENCRYPTION_KEY")
    if env_key:
        return env_key.encode() if isinstance(env_key, str) else env_key

    if KEY_FILE.exists():
        return KEY_FILE.read_bytes().strip()

    key = Fernet.generate_key()
    KEY_FILE.write_bytes(key)
    return key


try:
    key = load_key()
    cipher = Fernet(key)
except Exception as exc:
    raise RuntimeError(
        "Invalid encryption key. Set ENCRYPTION_KEY or remove/replace encryption.key with a valid Fernet key."
    ) from exc


def encrypt_data(data: str) -> str:
    return cipher.encrypt(data.encode()).decode()

def decrypt_data(encrypted_data: str) -> str:
    try:
        return cipher.decrypt(encrypted_data.encode()).decode()
    except InvalidToken as exc:
        raise ValueError("Unable to decrypt data: invalid encryption token or mismatched key") from exc