backend/security/vault_builder.py

import asyncio
import shutil
import tempfile
import os
import sys
import subprocess
from pathlib import Path
from typing import Literal
from backend.services.usb_vault import encrypt_directory
from backend.security.vault_sanitizer import strip_personal_data

async def stage_full_environment(source_path: str, staging_dir: str):
    omega_root = source_path
    
    # ── 1. MAIN CODEBASE ─────────────────────────────────────────────
    if getattr(sys, 'frozen', False):
        shutil.copy2(sys.executable, os.path.join(staging_dir, "JARVIS_OMEGA.exe"))
    else:
        code_staging = os.path.join(staging_dir, "F.R.I.D.A.Y - OMEGA")
        shutil.copytree(omega_root, code_staging, dirs_exist_ok=True)

    # ── 2. HF DEPLOY STAGE SERVER ────────────────────────────────────
    hf_stage = os.path.join(omega_root, "scratch", "hf_deploy_stage")
    if os.path.exists(hf_stage):
        shutil.copytree(hf_stage, os.path.join(staging_dir, "hf_deploy_stage"), dirs_exist_ok=True)

    # ── 3. CLOUD DEPLOYMENT SERVER ───────────────────────────────────
    cloud_dir = os.path.join(omega_root, "cloud_deployment")
    if os.path.exists(cloud_dir):
        shutil.copytree(cloud_dir, os.path.join(staging_dir, "cloud_deployment"), dirs_exist_ok=True)

    # ── 4. ANDROID APK SERVER (phone guardian) ───────────────────────
    phone_dir = os.path.join(omega_root, "phone")
    if os.path.exists(phone_dir):
        shutil.copytree(phone_dir, os.path.join(staging_dir, "phone"), dirs_exist_ok=True)

    # ── 5. CONFIG + SECRETS ──────────────────────────────────────────
    for cfg_file in ["config.local.json", "config.json", ".env.example"]:
        cfg_path = os.path.join(omega_root, cfg_file)
        if os.path.exists(cfg_path):
            shutil.copy2(cfg_path, os.path.join(staging_dir, cfg_file))

    # ── 6. APPDATA: SQLite DB + allowlist + vault_keys ───────────────
    appdata_path = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS')
    if os.path.exists(appdata_path):
        shutil.copytree(appdata_path, os.path.join(staging_dir, "JARVIS_OS_APPDATA"), dirs_exist_ok=True)

    # ── 7. LOCAL memory.db (dev mode) & Chroma & Dumps ───────────────
    local_db = os.path.join(omega_root, "memory.db")
    if os.path.exists(local_db):
        shutil.copy2(local_db, os.path.join(staging_dir, "memory.db"))
        try:
            subprocess.run(f'sqlite3 "{local_db}" .dump > "{os.path.join(staging_dir, "sqlite_backup.sql")}"', shell=True)
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

    chroma_dir = os.path.join(omega_root, "memory_chroma")
    if os.path.exists(chroma_dir):
        shutil.copytree(chroma_dir, os.path.join(staging_dir, "chroma"), dirs_exist_ok=True)
    
    try:
        subprocess.run(f"pip freeze > '{os.path.join(staging_dir, 'requirements_frozen.txt')}'", shell=True)
    except Exception as e:
        import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")
    
    try:
        mongo_uri = os.environ.get("MONGO_URI", "")
        if not mongo_uri:
            env_path = os.path.join(omega_root, ".env")
            if os.path.exists(env_path):
                with open(env_path, "r", encoding="utf-8") as f:
                    for line in f:
                        if line.startswith("MONGO_URI="):
                            mongo_uri = line.split("=", 1)[1].strip()
                            break
        if mongo_uri:
            subprocess.run(f"mongoexport --uri='{mongo_uri}' --out='{os.path.join(staging_dir, 'mongo_backup.json')}'", shell=True)
    except Exception as e:
        import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

    # ── 8. FAMILY DEVICE ENROLLMENT DATA (Explicit Manifest Inclusion) ─
    enrollment_data = os.path.join(omega_root, "family_device_enrollment_data")
    if os.path.exists(enrollment_data):
        shutil.copytree(enrollment_data, os.path.join(staging_dir, "family_device_enrollment_data"), dirs_exist_ok=True)

async def build_vault(mode: Literal["full", "sanitized"], source_path: str, target_path: str):
    def _sync():
        # First we stage into a temp dir
        temp_base = tempfile.mkdtemp()
        staging_dir = os.path.join(temp_base, "staging")
        os.makedirs(staging_dir, exist_ok=True)
        return temp_base, staging_dir

    temp_base, staging_dir = await asyncio.to_thread(_sync)
    
    try:
        await asyncio.to_thread(lambda: asyncio.run(stage_full_environment(source_path, staging_dir)))
    except RuntimeError:
        await stage_full_environment(source_path, staging_dir)
        
    final_dir_to_encrypt = staging_dir

    if mode == "sanitized":
        # Run sanitizer which creates a new directory beside staging
        final_dir_to_encrypt = await strip_personal_data(Path(staging_dir))
        final_dir_to_encrypt = str(final_dir_to_encrypt)
        
    def _encrypt():
        encrypt_directory(final_dir_to_encrypt, target_path)
    
    await asyncio.to_thread(_encrypt)
    
    def _cleanup():
        shutil.rmtree(temp_base, ignore_errors=True)
    await asyncio.to_thread(_cleanup)

    return target_path