#!/usr/bin/env python3
# keyCocoon/keyBunker - Credential Management
# Copyright (c) 2025 Gradient Research LLC
# All rights reserved.

import json
import sys
import struct
import os
import platform
from datetime import datetime
import subprocess
import time

# Ensure sibling modules (keyBunker.py) are importable when launched by the
# embedded webSlinger Python, whose python312._pth puts the interpreter in
# isolated mode and does NOT auto-add the script's directory to sys.path.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# The embedded webSlinger Python's ._pth file adds site-packages to sys.path
# raw, but does NOT process .pth files inside it. pywin32 ships pywin32.pth,
# which adds win32/, win32/lib/, and runs pywin32_bootstrap (sets DLL search
# dir). Without that processing, `import win32file` fails. Call addsitedir on
# the embedded site-packages so pywin32 becomes importable. No-op on normal
# Python distributions where site already processed it.
import site as _site
_emb_site_packages = os.path.join(os.path.dirname(os.path.abspath(sys.executable)), '..', 'site-packages')
if os.path.isdir(_emb_site_packages):
    _site.addsitedir(_emb_site_packages)

# Critical Windows fix: Set binary I/O mode
if sys.platform == "win32":
    import msvcrt
    msvcrt.setmode(sys.stdin.fileno(), os.O_BINARY)
    msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY)

# Import KeyBunker
try:
    from keyBunker import keyBunker
    KEYBUNKER_AVAILABLE = True
except ImportError as e:
    KEYBUNKER_AVAILABLE = False
    KEYBUNKER_ERROR = str(e)

# Simple logging to file
def log_to_file(message):
    try:
        with open(os.path.expanduser('~/keybunker_native_host.log'), 'a', encoding='utf-8') as f:
            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            f.write(f"[{timestamp}] {message}\n")
            f.flush()
    except:
        pass

def send_message(message):
    """Send message to Chrome"""
    # Filter credentials from logging
    log_message = message.copy() if isinstance(message, dict) else message
    if isinstance(log_message, dict) and 'credentials' in log_message:
        log_message = {**log_message, 'credentials': '[FILTERED]'}
    
    log_to_file(f"Sending: {log_message}")
    
    json_message = json.dumps(message)
    encoded_message = json_message.encode('utf-8')
    
    message_length = len(encoded_message)
    sys.stdout.buffer.write(struct.pack('@I', message_length))
    sys.stdout.buffer.write(encoded_message)
    sys.stdout.buffer.flush()
    
    log_to_file(f"Sent {message_length} bytes")

def read_message():
    """Read message from Chrome"""
    try:
        text_length_bytes = sys.stdin.buffer.read(4)
        
        if len(text_length_bytes) == 0:
            log_to_file("No input - Chrome disconnected")
            return None
        
        if len(text_length_bytes) != 4:
            log_to_file(f"Invalid length header: got {len(text_length_bytes)} bytes")
            return None
        
        text_length = struct.unpack('@I', text_length_bytes)[0]
        log_to_file(f"Message length: {text_length}")
        
        if text_length > 1024 * 1024:
            log_to_file(f"Message too large: {text_length}")
            return None
        
        text = sys.stdin.buffer.read(text_length).decode('utf-8')
        log_to_file(f"Received: {text}")
        
        message = json.loads(text)
        return message
        
    except Exception as e:
        log_to_file(f"Error reading message: {e}")
        return None

def handle_message(message):
    """Handle incoming message"""
    if not message:
        return {"success": False, "error": "No message received"}
    
    # Support both 'command' and 'action' message formats
    command = message.get('command') or message.get('action', 'unknown')
    log_to_file(f"Handling command: {command}")
    
    if command == 'ping':
        return {
            'success': True,
            'message': 'Pong! Native host is working correctly',
            'timestamp': datetime.now().isoformat(),
            'platform': platform.system(),
            'python_version': platform.python_version(),
            'received': message
        }
    
    elif command == 'status':
        if not KEYBUNKER_AVAILABLE:
            return {'success': False, 'available': False, 'authenticated': False, 'message': f'KeyBunker not available: {KEYBUNKER_ERROR}'}
        
        try:
            kb = keyBunker()
            # Actually try to get session key from IPC (daemon)
            try:
                session_key = kb._get_session_key_from_ipc()
                if session_key:
                    return {'success': True, 'available': True, 'authenticated': True, 'message': 'Authenticated and ready via daemon'}
                else:
                    return {'success': True, 'available': True, 'authenticated': False, 'message': 'Please enter your master password'}
            except Exception as e:
                return {'success': True, 'available': True, 'authenticated': False, 'message': 'Please enter your master password'}

        except Exception as e:
            return {'success': True, 'available': True, 'authenticated': False, 'message': 'Please enter your master password'}   
            
    elif command == 'system_info':
        return {
            'success': True,
            'system': {
                'platform': platform.platform(),
                'system': platform.system(),
                'release': platform.release(),
                'machine': platform.machine(),
                'python_version': platform.python_version(),
                'keybunker_available': KEYBUNKER_AVAILABLE,
                'keybunker_error': KEYBUNKER_ERROR,
                'working_directory': os.getcwd(),
                'script_directory': os.path.dirname(os.path.abspath(__file__))
            }
        }
            
    elif command == 'register':
        if not KEYBUNKER_AVAILABLE:
            return {'success': False, 'message': f'KeyBunker not available: {KEYBUNKER_ERROR}'}
        
        password = message.get('password')
        if not password:
            return {'success': False, 'message': 'Password required for registration'}
        
        try:
            # Check if daemon is already running
            kb = keyBunker()
            try:
                session_key = kb._get_session_key_from_ipc()
                if session_key:
                    return {'success': True, 'message': 'KeyBunker already authenticated via daemon'}
            except Exception:
                pass  # No daemon running, proceed to start one
            
            # Start daemon process
            import subprocess
            daemon_script = os.path.join(os.path.dirname(__file__), 'keybunker_daemon.py')
            python_exe = os.path.join(os.getenv('LOCALAPPDATA'), 'webSlinger', 'python', 'python.exe')
            
            log_to_file(f"Starting daemon: {python_exe} {daemon_script}")
            
            # Start daemon in background
            # Get timeout from message or use default
            timeout_minutes = message.get('timeout_minutes', 60)  # Default 30 minutes
            
            # Start daemon in background
            process = subprocess.Popen([
                python_exe, '-B', daemon_script, password, str(timeout_minutes)
            ], creationflags=subprocess.CREATE_NO_WINDOW)
            
            log_to_file(f"Daemon started with PID: {process.pid}")
            
            # Wait a moment for daemon to initialize
            time.sleep(2)
            
            # Verify daemon is working
            try:
                kb_test = keyBunker()
                session_key = kb_test._get_session_key_from_ipc()
                if session_key:
                    return {'success': True, 'message': 'KeyBunker daemon started and authenticated'}
                else:
                    return {'success': False, 'message': 'Daemon started but no session key available'}
            except Exception as e:
                return {'success': False, 'message': f'Daemon verification failed: {str(e)}'}
                
        except Exception as e:
            log_to_file(f"Register error: {str(e)}")
            return {'success': False, 'message': f'Registration error: {str(e)}'}

    elif command == 'reset':
        """Reset authentication by sending reset command to daemon"""
        try:
            if sys.platform == 'win32':
                try:
                    import win32file
                    import win32pipe
                    import pywintypes

                    pipe_name = r'\\.\pipe\keybunker_cmd'
                    log_to_file("Connecting to daemon command pipe")

                    # Connect to daemon's command pipe
                    handle = win32file.CreateFile(
                        pipe_name,
                        win32file.GENERIC_READ | win32file.GENERIC_WRITE,
                        0,
                        None,
                        win32file.OPEN_EXISTING,
                        0,
                        None
                    )

                    # Send reset command
                    win32file.WriteFile(handle, b'reset')

                    # Read response
                    result, data = win32file.ReadFile(handle, 1024)
                    response = data.decode('utf-8')

                    win32file.CloseHandle(handle)

                    log_to_file(f"Reset command sent, response: {response}")

                    if response == 'OK':
                        return {'success': True, 'message': 'Authentication reset successfully'}
                    else:
                        return {'success': False, 'message': f'Reset failed: {response}'}

                except pywintypes.error as e:
                    # If pipe doesn't exist, daemon isn't running - that's fine
                    if e.winerror == 2:  # ERROR_FILE_NOT_FOUND
                        log_to_file("Daemon not running (pipe not found)")
                        return {'success': True, 'message': 'No active session to reset'}
                    else:
                        log_to_file(f"Pipe error: {e}")
                        return {'success': False, 'message': f'Reset error: {str(e)}'}

            else:
                # Unix/Mac - use socket or signal approach
                return {'success': False, 'message': 'Reset not implemented for this platform'}

        except Exception as e:
            log_to_file(f"Reset error: {str(e)}")
            return {'success': False, 'message': f'Reset error: {str(e)}'}

    elif command == 'registerScript':
        if not KEYBUNKER_AVAILABLE:
            return {'success': False, 'message': f'KeyBunker not available: {KEYBUNKER_ERROR}'}

        script_name = message.get('scriptName')
        script_hash = message.get('scriptHash')
        manifest = message.get('manifest')

        if not script_name or not script_hash or not manifest:
            return {'success': False, 'message': 'scriptName, scriptHash, and manifest are required'}
        if not isinstance(script_hash, str) or len(script_hash) != 64:
            return {'success': False, 'message': f'scriptHash must be 64-char hex, got len={len(script_hash) if isinstance(script_hash, str) else "non-string"}'}

        try:
            kb = keyBunker()
            session_key = kb._get_session_key_from_ipc()
            if not session_key:
                return {'success': False, 'message': 'KeyBunker not authenticated'}
            kb._session_key = session_key
            kb._locate_vault()
            kb._read_vault()
            kb.registerScript(script_name, script_hash, manifest)
            return {'success': True, 'message': f'Script "{script_name}" registered'}
        except Exception as e:
            log_to_file(f"registerScript error: {str(e)}")
            return {'success': False, 'message': f'registerScript error: {str(e)}'}

    elif command == 'getCredentials':
        if not KEYBUNKER_AVAILABLE:
            return {'success': False, 'message': f'KeyBunker not available: {KEYBUNKER_ERROR}'}

        domain = message.get('domain')
        username = message.get('username')
        public_key_b64 = message.get('publicKey')
        script_name = message.get('scriptName')
        script_hash = message.get('scriptHash')

        if not domain or not username:
            return {'success': False, 'message': 'Domain and username required'}

        if not public_key_b64:
            return {'success': False, 'message': 'Public key required for encryption'}

        try:
            # Import public key
            from cryptography.hazmat.primitives.asymmetric import padding
            from cryptography.hazmat.primitives import serialization, hashes
            import base64

            public_key_bytes = base64.b64decode(public_key_b64)
            public_key = serialization.load_der_public_key(public_key_bytes)

            kb = keyBunker()

            # Try IPC authentication
            try:
                session_key = kb._get_session_key_from_ipc()
                if session_key:
                    kb._session_key = session_key
                    kb._locate_vault()
                    kb._read_vault()
                else:
                    return {'success': False, 'message': 'KeyBunker not authenticated'}
            except Exception as e:
                return {'success': False, 'message': f'KeyBunker IPC failed: {str(e)}'}

            # Script-registry integrity gate. Only enforced when the caller
            # supplies a script identity (automation phase). Session recording
            # happens before any script exists, so absence of both fields
            # means "manual/recording context" - allow credential release.
            if script_name or script_hash:
                if not script_name or not script_hash:
                    return {'success': False, 'message': 'scriptName and scriptHash must both be provided when either is set'}
                registration = kb.getScriptRegistration(script_name)
                if not registration:
                    return {'success': False, 'message': f'Script "{script_name}" not registered. Approve the script in the webSlinger wizard before running.'}
                if registration['hash'] != script_hash:
                    return {'success': False, 'message': f'Script "{script_name}" hash mismatch. The script bytes differ from the version that was approved. Re-approve the script to register the new version.'}
                manifest = registration.get('manifest') or {}
                cred_domains = manifest.get('credentialDomains') or []
                domain_norm = domain.strip().lower().lstrip('.').rstrip('/')
                allowed = any(
                    (entry.get('domain') or '').strip().lower().lstrip('.').rstrip('/') == domain_norm
                    for entry in cred_domains
                )
                if not allowed:
                    return {'success': False, 'message': f'Domain "{domain}" not in approved credential domains for script "{script_name}".'}

            # Get credentials
            credentials = kb.getCredentialsForWebsite(domain, username)
            
            # Encrypt password with public key
            if 'password' in credentials:
                encrypted = public_key.encrypt(
                    credentials['password'].encode(),
                    padding.OAEP(
                        mgf=padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                credentials['password'] = base64.b64encode(encrypted).decode()
            
            return {'success': True, 'credentials': credentials}
            
        except Exception as e:
            return {'success': False, 'message': f'Credential error: {str(e)}'}
        
    elif command == 'getTOTP':
        if not KEYBUNKER_AVAILABLE:
            return {'success': False, 'message': f'KeyBunker not available: {KEYBUNKER_ERROR}'}

        domain = message.get('domain')
        username = message.get('username')
        script_name = message.get('scriptName')
        script_hash = message.get('scriptHash')

        if not domain or not username:
            return {'success': False, 'message': 'Domain and username required'}

        try:
            kb = keyBunker()

            # Try IPC authentication
            try:
                session_key = kb._get_session_key_from_ipc()
                if session_key:
                    kb._session_key = session_key
                    kb._locate_vault()
                    kb._read_vault()
                else:
                    return {'success': False, 'message': 'KeyBunker not authenticated - start your AI agent with kb.register(password)'}
            except Exception as e:
                return {'success': False, 'message': f'KeyBunker IPC failed: {str(e)}'}

            # Script-registry integrity gate (mirrors getCredentials). Only
            # enforced when caller supplies a script identity. Absence of both
            # fields = recording/manual context, allow normal release.
            if script_name or script_hash:
                if not script_name or not script_hash:
                    return {'success': False, 'message': 'scriptName and scriptHash must both be provided when either is set'}
                registration = kb.getScriptRegistration(script_name)
                if not registration:
                    return {'success': False, 'message': f'Script "{script_name}" not registered.'}
                if registration['hash'] != script_hash:
                    return {'success': False, 'message': f'Script "{script_name}" hash mismatch. Re-approve the script to register the new version.'}
                manifest = registration.get('manifest') or {}
                cred_domains = manifest.get('credentialDomains') or []
                domain_norm = domain.strip().lower().lstrip('.').rstrip('/')
                allowed = any(
                    (entry.get('domain') or '').strip().lower().lstrip('.').rstrip('/') == domain_norm
                    and 'totp' in (entry.get('type') or [])
                    for entry in cred_domains
                )
                if not allowed:
                    return {'success': False, 'message': f'Domain "{domain}" not approved for TOTP release in script "{script_name}".'}

            # Generate TOTP code
            totp_code = kb.getTOTPForWebsite(domain, username)
            return {'success': True, 'code': totp_code}
            
        except Exception as e:
            return {'success': False, 'message': f'TOTP error: {str(e)}'}

    else:
        return {
            'success': False,
            'error': f'Unknown command: {command}',
            'available_commands': ['ping', 'status', 'getCredentials', 'getTOTP', 'registerScript', 'system_info']
        }

def main():
    log_to_file("=== KEYBUNKER NATIVE HOST STARTED ===")
    log_to_file(f"Platform: {platform.system()}")
    log_to_file(f"Python: {platform.python_version()}")
    log_to_file(f"KeyBunker available: {KEYBUNKER_AVAILABLE}")
    log_to_file(f"Args: {sys.argv}")
    
    try:
        message = read_message()
        
        if message is not None:
            response = handle_message(message)
            send_message(response)
            log_to_file("Message handled successfully")
        else:
            log_to_file("No valid message received")
            
    except Exception as e:
        log_to_file(f"Fatal error: {e}")
        error_response = {
            'success': False,
            'error': f'Native host error: {str(e)}'
        }
        try:
            send_message(error_response)
        except:
            pass
    
    log_to_file("=== KEYBUNKER NATIVE HOST EXITING ===")

if __name__ == '__main__':
    main()