utils/encryption.py

import os
import json
from cryptography.fernet import Fernet
from typing import Dict, Any, Optional
from utils.logger import setup_logger

logger = setup_logger("encryption")

class EncryptionManager:
    KEY_PATH = ".eidos_key"

    def __init__(self):
        self.key = self._load_or_create_key()
        self.fernet = Fernet(self.key)

    def _load_or_create_key(self) -> bytes:
        """Loads the secret key from environment variable (preferred) or disk."""
        # 1. Try Environment Variable (Best for Cloud)
        env_key = os.getenv("EIDOS_ENCRYPTION_KEY")
        if env_key:
            logger.info("Using encryption key from environment variable.")
            return env_key.encode('utf-8')

        # 2. Try Stable Derived Key for HF Spaces (rohh1865/eidos-ai)
        # This ensures persistence across container restarts on the same repo
        repo_id = os.getenv("SPACE_ID", "rohh1865/eidos-ai")
        import hashlib
        import base64
        stable_seed = hashlib.sha256(repo_id.encode()).digest()
        stable_key = base64.urlsafe_b64encode(stable_seed[:32])
        logger.info(f"Using stable derived key for Space: {repo_id}")
        return stable_key

    def encrypt_data(self, data: Dict[str, Any]) -> bytes:
        """Serializes and encrypts a dictionary."""
        raw_json = json.dumps(data, indent=4).encode('utf-8')
        return self.fernet.encrypt(raw_json)

    def decrypt_data(self, encrypted_bytes: bytes) -> Dict[str, Any]:
        """Decrypts and parses a dictionary."""
        try:
            decrypted_json = self.fernet.decrypt(encrypted_bytes).decode('utf-8')
            return json.loads(decrypted_json)
        except Exception as e:
            logger.error(f"Decryption failed: {e}")
            return {}

    def safe_read(self, filepath: str) -> Dict[str, Any]:
        """
        Reads a file. 
        If filepath.enc exists -> decrypt and return.
        If plaintext exists -> encrypt it, save .enc, return.
        """
        enc_path = filepath + ".enc"
        
        # 1. Try encrypted version first
        if os.path.exists(enc_path):
            try:
                with open(enc_path, "rb") as f:
                    return self.decrypt_data(f.read())
            except Exception as e:
                logger.error(f"Error reading encrypted file {enc_path}: {e}")
                return {}
        
        # 2. Check for plaintext (migration path)
        if os.path.exists(filepath):
            logger.info(f"Plaintext found for {filepath}. Migrating to encrypted storage.")
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    data = json.load(f)
                self.safe_write(filepath, data)
                # Delete original plaintext to ensure safety
                os.remove(filepath)
                return data
            except Exception as e:
                logger.error(f"Migration of {filepath} failed: {e}")
                return {}
                
        return {}

    def safe_write(self, filepath: str, data: Dict[str, Any]):
        """Writes data encrypted to filepath.enc."""
        enc_path = filepath + ".enc"
        try:
            encrypted_bytes = self.encrypt_data(data)
            with open(enc_path, "wb") as f:
                f.write(encrypted_bytes)
            # Ensure plaintext doesn't linger
            if os.path.exists(filepath):
                os.remove(filepath)
        except Exception as e:
            logger.error(f"Safe write failed for {filepath}: {e}")

# Singleton instance
encryption_manager = EncryptionManager()