# backend/services/usb_vault.py import os import sqlite3 import logging import base64 from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from backend.services.usb_monitor import get_db_path from enum import Enum class MissingDomainKeyError(Exception): pass class KeyDomain(str, Enum): CHAT = "chat" CLOUD = "cloud" GAMING_COACH = "gaming_coach" AUTO_UPGRADE_PC = "auto_upgrade_pc" AUTO_UPGRADE_CLOUD = "auto_upgrade_cloud" RESEARCH_ENGINE = "research_engine" CAPTCHA_SOLVER_PC = "captcha_solver_pc" CAPTCHA_SOLVER_CLOUD = "captcha_solver_cloud" IMAGE_GENERATION_PC = "image_generation_pc" IMAGE_GENERATION_CLOUD = "image_generation_cloud" OSINT_PROTOCOL_PC = "osint_protocol_pc" OSINT_PROTOCOL_CLOUD = "osint_protocol_cloud" SUPERVISOR_HEAL_PC = "supervisor_heal_pc" SUPERVISOR_HEAL_CLOUD = "supervisor_heal_cloud" HEAD_SUPERVISOR_PC = "head_supervisor_pc" HEAD_SUPERVISOR_CLOUD = "head_supervisor_cloud" CLOUD_SYNC_DOMAIN = "cloud_sync_domain" KEY_DOMAIN_ENV_MAP = { KeyDomain.CHAT: "GOOGLE_API_KEY_CHAT", KeyDomain.CLOUD: "GOOGLE_API_KEY_CLOUD", KeyDomain.GAMING_COACH: "GOOGLE_API_KEY_GAMING_COACH", KeyDomain.AUTO_UPGRADE_PC: "GOOGLE_API_KEY_AUTO_UPGRADE_PC", KeyDomain.AUTO_UPGRADE_CLOUD: "GOOGLE_API_KEY_AUTO_UPGRADE_CLOUD", KeyDomain.RESEARCH_ENGINE: "GOOGLE_API_KEY_RESEARCH_ENGINE", KeyDomain.CAPTCHA_SOLVER_PC: "GOOGLE_API_KEY_CAPTCHA_SOLVER_PC", KeyDomain.CAPTCHA_SOLVER_CLOUD: "GOOGLE_API_KEY_CAPTCHA_SOLVER_CLOUD", KeyDomain.IMAGE_GENERATION_PC: "GOOGLE_API_KEY_IMAGE_GENERATION_PC", KeyDomain.IMAGE_GENERATION_CLOUD: "GOOGLE_API_KEY_IMAGE_GENERATION_CLOUD", KeyDomain.OSINT_PROTOCOL_PC: "GOOGLE_API_KEY_OSINT_PROTOCOL_PC", KeyDomain.OSINT_PROTOCOL_CLOUD: "GOOGLE_API_KEY_OSINT_PROTOCOL_CLOUD", KeyDomain.SUPERVISOR_HEAL_PC: "GOOGLE_API_KEY_SUPERVISOR_HEAL_PC", KeyDomain.SUPERVISOR_HEAL_CLOUD: "GOOGLE_API_KEY_SUPERVISOR_HEAL_CLOUD", KeyDomain.HEAD_SUPERVISOR_PC: "GOOGLE_API_KEY_HEAD_SUPERVISOR_PC", KeyDomain.HEAD_SUPERVISOR_CLOUD: "GOOGLE_API_KEY_HEAD_SUPERVISOR_CLOUD", KeyDomain.CLOUD_SYNC_DOMAIN: "GOOGLE_API_KEY_CLOUD_SYNC_DOMAIN", } CLOUD_ONLY_DOMAINS = { KeyDomain.CLOUD, KeyDomain.AUTO_UPGRADE_CLOUD, KeyDomain.RESEARCH_ENGINE, KeyDomain.CAPTCHA_SOLVER_CLOUD, KeyDomain.IMAGE_GENERATION_CLOUD, KeyDomain.OSINT_PROTOCOL_CLOUD, KeyDomain.SUPERVISOR_HEAL_CLOUD, KeyDomain.HEAD_SUPERVISOR_CLOUD, KeyDomain.CLOUD_SYNC_DOMAIN } def generate_key_from_password(password: str) -> str: kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=b"JARVIS_OMEGA_SALT_V15", iterations=480000, ) return base64.urlsafe_b64encode(kdf.derive(password.encode())).decode('utf-8') def _init_vault_db() -> str: """Initializes the vault keys table and retrieves or generates the AES Master Key.""" try: db_path = get_db_path() with sqlite3.connect(db_path) as conn: conn.execute('PRAGMA journal_mode=WAL') conn.execute("CREATE TABLE IF NOT EXISTS vault_keys (id INTEGER PRIMARY KEY, key_data TEXT)") conn.execute("CREATE TABLE IF NOT EXISTS vault_secrets (key_name TEXT PRIMARY KEY, encrypted_value BLOB)") conn.execute("CREATE TABLE IF NOT EXISTS vault_credentials (id INTEGER PRIMARY KEY AUTOINCREMENT, site TEXT, username TEXT, password_enc BLOB, captured_at TEXT, source TEXT CHECK(source IN ('sentinel','manual','internet')))") cursor = conn.cursor() cursor.execute("SELECT key_data FROM vault_keys WHERE id = 1") row = cursor.fetchone() target_key = generate_key_from_password("always iron man") if not row: conn.execute("INSERT INTO vault_keys (id, key_data) VALUES (1, ?)", (target_key,)) conn.commit() return target_key elif row[0] != target_key: conn.execute("UPDATE vault_keys SET key_data = ? WHERE id = 1", (target_key,)) conn.commit() return target_key return row[0] except Exception as e: logging.error(f"Vault DB Error: {e}") return generate_key_from_password("always iron man") MASTER_KEY = _init_vault_db() cipher_suite = Fernet(MASTER_KEY.encode('utf-8')) def set_secret(key_name: str, plain_value: str): """Encrypts and stores a string secret in the memory.db vault.""" encrypted_value = cipher_suite.encrypt(plain_value.encode('utf-8')) db_path = get_db_path() with sqlite3.connect(db_path) as conn: conn.execute('PRAGMA journal_mode=WAL') conn.execute("INSERT OR REPLACE INTO vault_secrets (key_name, encrypted_value) VALUES (?, ?)", (key_name, encrypted_value)) conn.commit() def get_secret(key_name: str) -> str: """Retrieves and decrypts a string secret from the memory.db vault. Returns None if not found.""" # First check actual environment variables for cloud deployment env_val = os.environ.get(key_name) if env_val: return env_val db_path = get_db_path() try: with sqlite3.connect(db_path) as conn: conn.execute('PRAGMA journal_mode=WAL') cursor = conn.cursor() cursor.execute("SELECT encrypted_value FROM vault_secrets WHERE key_name = ?", (key_name,)) row = cursor.fetchone() if row and row[0]: return cipher_suite.decrypt(row[0]).decode('utf-8') except Exception as e: logging.error(f"Error decrypting secret {key_name}: {e}") return None def resolve_vault_key(domain: KeyDomain) -> str: """Dynamically resolves a key based on the environment and domain constraints.""" import os env_var_name = KEY_DOMAIN_ENV_MAP[domain] # HF Spaces use actual environment variables natively key = os.environ.get(env_var_name) if not key: # Fallback to local SQLite vault key = get_secret(env_var_name) # BACKWARD COMPATIBILITY / SINGLE-KEY FALLBACK: every domain here is a Google # (Gemini) key slot. If a caller hasn't provisioned the dedicated per-domain # key, fall back to the generic GEMINI_API_KEY rather than crashing the # feature. The dedicated key is always preferred (checked above) so the # multi-key rate-limit partitioning still applies whenever it's configured; # this only rescues the common single-key desktop user, for whom otherwise # only chat/cloud worked and gaming-coach / image-gen / self-heal / etc. # raised MissingDomainKeyError. if not key: key = os.environ.get("GEMINI_API_KEY") or get_secret("GEMINI_API_KEY") if not key: location = "HF Space Secrets" if domain in CLOUD_ONLY_DOMAINS else "local .env/vault" raise MissingDomainKeyError( f"No API key configured for domain '{domain.value}'. " f"Add {env_var_name} to your {location} before this feature will work." ) return key def get_recovery_key() -> str: """Returns the Master Key. This should be kept strictly secret.""" return MASTER_KEY def encrypt_file(source_path: str, dest_path: str): """Encrypts a file from the PC and writes it to the Pendrive target.""" with open(source_path, 'rb') as f: file_data = f.read() encrypted_data = cipher_suite.encrypt(file_data) with open(dest_path, 'wb') as f: f.write(encrypted_data) logging.info(f"VAULT: Successfully encrypted {source_path} to {dest_path}") def decrypt_file(source_path: str, dest_path: str): """Decrypts a file from the Pendrive target and restores it to the PC.""" with open(source_path, 'rb') as f: encrypted_data = f.read() try: decrypted_data = cipher_suite.decrypt(encrypted_data) except Exception as e: logging.error(f"VAULT DECRYPTION ERROR: Invalid key or corrupted file. {e}") raise ValueError("Decryption failed. The file is corrupted or you are using the wrong JARVIS PC.") with open(dest_path, 'wb') as f: f.write(decrypted_data) logging.info(f"VAULT: Successfully decrypted {source_path} to {dest_path}") import shutil import tempfile def encrypt_directory(source_dir: str, dest_vault_file: str): """Zips an entire directory, encrypts it, and writes it to the pendrive.""" with tempfile.TemporaryDirectory() as tmpdirname: zip_path = os.path.join(tmpdirname, "archive") shutil.make_archive(zip_path, 'zip', source_dir) zip_file = zip_path + ".zip" encrypt_file(zip_file, dest_vault_file) logging.info(f"VAULT: Successfully zipped and encrypted directory {source_dir} to {dest_vault_file}") def decrypt_directory(source_vault_file: str, dest_dir: str): """Decrypts a vault file from the pendrive and extracts the directory to the PC.""" with tempfile.TemporaryDirectory() as tmpdirname: zip_file = os.path.join(tmpdirname, "decrypted.zip") decrypt_file(source_vault_file, zip_file) shutil.unpack_archive(zip_file, dest_dir, 'zip') logging.info(f"VAULT: Successfully decrypted and unpacked {source_vault_file} to {dest_dir}")