"""
SSH Connection Manager for Pentest MCP Server.
Handles SSH connections with auto-reconnection and health checks.
"""
import asyncio
import logging
import time
from typing import Optional, Dict, Any
import asyncssh
from .config import Config
logger = logging.getLogger(__name__)
class SSHConnectionError(Exception):
"""Custom exception for SSH connection issues."""
pass
class SSHManager:
"""Manages SSH connection to your OS with auto-reconnection."""
def __init__(self, config: Config = None):
self.config = config or Config()
self.connection: Optional[asyncssh.SSHClientConnection] = None
self.last_keepalive = time.time()
self.connection_attempts = 0
self._connecting = False
async def connect(self) -> bool:
"""
Establish SSH connection to your OS.
Returns:
bool: True if connection successful, False otherwise
"""
if self._connecting:
logger.info("Connection attempt already in progress")
return False
self._connecting = True
try:
logger.info(f"Connecting to {self.config.TARGET_HOST}:{self.config.TARGET_PORT}")
# Get authentication method
auth_method = self.config.get_ssh_auth_method()
# Establish connection
self.connection = await asyncssh.connect(
self.config.TARGET_HOST,
port=self.config.TARGET_PORT,
username=self.config.TARGET_USER,
known_hosts=None, # Disable known_hosts checking for pentesting
**auth_method
)
logger.info("SSH connection established successfully")
self.connection_attempts = 0
self.last_keepalive = time.time()
return True
except Exception as e:
logger.error(f"Failed to connect: {e}")
self.connection = None
raise SSHConnectionError(f"Connection failed: {e}")
finally:
self._connecting = False
async def disconnect(self):
"""Disconnect from SSH."""
if self.connection and not self.connection.is_closed():
self.connection.close()
await self.connection.wait_closed()
self.connection = None
logger.info("SSH connection closed")
async def is_connected(self) -> bool:
"""
Check if SSH connection is alive.
Returns:
bool: True if connected and responsive
"""
if not self.connection or self.connection.is_closed():
return False
try:
# Test connection with a simple command
result = await self.connection.run("echo ping", timeout=5)
return result.exit_status == 0
except Exception as e:
logger.warning(f"Connection check failed: {e}")
return False
async def ensure_connected(self) -> bool:
"""
Ensure SSH connection is alive, reconnect if needed.
Returns:
bool: True if connected after check/reconnect
"""
if await self.is_connected():
return True
logger.info("Connection lost, attempting to reconnect...")
return await self.reconnect()
async def reconnect(self) -> bool:
"""
Reconnect with exponential backoff.
Returns:
bool: True if reconnection successful
"""
max_attempts = self.config.MAX_RECONNECT_ATTEMPTS
base_delay = self.config.RECONNECT_DELAY_BASE
for attempt in range(max_attempts):
self.connection_attempts = attempt + 1
try:
if self.connection:
await self.disconnect()
# Exponential backoff delay
if attempt > 0:
delay = base_delay ** attempt
logger.info(f"Waiting {delay}s before reconnect attempt {attempt + 1}/{max_attempts}")
await asyncio.sleep(delay)
# Attempt connection
await self.connect()
logger.info(f"Reconnection successful after {attempt + 1} attempts")
return True
except Exception as e:
logger.error(f"Reconnect attempt {attempt + 1}/{max_attempts} failed: {e}")
logger.error(f"Failed to reconnect after {max_attempts} attempts")
return False
async def run_command(self, command: str, timeout: int = None) -> asyncssh.SSHCompletedProcess:
"""
Execute a command via SSH.
Args:
command: Command to execute
timeout: Command timeout in seconds
Returns:
SSHCompletedProcess: Command result
Raises:
SSHConnectionError: If connection fails
"""
if not await self.ensure_connected():
raise SSHConnectionError("Unable to establish SSH connection")
timeout = timeout or self.config.DEFAULT_TIMEOUT
try:
result = await self.connection.run(command, timeout=timeout)
self.last_keepalive = time.time()
return result
except Exception as e:
logger.error(f"Command execution failed: {e}")
raise SSHConnectionError(f"Command failed: {e}")
async def create_process(self, command: str) -> asyncssh.SSHClientProcess:
"""
Create an interactive SSH process.
Args:
command: Command to start
Returns:
SSHClientProcess: Interactive process
"""
if not await self.ensure_connected():
raise SSHConnectionError("Unable to establish SSH connection")
try:
process = await self.connection.create_process(command)
return process
except Exception as e:
logger.error(f"Process creation failed: {e}")
raise SSHConnectionError(f"Process creation failed: {e}")
async def start_sftp_client(self) -> asyncssh.SFTPClient:
"""
Start SFTP client for file operations.
Returns:
SFTPClient: SFTP client instance
"""
if not await self.ensure_connected():
raise SSHConnectionError("Unable to establish SSH connection")
try:
sftp = await self.connection.start_sftp_client()
return sftp
except Exception as e:
logger.error(f"SFTP client creation failed: {e}")
raise SSHConnectionError(f"SFTP failed: {e}")
async def get_system_info(self) -> Dict[str, Any]:
"""
Get system information with structured, machine-readable format.
Returns:
dict: System information with detailed memory and disk data
"""
try:
# Get basic system info
basic_commands = {
"hostname": "hostname",
"uptime": "uptime",
"kernel": "uname -r",
"cpu_info": "nproc",
"disk": "df -h / | awk 'NR==2{print $5,$3,$2}'"
}
info = {}
# Use configurable timeout instead of hardcoded value
timeout = getattr(self.config, 'DEFAULT_TIMEOUT', 30)
for key, cmd in basic_commands.items():
try:
result = await self.run_command(cmd, timeout=timeout)
info[key] = result.stdout.strip()
except Exception as e:
logger.warning(f"Failed to get {key}: {e}")
info[key] = "unknown"
# Get structured memory information
try:
# Get memory values: total, used, available (in MB)
mem_result = await self.run_command(
"free -m | awk 'NR==2{print $2,$3,$7}'", timeout=timeout
)
mem_values = mem_result.stdout.strip().split()
if len(mem_values) != 3:
raise ValueError(f"Expected 3 memory values, got {len(mem_values)}: {mem_values}")
total_mb, used_mb, available_mb = map(int, mem_values)
# Convert MB to GB (1024 MB = 1 GB)
mb_to_gb = 1024
usage_percent = round((used_mb / total_mb) * 100, 1)
used_gb = round(used_mb / mb_to_gb, 1)
total_gb = round(total_mb / mb_to_gb, 1)
available_gb = round(available_mb / mb_to_gb, 1)
memory_info = {
"usage_percent": usage_percent,
"used_mb": used_mb,
"total_mb": total_mb,
"available_mb": available_mb,
"used_gb": used_gb,
"total_gb": total_gb,
"available_gb": available_gb,
"summary": f"{used_gb}GB used of {total_gb}GB total ({usage_percent}% used, {available_gb}GB available)"
}
info["memory"] = memory_info
except Exception as e:
logger.warning(f"Failed to get structured memory info: {e}")
# Fallback to old format if structured parsing fails
try:
result = await self.run_command(
"free -m | awk 'NR==2{printf \"%.1f,%.1f,%.1f\", $3*100/$2, $3, $2}'", timeout=timeout
)
info["memory"] = result.stdout.strip() + " (legacy format)"
except Exception:
info["memory"] = "unknown"
# Improve disk information formatting
if "disk" in info and info["disk"] != "unknown":
try:
# Parse disk info format: "percentage used total" -> structured format
disk_parts = info["disk"].split()
if len(disk_parts) == 3:
usage_percent = disk_parts[0]
used = disk_parts[1]
total = disk_parts[2]
# Calculate available using safer method
try:
def parse_size(size_str):
"""Parse size string like '17G', '1.5T', '512M' to GB."""
if not size_str or len(size_str) < 2:
return 0.0
unit = size_str[-1].upper()
value = float(size_str[:-1])
# Convert to GB
if unit == 'T':
return value * 1024 # TB to GB
elif unit == 'G':
return value
elif unit == 'M':
return value / 1024 # MB to GB
elif unit == 'K':
return value / (1024 * 1024) # KB to GB
else:
# Assume bytes
return value / (1024 * 1024 * 1024)
used_gb = parse_size(used)
total_gb = parse_size(total)
available_gb = total_gb - used_gb
# Format available size in appropriate unit
if available_gb >= 1024:
available_str = f"{available_gb/1024:.1f}T"
elif available_gb >= 1:
available_str = f"{available_gb:.1f}G"
else:
available_str = f"{available_gb*1024:.0f}M"
disk_info = {
"usage_percent": usage_percent,
"used": used,
"total": total,
"available": available_str,
"summary": f"Disk: {used} used of {total} total ({usage_percent} full, {available_str} available)"
}
info["disk"] = disk_info
except (ValueError, IndexError) as calc_error:
logger.warning(f"Failed to calculate disk available space: {calc_error}")
# Fallback: structured format without calculated available
disk_info = {
"usage_percent": usage_percent,
"used": used,
"total": total,
"summary": f"Disk: {used} used of {total} total ({usage_percent} full)"
}
info["disk"] = disk_info
except Exception as disk_error:
logger.warning(f"Failed to structure disk info: {disk_error}")
# Keep original format as fallback
return info
except Exception as e:
logger.error(f"Failed to get system info: {e}")
return {"error": str(e)}
def get_connection_status(self) -> Dict[str, Any]:
"""
Get current connection status.
Returns:
dict: Connection status information
"""
return {
"connected": self.connection is not None and not self.connection.is_closed(),
"host": self.config.TARGET_HOST,
"port": self.config.TARGET_PORT,
"user": self.config.TARGET_USER,
"connection_attempts": self.connection_attempts,
"last_keepalive": self.last_keepalive
}
