Add terminal session management with heartbeat monitoring, idle timeout detection, session reuse logic, and command history panel UI with search and filtering capabilities
Tests / Backend Tests (Python) (3.10) (push) Has been cancelled
Tests / Backend Tests (Python) (3.11) (push) Has been cancelled
Tests / Backend Tests (Python) (3.12) (push) Has been cancelled
Tests / Frontend Tests (JS) (push) Has been cancelled
Tests / Integration Tests (push) Has been cancelled
Tests / All Tests Passed (push) Has been cancelled
Tests / Backend Tests (Python) (3.10) (push) Has been cancelled
Tests / Backend Tests (Python) (3.11) (push) Has been cancelled
Tests / Backend Tests (Python) (3.12) (push) Has been cancelled
Tests / Frontend Tests (JS) (push) Has been cancelled
Tests / Integration Tests (push) Has been cancelled
Tests / All Tests Passed (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,322 @@
|
||||
"""
|
||||
Terminal Command Logger - Captures and logs validated commands from terminal sessions.
|
||||
|
||||
This module provides:
|
||||
- Command buffer management per session
|
||||
- Enter key detection to capture complete commands
|
||||
- Integration with CommandPolicy for validation
|
||||
- Async database logging
|
||||
|
||||
Key features:
|
||||
- Handles backspace, Ctrl+U, and Enter sequences
|
||||
- Only logs commands that pass security validation
|
||||
- Thread-safe buffer management
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
from app.security.command_policy import CommandPolicy, CommandPolicyResult, get_command_policy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Configuration
|
||||
TERMINAL_COMMAND_RETENTION_DAYS = int(os.environ.get("TERMINAL_COMMAND_RETENTION_DAYS", "30"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionContext:
|
||||
"""Context for a terminal session's command buffer."""
|
||||
session_id: str
|
||||
host_id: str
|
||||
host_name: str
|
||||
user_id: Optional[str]
|
||||
username: Optional[str]
|
||||
|
||||
# Command buffer (characters waiting for Enter)
|
||||
buffer: str = ""
|
||||
|
||||
# Stats
|
||||
commands_logged: int = 0
|
||||
commands_blocked: int = 0
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class TerminalCommandLogger:
|
||||
"""
|
||||
Manages command buffers for terminal sessions and logs validated commands.
|
||||
|
||||
Usage:
|
||||
logger = TerminalCommandLogger()
|
||||
|
||||
# Create session context
|
||||
logger.create_session(session_id, host_id, host_name, user_id, username)
|
||||
|
||||
# Process input bytes from terminal
|
||||
async for result in logger.process_input(session_id, input_bytes):
|
||||
if result.should_log:
|
||||
# Command was logged to database
|
||||
pass
|
||||
|
||||
# Cleanup when session ends
|
||||
logger.remove_session(session_id)
|
||||
"""
|
||||
|
||||
def __init__(self, policy: Optional[CommandPolicy] = None):
|
||||
self._sessions: Dict[str, SessionContext] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
self._policy = policy or get_command_policy()
|
||||
|
||||
# Callback for logging to database (injected)
|
||||
self._log_callback: Optional[Callable] = None
|
||||
|
||||
def set_log_callback(self, callback: Callable):
|
||||
"""Set the callback function for logging commands to database."""
|
||||
self._log_callback = callback
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
session_id: str,
|
||||
host_id: str,
|
||||
host_name: str,
|
||||
user_id: Optional[str] = None,
|
||||
username: Optional[str] = None,
|
||||
) -> SessionContext:
|
||||
"""Create a new session context for command buffering."""
|
||||
async with self._lock:
|
||||
ctx = SessionContext(
|
||||
session_id=session_id,
|
||||
host_id=host_id,
|
||||
host_name=host_name,
|
||||
user_id=user_id,
|
||||
username=username,
|
||||
)
|
||||
self._sessions[session_id] = ctx
|
||||
logger.debug(f"Created command logger session: {session_id[:8]}... for {host_name}")
|
||||
return ctx
|
||||
|
||||
async def remove_session(self, session_id: str) -> Optional[SessionContext]:
|
||||
"""Remove a session context."""
|
||||
async with self._lock:
|
||||
ctx = self._sessions.pop(session_id, None)
|
||||
if ctx:
|
||||
logger.debug(
|
||||
f"Removed command logger session: {session_id[:8]}... "
|
||||
f"(logged={ctx.commands_logged}, blocked={ctx.commands_blocked})"
|
||||
)
|
||||
return ctx
|
||||
|
||||
def get_session(self, session_id: str) -> Optional[SessionContext]:
|
||||
"""Get a session context."""
|
||||
return self._sessions.get(session_id)
|
||||
|
||||
def _handle_control_char(self, ctx: SessionContext, char: bytes) -> bool:
|
||||
"""
|
||||
Handle control characters in the input stream.
|
||||
|
||||
Returns True if the character was handled, False otherwise.
|
||||
"""
|
||||
if char in (b'\x7f', b'\b'): # Backspace (DEL or BS)
|
||||
if ctx.buffer:
|
||||
ctx.buffer = ctx.buffer[:-1]
|
||||
return True
|
||||
|
||||
if char == b'\x15': # Ctrl+U - clear line
|
||||
ctx.buffer = ""
|
||||
return True
|
||||
|
||||
if char == b'\x17': # Ctrl+W - delete word
|
||||
# Delete last word (back to space or start)
|
||||
ctx.buffer = ctx.buffer.rstrip()
|
||||
if ' ' in ctx.buffer:
|
||||
ctx.buffer = ctx.buffer[:ctx.buffer.rfind(' ') + 1]
|
||||
else:
|
||||
ctx.buffer = ""
|
||||
return True
|
||||
|
||||
if char == b'\x03': # Ctrl+C - interrupt
|
||||
ctx.buffer = ""
|
||||
return True
|
||||
|
||||
if char == b'\x04': # Ctrl+D - EOF
|
||||
ctx.buffer = ""
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def process_input(
|
||||
self,
|
||||
session_id: str,
|
||||
data: bytes,
|
||||
) -> list[CommandPolicyResult]:
|
||||
"""
|
||||
Process input bytes from a terminal session.
|
||||
|
||||
This method:
|
||||
1. Adds printable characters to the buffer
|
||||
2. Handles control characters (backspace, Ctrl+U, etc.)
|
||||
3. When Enter is detected, evaluates the command and logs if valid
|
||||
|
||||
Args:
|
||||
session_id: The terminal session ID
|
||||
data: Raw bytes from the terminal input stream
|
||||
|
||||
Returns:
|
||||
List of CommandPolicyResult for any commands that were evaluated
|
||||
"""
|
||||
ctx = self._sessions.get(session_id)
|
||||
if not ctx:
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
# Process byte by byte
|
||||
i = 0
|
||||
while i < len(data):
|
||||
char = data[i:i+1]
|
||||
|
||||
# Check for Enter (CR or LF)
|
||||
if char in (b'\r', b'\n'):
|
||||
# Skip CRLF sequence
|
||||
if char == b'\r' and i + 1 < len(data) and data[i+1:i+2] == b'\n':
|
||||
i += 1
|
||||
|
||||
# Process the command
|
||||
command = ctx.buffer.strip()
|
||||
ctx.buffer = ""
|
||||
|
||||
if command:
|
||||
result = await self._evaluate_and_log(ctx, command)
|
||||
results.append(result)
|
||||
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Handle control characters
|
||||
if self._handle_control_char(ctx, char):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Skip other control characters and escape sequences
|
||||
if char[0] < 32 or char == b'\x7f':
|
||||
# Check for escape sequence
|
||||
if char == b'\x1b' and i + 1 < len(data):
|
||||
# Skip escape sequences (arrow keys, etc.)
|
||||
i += 1
|
||||
while i < len(data) and data[i:i+1] not in (b'', b'\r', b'\n'):
|
||||
next_char = data[i:i+1]
|
||||
if next_char.isalpha() or next_char == b'~':
|
||||
i += 1
|
||||
break
|
||||
i += 1
|
||||
continue
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Add printable character to buffer
|
||||
try:
|
||||
ctx.buffer += char.decode('utf-8', errors='ignore')
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
i += 1
|
||||
|
||||
return results
|
||||
|
||||
async def _evaluate_and_log(
|
||||
self,
|
||||
ctx: SessionContext,
|
||||
command: str,
|
||||
) -> CommandPolicyResult:
|
||||
"""Evaluate a command against policy and log if valid."""
|
||||
result = self._policy.evaluate(command)
|
||||
|
||||
if result.should_log:
|
||||
ctx.commands_logged += 1
|
||||
logger.debug(f"Command logged for {ctx.host_name}: {result.masked_command[:50]}...")
|
||||
|
||||
# Call the log callback if set
|
||||
if self._log_callback:
|
||||
try:
|
||||
await self._log_callback(
|
||||
host_id=ctx.host_id,
|
||||
host_name=ctx.host_name,
|
||||
user_id=ctx.user_id,
|
||||
username=ctx.username,
|
||||
session_id=ctx.session_id,
|
||||
command=result.masked_command,
|
||||
command_hash=result.command_hash,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to log command to database: {e}")
|
||||
|
||||
elif result.is_blocked:
|
||||
ctx.commands_blocked += 1
|
||||
logger.info(f"Command blocked for {ctx.host_name}: {result.reason}")
|
||||
|
||||
# Optionally log blocked command metadata (not the command itself)
|
||||
if self._log_callback:
|
||||
try:
|
||||
await self._log_callback(
|
||||
host_id=ctx.host_id,
|
||||
host_name=ctx.host_name,
|
||||
user_id=ctx.user_id,
|
||||
username=ctx.username,
|
||||
session_id=ctx.session_id,
|
||||
command="[BLOCKED]",
|
||||
command_hash="blocked",
|
||||
is_blocked=True,
|
||||
blocked_reason=result.reason,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to log blocked command metadata: {e}")
|
||||
|
||||
else:
|
||||
# Command not in allowlist - silently skip
|
||||
logger.debug(f"Command not in allowlist for {ctx.host_name}: {command[:30]}...")
|
||||
|
||||
return result
|
||||
|
||||
async def flush_buffer(self, session_id: str) -> Optional[CommandPolicyResult]:
|
||||
"""
|
||||
Flush any remaining buffer content as a command.
|
||||
|
||||
Useful when session ends to capture any incomplete command.
|
||||
"""
|
||||
ctx = self._sessions.get(session_id)
|
||||
if not ctx or not ctx.buffer.strip():
|
||||
return None
|
||||
|
||||
command = ctx.buffer.strip()
|
||||
ctx.buffer = ""
|
||||
return await self._evaluate_and_log(ctx, command)
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""Get statistics about command logging."""
|
||||
total_logged = sum(ctx.commands_logged for ctx in self._sessions.values())
|
||||
total_blocked = sum(ctx.commands_blocked for ctx in self._sessions.values())
|
||||
|
||||
return {
|
||||
"active_sessions": len(self._sessions),
|
||||
"total_commands_logged": total_logged,
|
||||
"total_commands_blocked": total_blocked,
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Global instance
|
||||
# ============================================================================
|
||||
_command_logger_instance: Optional[TerminalCommandLogger] = None
|
||||
|
||||
|
||||
def get_command_logger() -> TerminalCommandLogger:
|
||||
"""Get or create the global command logger instance."""
|
||||
global _command_logger_instance
|
||||
|
||||
if _command_logger_instance is None:
|
||||
_command_logger_instance = TerminalCommandLogger()
|
||||
|
||||
return _command_logger_instance
|
||||
@@ -0,0 +1,396 @@
|
||||
"""
|
||||
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
|
||||
@@ -6,6 +6,8 @@ This service handles:
|
||||
- Spawning ttyd processes for SSH connections
|
||||
- Managing session lifecycle (creation, expiration, cleanup)
|
||||
- Port allocation for ttyd instances
|
||||
- Garbage collection of stale/idle sessions
|
||||
- Session reuse for same user/host/mode
|
||||
"""
|
||||
import asyncio
|
||||
import hashlib
|
||||
@@ -16,22 +18,39 @@ import shutil
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Dict, Optional, Tuple
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Any, Tuple, Optional
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Configuration
|
||||
# Configuration - Session TTL and limits
|
||||
TERMINAL_SESSION_TTL_MINUTES = int(os.environ.get("TERMINAL_SESSION_TTL_MINUTES", "30"))
|
||||
TERMINAL_SESSION_TTL_SECONDS = TERMINAL_SESSION_TTL_MINUTES * 60
|
||||
TERMINAL_MAX_SESSIONS_PER_USER = int(os.environ.get("TERMINAL_MAX_SESSIONS_PER_USER", "3"))
|
||||
|
||||
# Idle timeout - sessions without heartbeat for this long are considered dead
|
||||
TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS = int(os.environ.get("TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS", "120"))
|
||||
|
||||
# Heartbeat interval - how often client should send heartbeat (slightly less than idle timeout)
|
||||
TERMINAL_HEARTBEAT_INTERVAL_SECONDS = int(os.environ.get("TERMINAL_HEARTBEAT_INTERVAL_SECONDS", "15"))
|
||||
|
||||
# GC interval - how often to run garbage collection
|
||||
TERMINAL_GC_INTERVAL_SECONDS = int(os.environ.get("TERMINAL_GC_INTERVAL_SECONDS", "30"))
|
||||
|
||||
# Port range for ttyd instances
|
||||
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"))
|
||||
|
||||
# SSH configuration
|
||||
SSH_USER = os.environ.get("TERMINAL_SSH_USER", "automation")
|
||||
TTYD_PATH = os.environ.get("TTYD_PATH", "ttyd")
|
||||
|
||||
# ttyd bind interface (important for WSL/remote access)
|
||||
TERMINAL_TTYD_INTERFACE = os.environ.get("TERMINAL_TTYD_INTERFACE")
|
||||
|
||||
|
||||
class TerminalServiceError(Exception):
|
||||
"""Base exception for terminal service errors."""
|
||||
@@ -61,7 +80,12 @@ class TerminalService:
|
||||
- Generating secure session tokens
|
||||
- Allocating ports for ttyd instances
|
||||
- Spawning and managing ttyd processes
|
||||
- Cleaning up expired sessions
|
||||
- Cleaning up expired/idle sessions via GC
|
||||
- Session reuse for same user/host/mode
|
||||
|
||||
The GC task runs periodically to clean up:
|
||||
- Sessions that exceeded their TTL
|
||||
- Sessions that haven't received a heartbeat within idle timeout
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
@@ -73,6 +97,20 @@ class TerminalService:
|
||||
self._lock = asyncio.Lock()
|
||||
# Check if ttyd is available
|
||||
self._ttyd_available: Optional[bool] = None
|
||||
# GC task handle
|
||||
self._gc_task: Optional[asyncio.Task] = None
|
||||
self._gc_running: bool = False
|
||||
# Database session factory for GC (set during app startup)
|
||||
self._db_session_factory = None
|
||||
# Metrics for observability
|
||||
self._metrics = {
|
||||
"sessions_created": 0,
|
||||
"sessions_reused": 0,
|
||||
"sessions_closed_user": 0,
|
||||
"sessions_gc_expired": 0,
|
||||
"sessions_gc_idle": 0,
|
||||
"session_limit_hits": 0,
|
||||
}
|
||||
|
||||
def check_ttyd_available(self) -> bool:
|
||||
"""Check if ttyd binary is available."""
|
||||
@@ -87,6 +125,108 @@ class TerminalService:
|
||||
|
||||
return self._ttyd_available
|
||||
|
||||
def set_db_session_factory(self, factory):
|
||||
"""Set the database session factory for GC operations."""
|
||||
self._db_session_factory = factory
|
||||
|
||||
async def start_gc_task(self):
|
||||
"""Start the garbage collection background task."""
|
||||
if self._gc_task is not None:
|
||||
return
|
||||
|
||||
self._gc_running = True
|
||||
self._gc_task = asyncio.create_task(self._gc_loop())
|
||||
logger.info(f"Terminal session GC started (interval={TERMINAL_GC_INTERVAL_SECONDS}s, idle_timeout={TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS}s)")
|
||||
|
||||
async def stop_gc_task(self):
|
||||
"""Stop the garbage collection background task."""
|
||||
self._gc_running = False
|
||||
if self._gc_task:
|
||||
self._gc_task.cancel()
|
||||
try:
|
||||
await self._gc_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._gc_task = None
|
||||
logger.info("Terminal session GC stopped")
|
||||
|
||||
async def _gc_loop(self):
|
||||
"""Background loop that periodically cleans up stale sessions."""
|
||||
while self._gc_running:
|
||||
try:
|
||||
await asyncio.sleep(TERMINAL_GC_INTERVAL_SECONDS)
|
||||
if self._db_session_factory:
|
||||
await self._run_gc_cycle()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.exception(f"Error in terminal GC cycle: {e}")
|
||||
|
||||
async def _run_gc_cycle(self):
|
||||
"""Run a single GC cycle to clean up stale sessions."""
|
||||
from app.crud.terminal_session import TerminalSessionRepository
|
||||
from app.models.terminal_session import CLOSE_REASON_TTL, CLOSE_REASON_IDLE
|
||||
|
||||
async with self._db_session_factory() as db_session:
|
||||
repo = TerminalSessionRepository(db_session)
|
||||
|
||||
# Get all stale sessions (expired or idle)
|
||||
stale_sessions = await repo.list_stale_sessions(
|
||||
ttl_seconds=TERMINAL_SESSION_TTL_SECONDS,
|
||||
idle_timeout_seconds=TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS
|
||||
)
|
||||
|
||||
if not stale_sessions:
|
||||
return
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
expired_count = 0
|
||||
idle_count = 0
|
||||
|
||||
for session in stale_sessions:
|
||||
# Determine reason: TTL expired or idle
|
||||
expires_at = self._to_utc_aware(session.expires_at)
|
||||
is_expired = bool(expires_at and expires_at <= now)
|
||||
|
||||
# Terminate the ttyd process
|
||||
await self.terminate_session(session.id)
|
||||
# Release the port
|
||||
await self.release_port(session.ttyd_port)
|
||||
|
||||
# Mark in DB with appropriate reason
|
||||
if is_expired:
|
||||
await repo.mark_expired(session.id)
|
||||
expired_count += 1
|
||||
self._metrics["sessions_gc_expired"] += 1
|
||||
logger.info(f"session_gc_closed session={session.id[:8]}... host={session.host_name} reason=ttl")
|
||||
else:
|
||||
await repo.mark_idle(session.id)
|
||||
idle_count += 1
|
||||
self._metrics["sessions_gc_idle"] += 1
|
||||
logger.info(f"session_gc_closed session={session.id[:8]}... host={session.host_name} reason=idle")
|
||||
|
||||
await db_session.commit()
|
||||
|
||||
if expired_count or idle_count:
|
||||
logger.info(f"Terminal GC cycle: cleaned {expired_count} expired, {idle_count} idle sessions")
|
||||
|
||||
def _to_utc_aware(self, dt: Optional[datetime]) -> Optional[datetime]:
|
||||
"""Normalize datetimes from DB (SQLite may return naive) into UTC-aware."""
|
||||
if dt is None:
|
||||
return None
|
||||
if dt.tzinfo is None:
|
||||
return dt.replace(tzinfo=timezone.utc)
|
||||
return dt.astimezone(timezone.utc)
|
||||
|
||||
def get_metrics(self) -> Dict[str, Any]:
|
||||
"""Get service metrics for observability."""
|
||||
return {
|
||||
**self._metrics,
|
||||
"active_processes": len(self._processes),
|
||||
"allocated_ports": len(self._allocated_ports),
|
||||
"gc_running": self._gc_running,
|
||||
}
|
||||
|
||||
def generate_session_id(self) -> str:
|
||||
"""Generate a unique session ID."""
|
||||
return secrets.token_hex(32)
|
||||
@@ -121,6 +261,14 @@ class TerminalService:
|
||||
sock.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def is_session_process_alive(self, session_id: str) -> bool:
|
||||
"""Check if the ttyd process for a session is still running."""
|
||||
async with self._lock:
|
||||
process = self._processes.get(session_id)
|
||||
if process is None:
|
||||
return False
|
||||
return process.poll() is None
|
||||
|
||||
async def allocate_port(self, session_id: str) -> int:
|
||||
"""
|
||||
@@ -217,6 +365,8 @@ class TerminalService:
|
||||
f"{SSH_USER}@{host_ip}",
|
||||
]
|
||||
|
||||
ttyd_interface = TERMINAL_TTYD_INTERFACE
|
||||
|
||||
# Note: --credential removed to avoid HTTP basic auth popup
|
||||
# Security is handled by session token validation in the connect/popout routes
|
||||
cmd = [
|
||||
@@ -224,8 +374,10 @@ class TerminalService:
|
||||
"--once",
|
||||
f"--port={port}",
|
||||
"--writable",
|
||||
*ssh_cmd,
|
||||
]
|
||||
if ttyd_interface:
|
||||
cmd.extend(["--interface", ttyd_interface])
|
||||
cmd.extend(ssh_cmd)
|
||||
|
||||
logger.info(
|
||||
f"Spawning ttyd for session {session_id[:8]}... on port {port} -> {SSH_USER}@{host_ip} (key={ssh_key_path})"
|
||||
@@ -328,10 +480,69 @@ class TerminalService:
|
||||
# Mark as expired in DB
|
||||
await repo.mark_expired(session.id)
|
||||
count += 1
|
||||
logger.info(f"Cleaned up expired session {session.id[:8]}... for host {session.host_name}")
|
||||
self._metrics["sessions_gc_expired"] += 1
|
||||
logger.info(f"session_gc_closed session={session.id[:8]}... host={session.host_name} reason=ttl")
|
||||
|
||||
return count
|
||||
|
||||
async def cleanup_idle_sessions(self, db_session) -> int:
|
||||
"""
|
||||
Clean up idle terminal sessions (no heartbeat within timeout).
|
||||
|
||||
Returns:
|
||||
Number of sessions cleaned up
|
||||
"""
|
||||
from app.crud.terminal_session import TerminalSessionRepository
|
||||
|
||||
repo = TerminalSessionRepository(db_session)
|
||||
idle_sessions = await repo.list_idle(TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS)
|
||||
|
||||
count = 0
|
||||
for session in idle_sessions:
|
||||
# Terminate the ttyd process
|
||||
await self.terminate_session(session.id)
|
||||
# Release the port
|
||||
await self.release_port(session.ttyd_port)
|
||||
# Mark as idle in DB
|
||||
await repo.mark_idle(session.id)
|
||||
count += 1
|
||||
self._metrics["sessions_gc_idle"] += 1
|
||||
logger.info(f"session_gc_closed session={session.id[:8]}... host={session.host_name} reason=idle")
|
||||
|
||||
return count
|
||||
|
||||
async def close_session_with_cleanup(
|
||||
self,
|
||||
session_id: str,
|
||||
port: int,
|
||||
reason: str = "user_close"
|
||||
) -> bool:
|
||||
"""
|
||||
Close a session with full cleanup (process + port).
|
||||
|
||||
Returns:
|
||||
True if session was cleaned up
|
||||
"""
|
||||
terminated = await self.terminate_session(session_id)
|
||||
await self.release_port(port)
|
||||
|
||||
if reason == "user_close":
|
||||
self._metrics["sessions_closed_user"] += 1
|
||||
|
||||
logger.info(f"session_closed session={session_id[:8]}... reason={reason}")
|
||||
return terminated
|
||||
|
||||
def record_session_created(self, reused: bool = False):
|
||||
"""Record a session creation for metrics."""
|
||||
if reused:
|
||||
self._metrics["sessions_reused"] += 1
|
||||
else:
|
||||
self._metrics["sessions_created"] += 1
|
||||
|
||||
def record_session_limit_hit(self):
|
||||
"""Record a session limit hit for metrics."""
|
||||
self._metrics["session_limit_hits"] += 1
|
||||
|
||||
def get_active_session_count(self) -> int:
|
||||
"""Get the number of active ttyd processes."""
|
||||
return len(self._processes)
|
||||
@@ -366,3 +577,21 @@ terminal_service = TerminalService()
|
||||
def get_terminal_service() -> TerminalService:
|
||||
"""Get the global terminal service instance."""
|
||||
return terminal_service
|
||||
|
||||
|
||||
# Export configuration for use in routes
|
||||
__all__ = [
|
||||
"terminal_service",
|
||||
"get_terminal_service",
|
||||
"TerminalService",
|
||||
"TerminalServiceError",
|
||||
"HostNotReadyError",
|
||||
"SessionLimitExceededError",
|
||||
"TtydNotAvailableError",
|
||||
"TERMINAL_SESSION_TTL_MINUTES",
|
||||
"TERMINAL_SESSION_TTL_SECONDS",
|
||||
"TERMINAL_MAX_SESSIONS_PER_USER",
|
||||
"TERMINAL_SESSION_IDLE_TIMEOUT_SECONDS",
|
||||
"TERMINAL_HEARTBEAT_INTERVAL_SECONDS",
|
||||
"TERMINAL_GC_INTERVAL_SECONDS",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user