homelab_automation/app/services/terminal_gateway.py

"""
Terminal Gateway - PTY proxy for command interception and logging.

This service provides a PTY gateway that:
1. Creates a pseudo-terminal (PTY) for SSH connections
2. Intercepts input to capture commands before Enter
3. Logs validated commands to the database
4. Exposes the PTY to ttyd for web terminal access

Architecture:
    Browser <--> ttyd <--> Terminal Gateway (PTY proxy) <--> SSH <--> Remote Host
    
The gateway intercepts stdin going to SSH, reconstructs commands,
and logs them when Enter is detected.
"""
import asyncio
import logging
import os
import pty
import select
import shutil
import signal
import subprocess
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Callable, Dict, Optional, Tuple

from app.core.config import settings
from app.services.terminal_command_logger import TerminalCommandLogger, get_command_logger

logger = logging.getLogger(__name__)

# Configuration
SSH_USER = os.environ.get("TERMINAL_SSH_USER", "automation")
TTYD_PATH = os.environ.get("TTYD_PATH", "ttyd")
TERMINAL_PORT_RANGE_START = int(os.environ.get("TERMINAL_PORT_RANGE_START", "7680"))
TERMINAL_PORT_RANGE_END = int(os.environ.get("TERMINAL_PORT_RANGE_END", "7700"))


@dataclass
class GatewaySession:
    """Represents an active terminal gateway session."""
    session_id: str
    host_id: str
    host_name: str
    host_ip: str
    user_id: Optional[str]
    username: Optional[str]
    port: int
    
    # Process handles
    ttyd_process: Optional[subprocess.Popen] = None
    ssh_master_fd: Optional[int] = None
    ssh_slave_fd: Optional[int] = None
    ssh_pid: Optional[int] = None
    
    # Gateway script path
    gateway_script: Optional[str] = None
    
    # Stats
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    bytes_in: int = 0
    bytes_out: int = 0


class TerminalGateway:
    """
    Terminal Gateway service that manages PTY-based terminal sessions with command logging.
    
    This gateway creates a wrapper script that:
    1. Opens a PTY
    2. Runs SSH through the PTY
    3. Intercepts stdin to log commands
    4. Is then served by ttyd
    """
    
    def __init__(self):
        self._sessions: Dict[str, GatewaySession] = {}
        self._lock = asyncio.Lock()
        self._command_logger = get_command_logger()
        self._log_callback: Optional[Callable] = None
        
        # Check if we're on a Unix system (PTY requires Unix)
        self._pty_available = hasattr(pty, 'openpty') and sys.platform != 'win32'
        
        if not self._pty_available:
            logger.warning(
                "PTY not available on this platform. "
                "Command logging will be limited. "
                "Consider running on Linux for full functionality."
            )
    
    def set_log_callback(self, callback: Callable):
        """Set the callback function for logging commands to database."""
        self._log_callback = callback
        self._command_logger.set_log_callback(callback)
    
    def _find_ssh_key(self) -> Optional[str]:
        """Find available SSH key for terminal sessions."""
        key_candidates = []
        env_key = os.environ.get("TERMINAL_SSH_KEY_PATH")
        if env_key:
            key_candidates.append(env_key)
        if getattr(settings, "ssh_key_path", None):
            key_candidates.append(settings.ssh_key_path)
        key_candidates.append(os.path.expanduser("~/.ssh/id_automation_ansible"))
        key_candidates.append(os.path.expanduser("~/.ssh/id_rsa"))
        
        return next((p for p in key_candidates if p and os.path.exists(p)), None)
    
    def _create_gateway_script(
        self,
        session_id: str,
        host_ip: str,
        ssh_key_path: str,
    ) -> str:
        """
        Create a gateway script that wraps SSH and enables command logging.
        
        On Unix systems, this script uses a Python-based PTY wrapper.
        On Windows, we fall back to direct SSH (no command logging).
        """
        import tempfile
        
        # Create a temporary script
        script_dir = tempfile.mkdtemp(prefix=f"terminal_gateway_{session_id[:8]}_")
        script_path = os.path.join(script_dir, "gateway.sh")
        
        if self._pty_available:
            # Unix: Create a shell script that we can potentially enhance later
            # For now, just run SSH directly - the command logging happens via
            # the terminal_command_logger when we upgrade to full PTY interception
            script_content = f'''#!/bin/bash
# Terminal Gateway Script for session {session_id[:8]}
# This wraps SSH and could be enhanced for additional logging

exec ssh -i "{ssh_key_path}" \\
    -o BatchMode=no \\
    -o PreferredAuthentications=publickey,keyboard-interactive,password \\
    -o ConnectTimeout=10 \\
    -o ServerAliveInterval=30 \\
    -o ServerAliveCountMax=2 \\
    -o StrictHostKeyChecking=accept-new \\
    -o UserKnownHostsFile=/dev/null \\
    {SSH_USER}@{host_ip}
'''
        else:
            # Windows: PowerShell script
            script_path = os.path.join(script_dir, "gateway.ps1")
            script_content = f'''# Terminal Gateway Script for session {session_id[:8]}
ssh -i "{ssh_key_path}" `
    -o BatchMode=no `
    -o PreferredAuthentications=publickey,keyboard-interactive,password `
    -o ConnectTimeout=10 `
    -o ServerAliveInterval=30 `
    -o ServerAliveCountMax=2 `
    -o StrictHostKeyChecking=accept-new `
    -o UserKnownHostsFile=/dev/null `
    {SSH_USER}@{host_ip}
'''
        
        with open(script_path, 'w') as f:
            f.write(script_content)
        
        if self._pty_available:
            os.chmod(script_path, 0o755)
        
        return script_path
    
    async def create_session(
        self,
        session_id: str,
        host_id: str,
        host_name: str,
        host_ip: str,
        port: int,
        user_id: Optional[str] = None,
        username: Optional[str] = None,
    ) -> GatewaySession:
        """
        Create a new terminal gateway session.
        
        This sets up the command logger context for the session.
        The actual ttyd process is spawned separately.
        """
        async with self._lock:
            session = GatewaySession(
                session_id=session_id,
                host_id=host_id,
                host_name=host_name,
                host_ip=host_ip,
                user_id=user_id,
                username=username,
                port=port,
            )
            self._sessions[session_id] = session
        
        # Create command logger session
        await self._command_logger.create_session(
            session_id=session_id,
            host_id=host_id,
            host_name=host_name,
            user_id=user_id,
            username=username,
        )
        
        logger.info(f"Created gateway session {session_id[:8]}... for {host_name} ({host_ip})")
        return session
    
    async def spawn_ttyd_with_gateway(
        self,
        session: GatewaySession,
        token: str,
    ) -> Optional[int]:
        """
        Spawn ttyd with a gateway wrapper script.
        
        Returns:
            PID of ttyd process, or None if failed
        """
        if not shutil.which(TTYD_PATH):
            raise RuntimeError("ttyd is not installed or not in PATH")
        
        ssh_key_path = self._find_ssh_key()
        if not ssh_key_path:
            raise RuntimeError("No SSH key found for terminal sessions")
        
        # Create gateway script
        gateway_script = self._create_gateway_script(
            session.session_id,
            session.host_ip,
            ssh_key_path,
        )
        session.gateway_script = gateway_script
        
        # Build ttyd command
        if self._pty_available:
            cmd = [
                TTYD_PATH,
                "--once",
                f"--port={session.port}",
                "--writable",
                gateway_script,
            ]
        else:
            # Windows: Use ttyd with direct SSH command
            cmd = [
                TTYD_PATH,
                "--once",
                f"--port={session.port}",
                "--writable",
                "ssh",
                "-i", ssh_key_path,
                "-o", "BatchMode=no",
                "-o", "PreferredAuthentications=publickey,keyboard-interactive,password",
                "-o", "ConnectTimeout=10",
                "-o", "ServerAliveInterval=30",
                "-o", "ServerAliveCountMax=2",
                "-o", "StrictHostKeyChecking=accept-new",
                "-o", "UserKnownHostsFile=/dev/null",
                f"{SSH_USER}@{session.host_ip}",
            ]
        
        logger.info(
            f"Spawning ttyd gateway for session {session.session_id[:8]}... "
            f"on port {session.port} -> {SSH_USER}@{session.host_ip}"
        )
        
        try:
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                start_new_session=True,
            )
            
            session.ttyd_process = process
            
            # Wait briefly to check if process started successfully
            await asyncio.sleep(0.5)
            
            if process.poll() is not None:
                stdout, stderr = process.communicate(timeout=1)
                error_msg = stderr.decode() if stderr else stdout.decode() if stdout else "Unknown error"
                logger.error(f"ttyd gateway failed to start: {error_msg}")
                self._cleanup_gateway_script(session)
                return None
            
            logger.info(f"ttyd gateway started with PID {process.pid}")
            return process.pid
            
        except Exception as e:
            logger.exception(f"Failed to spawn ttyd gateway: {e}")
            self._cleanup_gateway_script(session)
            return None
    
    def _cleanup_gateway_script(self, session: GatewaySession):
        """Clean up the gateway script directory."""
        if session.gateway_script:
            try:
                script_dir = os.path.dirname(session.gateway_script)
                if os.path.exists(script_dir):
                    import shutil
                    shutil.rmtree(script_dir, ignore_errors=True)
            except Exception as e:
                logger.warning(f"Failed to cleanup gateway script: {e}")
    
    async def terminate_session(self, session_id: str) -> bool:
        """
        Terminate a terminal gateway session.
        
        Returns:
            True if session was terminated, False if not found
        """
        async with self._lock:
            session = self._sessions.pop(session_id, None)
        
        if not session:
            return False
        
        # Flush any remaining command buffer
        await self._command_logger.flush_buffer(session_id)
        
        # Remove command logger session
        await self._command_logger.remove_session(session_id)
        
        # Terminate ttyd process
        if session.ttyd_process and session.ttyd_process.poll() is None:
            try:
                session.ttyd_process.terminate()
                try:
                    session.ttyd_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    session.ttyd_process.kill()
                    session.ttyd_process.wait(timeout=2)
            except Exception as e:
                logger.warning(f"Error terminating ttyd process: {e}")
        
        # Cleanup gateway script
        self._cleanup_gateway_script(session)
        
        logger.info(f"Terminated gateway session {session_id[:8]}...")
        return True
    
    async def process_terminal_input(
        self,
        session_id: str,
        data: bytes,
    ) -> None:
        """
        Process terminal input for command logging.
        
        This is called when we intercept input going to the terminal.
        In the current architecture with ttyd, this would require
        a WebSocket proxy to intercept. For now, this is a placeholder
        for future enhancement.
        """
        results = await self._command_logger.process_input(session_id, data)
        
        for result in results:
            if result.should_log:
                logger.debug(f"Command captured: {result.masked_command[:50]}...")
    
    def get_session(self, session_id: str) -> Optional[GatewaySession]:
        """Get a gateway session by ID."""
        return self._sessions.get(session_id)
    
    def get_active_session_count(self) -> int:
        """Get the number of active gateway sessions."""
        return len(self._sessions)
    
    def get_stats(self) -> dict:
        """Get statistics about the terminal gateway."""
        cmd_stats = self._command_logger.get_stats()
        return {
            "active_sessions": len(self._sessions),
            "pty_available": self._pty_available,
            **cmd_stats,
        }


# ============================================================================
# Global instance
# ============================================================================
_gateway_instance: Optional[TerminalGateway] = None


def get_terminal_gateway() -> TerminalGateway:
    """Get or create the global terminal gateway instance."""
    global _gateway_instance
    
    if _gateway_instance is None:
        _gateway_instance = TerminalGateway()
    
    return _gateway_instance