File size: 3,617 Bytes
ab616bf 5cce2f9 29c12d2 ab616bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | 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()
|