src/storage/encryption.py - develop - dsboikov/aiBoardRoom

src/storage/encryption.py
113 строк4 KB
Денис
Project updated
03 авг 2026, 12:39
03 авг 2026, 12:39d7da468
"""Fernet encryption with master key stored in OS keyring."""

from __future__ import annotations

import base64
import json
from typing import Any

import keyring
from cryptography.fernet import Fernet, InvalidToken

from config import KEYRING_MASTER_KEY, KEYRING_SERVICE
from utils.logger import get_logger

logger = get_logger("encryption")


class EncryptionError(Exception):
    """Raised when encryption or decryption fails."""


def _get_or_create_master_key() -> bytes:
    """Load Fernet master key from OS keyring, creating one if missing."""
    existing = keyring.get_password(KEYRING_SERVICE, KEYRING_MASTER_KEY)
    if existing:
        return existing.encode("utf-8")
    key = Fernet.generate_key()
    keyring.set_password(KEYRING_SERVICE, KEYRING_MASTER_KEY, key.decode("utf-8"))
    logger.info("Generated and stored new Fernet master key in OS keyring")
    return key


class SecretStore:
    """Encrypt secrets with Fernet; persist API keys in OS keyring."""

    def __init__(self) -> None:
        self._fernet = Fernet(_get_or_create_master_key())

    def encrypt(self, plaintext: str) -> str:
        """Encrypt a string and return a URL-safe base64 token."""
        try:
            token = self._fernet.encrypt(plaintext.encode("utf-8"))
            return token.decode("utf-8")
        except Exception as exc:  # noqa: BLE001
            raise EncryptionError(f"Encryption failed: {exc}") from exc

    def decrypt(self, token: str) -> str:
        """Decrypt a Fernet token back to plaintext."""
        try:
            return self._fernet.decrypt(token.encode("utf-8")).decode("utf-8")
        except InvalidToken as exc:
            raise EncryptionError("Invalid encryption token") from exc
        except Exception as exc:  # noqa: BLE001
            raise EncryptionError(f"Decryption failed: {exc}") from exc

    def encrypt_json(self, data: Any) -> str:
        """Serialize and encrypt a JSON-compatible object."""
        payload = json.dumps(data, ensure_ascii=False, separators=(",", ":"))
        return self.encrypt(payload)

    def decrypt_json(self, token: str) -> Any:
        """Decrypt and parse a JSON payload."""
        return json.loads(self.decrypt(token))

    def set_api_key(self, provider: str, api_key: str) -> None:
        """Store an API key in the OS credential store (encrypted)."""
        encrypted = self.encrypt(api_key)
        keyring.set_password(KEYRING_SERVICE, f"api-key:{provider}", encrypted)
        logger.info("Stored API key for provider=%s", provider)

    def get_api_key(self, provider: str) -> str | None:
        """Retrieve and decrypt an API key, or None if missing."""
        encrypted = keyring.get_password(KEYRING_SERVICE, f"api-key:{provider}")
        if not encrypted:
            return None
        try:
            return self.decrypt(encrypted)
        except EncryptionError:
            # Legacy: plain key stored before encryption was introduced
            logger.warning("Falling back to plaintext keyring value for %s", provider)
            return encrypted

    def delete_api_key(self, provider: str) -> None:
        """Remove an API key from the credential store."""
        try:
            keyring.delete_password(KEYRING_SERVICE, f"api-key:{provider}")
        except keyring.errors.PasswordDeleteError:
            pass

    def has_api_key(self, provider: str) -> bool:
        """Return True if a key exists for the provider."""
        return bool(keyring.get_password(KEYRING_SERVICE, f"api-key:{provider}"))

    def export_blob(self, data: Any) -> bytes:
        """Encrypt arbitrary JSON data into binary for file export."""
        token = self.encrypt_json(data)
        return base64.urlsafe_b64encode(token.encode("utf-8"))

    def import_blob(self, blob: bytes) -> Any:
        """Decrypt an export blob back to a Python object."""
        token = base64.urlsafe_b64decode(blob).decode("utf-8")
        return self.decrypt_json(token)


_store: SecretStore | None = None


def get_secret_store() -> SecretStore:
    """Return a process-wide SecretStore singleton."""
    global _store
    if _store is None:
        _store = SecretStore()
    return _store