ss-tools/backend/src/core/encryption.py

# #region EncryptionCore [C:5] [TYPE Module] [SEMANTICS encryption, fernet, key, crypto, secret]
# @defgroup Core Module group.
# @BRIEF Core encryption infrastructure: Fernet-based EncryptionManager for reversible secret storage.
# @LAYER Infrastructure
# @RELATION DEPENDS_ON -> [LoggerModule]
# @INVARIANT Encryption initialization never falls back to a hardcoded secret.
#            ENCRYPTION_KEY must be set via environment or .env before first use.
# @PRE Runtime environment has a valid Fernet ENCRYPTION_KEY.
# @POST EncryptionManager provides encrypt/decrypt operations backed by validated Fernet key.
# @SIDE_EFFECT Reads ENCRYPTION_KEY from process environment.
# @DATA_CONTRACT Input[str] -> Encrypted[str] -> Output[str]
# @RATIONALE Extracted from services/llm_provider.py to decouple core encryption from LLM domain
#            and enable reuse by ConfigManager for Superset password encryption (see [SEC:C-1]).

import base64
import os

from cryptography.fernet import Fernet

from .logger import belief_scope, logger


# #region _require_fernet_key [C:5] [TYPE Function]
# @BRIEF Load and validate the Fernet key used for secret encryption.
# @PRE  ENCRYPTION_KEY environment variable must be set to a valid Fernet key.
# @POST  Returns validated key bytes ready for Fernet initialization.
# @RELATION DEPENDS_ON -> [LoggerModule]
# @SIDE_EFFECT  Emits belief-state logs for missing or invalid encryption configuration.
# @DATA_CONTRACT  Input[ENCRYPTION_KEY:str] -> Output[bytes]
# @INVARIANT  Encryption initialization never falls back to a hardcoded secret.
def _require_fernet_key() -> bytes:
    with belief_scope("_require_fernet_key"):
        raw_key = os.getenv("ENCRYPTION_KEY", "").strip()
        if not raw_key:
            logger.explore(
                "Missing ENCRYPTION_KEY blocks EncryptionManager initialization"
            )
            raise RuntimeError(
                "ENCRYPTION_KEY must be set. Run ensure_encryption_key() "
                "during application startup before using EncryptionManager."
            )

        key = raw_key.encode()
        try:
            Fernet(key)
        except Exception as exc:
            logger.explore(
                "Invalid ENCRYPTION_KEY blocks EncryptionManager initialization"
            )
            raise RuntimeError("ENCRYPTION_KEY must be a valid Fernet key") from exc

        logger.reflect("Validated ENCRYPTION_KEY for EncryptionManager initialization")
        return key


# #endregion _require_fernet_key


# #region EncryptionManager [C:5] [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Handles encryption and decryption of sensitive data like API keys and passwords.
# @RELATION CALLS -> [_require_fernet_key]
# @PRE  ENCRYPTION_KEY is configured with a valid Fernet key before instantiation.
# @POST  Manager exposes reversible encrypt/decrypt operations for persisted secrets.
# @SIDE_EFFECT  Initializes Fernet cryptography state from process environment.
# @DATA_CONTRACT  Input[str] -> Output[str]
# @INVARIANT  Uses only a validated secret key from environment.
# @RATIONALE Fernet key validation (Fernet(key)) is O(1) but repeated ~90× per translation
#            run — cached at module level via get_encryption_manager() singleton to eliminate
#            redundant env-var reads and key validation.
#
# @TEST_CONTRACT  EncryptionManagerModel ->
# {
#   required_fields: {},
#   invariants: [
#       "encrypted data can be decrypted back to the original string"
#   ]
# }
# @TEST_FIXTURE  basic_encryption_cycle -> {"data": "my_secret_key"}
# @TEST_EDGE  decrypt_invalid_data -> raises Exception
# @TEST_EDGE  empty_string_encryption -> {"data": ""}
# @TEST_INVARIANT  symmetric_encryption -> verifies: [basic_encryption_cycle, empty_string_encryption]
class EncryptionManager:
    # region EncryptionManager_init [TYPE Function]
    # @PURPOSE: Initialize the encryption manager with a Fernet key.
    # @PRE  ENCRYPTION_KEY env var must be set to a valid Fernet key.
    # @POST  Fernet instance ready for encryption/decryption.
    def __init__(self):
        self.key = _require_fernet_key()
        self.fernet = Fernet(self.key)

    # endregion EncryptionManager_init

    # region encrypt [TYPE Function]
    # @PURPOSE: Encrypt a plaintext string.
    # @PRE  data must be a non-empty string.
    # @POST  Returns encrypted string.
    def encrypt(self, data: str) -> str:
        with belief_scope("encrypt"):
            return self.fernet.encrypt(data.encode()).decode()

    # endregion encrypt

    # region decrypt [TYPE Function]
    # @PURPOSE: Decrypt an encrypted string.
    # @PRE  encrypted_data must be a valid Fernet-encrypted string.
    # @POST  Returns original plaintext string.
    def decrypt(self, encrypted_data: str) -> str:
        with belief_scope("decrypt"):
            return self.fernet.decrypt(encrypted_data.encode()).decode()

    # endregion decrypt


# #endregion EncryptionManager


# Module-level singleton — avoids redundant env reads and Fernet construction
# across ~90 LLMProviderService instantiations per translation run.
_encryption_manager: EncryptionManager | None = None


# #region is_fernet_token [C:2] [TYPE Function] [SEMANTICS encryption,validation]
# @BRIEF Check whether a string looks like a valid Fernet-encrypted token.
# @PRE  value is a string.
# @POST Returns True if value is base64-decodable with decoded length >= 32 bytes.
# @RATIONALE Differentiates Fernet-encrypted values from plaintext without decrypting.
#            Used by ConnectionService and ConfigManager to handle plaintext legacy data
#            vs. data corrupted by key mismatch.
# @REJECTED Trying decryption as detection was rejected — it logs noisy "Assumption
#           violated" warnings for every plaintext value and cannot distinguish between
#           "not encrypted" and "encrypted with a different key".
MIN_FERNET_TOKEN_LENGTH = 80


def is_fernet_token(value: str) -> bool:
    if not value or len(value) < MIN_FERNET_TOKEN_LENGTH:
        return False
    try:
        decoded = base64.urlsafe_b64decode(value.encode())
    except Exception:
        return False
    return len(decoded) >= 32
# #endregion is_fernet_token


# #region get_encryption_manager [C:4] [TYPE Function] [SEMANTICS encryption,singleton,cache]
# @BRIEF Return the process-wide EncryptionManager singleton.
# @PRE ENCRYPTION_KEY env var is set to a valid Fernet key.
# @POST Returns the same EncryptionManager instance on every call (idempotent).
# @SIDE_EFFECT Creates EncryptionManager on first call (reads env, validates Fernet key).
# @RATIONALE Eliminates redundant ENCRYPTION_KEY reads and Fernet key validation —
#            LLMProviderService was creating ~90 EncryptionManager instances per
#            translation run (2 per batch). The singleton reduces this to 1.
def get_encryption_manager() -> EncryptionManager:
    """Return the process-wide EncryptionManager singleton.

    Creates on first call; returns cached instance thereafter.
    Eliminates redundant ENCRYPTION_KEY reads and Fernet key validation.
    """
    global _encryption_manager
    if _encryption_manager is None:
        _encryption_manager = EncryptionManager()
    return _encryption_manager
# #endregion get_encryption_manager

# #endregion EncryptionCore