mcp-ssh-vps 0.4.1__py3-none-any.whl

sshmcp/resources/status.py

@@ -0,0 +1,160 @@
+"""MCP Resource for server status."""
+
+import structlog
+
+from sshmcp.config import get_machine
+from sshmcp.security.audit import get_audit_logger
+from sshmcp.ssh.pool import get_pool
+
+logger = structlog.get_logger()
+
+
+def get_status(host: str) -> dict:
+    """
+    Get status of VPS server.
+
+    Returns server availability, hostname, and service statuses.
+
+    Args:
+        host: Name of the host from machines.json configuration.
+
+    Returns:
+        Dictionary with:
+        - hostname: Server hostname
+        - status: Server status (online/offline)
+        - services: List of service statuses
+        - system_info: Basic system information
+
+    Raises:
+        ValueError: If host not found.
+        RuntimeError: If status cannot be retrieved.
+
+    Example:
+        Resource URI: vps://production-server/status
+    """
+    audit = get_audit_logger()
+
+    # Get machine configuration
+    try:
+        machine = get_machine(host)
+    except Exception as e:
+        raise ValueError(f"Host not found: {host}") from e
+
+    pool = get_pool()
+    pool.register_machine(machine)
+
+    status_info = {
+        "hostname": "",
+        "status": "unknown",
+        "host": host,
+        "services": [],
+        "system_info": {},
+    }
+
+    try:
+        client = pool.get_client(host)
+        try:
+            # Server is online if we can connect
+            status_info["status"] = "online"
+
+            # Get hostname
+            hostname_result = client.execute("hostname")
+            status_info["hostname"] = hostname_result.stdout.strip()
+
+            # Get basic system info
+            uname_result = client.execute("uname -a")
+            status_info["system_info"]["uname"] = uname_result.stdout.strip()
+
+            # Get OS info
+            os_result = client.execute(
+                "cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d= -f2 | tr -d '\"'"
+            )
+            if os_result.exit_code == 0 and os_result.stdout.strip():
+                status_info["system_info"]["os"] = os_result.stdout.strip()
+
+            # Get running services (try systemd first)
+            services_result = client.execute(
+                "systemctl list-units --type=service --state=running --no-pager --no-legend 2>/dev/null | head -20"
+            )
+
+            if services_result.exit_code == 0 and services_result.stdout.strip():
+                status_info["services"] = _parse_systemd_services(
+                    services_result.stdout
+                )
+            else:
+                # Try pm2
+                pm2_result = client.execute("pm2 jlist 2>/dev/null")
+                if pm2_result.exit_code == 0 and pm2_result.stdout.strip():
+                    status_info["services"] = _parse_pm2_services(pm2_result.stdout)
+
+            audit.log(
+                event="status_read",
+                host=host,
+            )
+
+            return status_info
+
+        finally:
+            pool.release_client(client)
+
+    except Exception as e:
+        # If we can't connect, server is offline
+        status_info["status"] = "offline"
+        status_info["error"] = str(e)
+
+        audit.log(
+            event="status_read_failed",
+            host=host,
+            error=str(e),
+        )
+
+        return status_info
+
+
+def _parse_systemd_services(output: str) -> list[dict]:
+    """Parse systemd services from systemctl output."""
+    services = []
+
+    for line in output.strip().split("\n"):
+        if not line.strip():
+            continue
+
+        parts = line.split()
+        if len(parts) >= 1:
+            service_name = parts[0]
+            # Remove .service suffix
+            if service_name.endswith(".service"):
+                service_name = service_name[:-8]
+
+            services.append(
+                {
+                    "name": service_name,
+                    "status": "running",
+                    "manager": "systemd",
+                }
+            )
+
+    return services
+
+
+def _parse_pm2_services(output: str) -> list[dict]:
+    """Parse pm2 services from pm2 jlist output."""
+    import json
+
+    services = []
+
+    try:
+        data = json.loads(output)
+        for app in data:
+            services.append(
+                {
+                    "name": app.get("name", "unknown"),
+                    "status": app.get("pm2_env", {}).get("status", "unknown"),
+                    "manager": "pm2",
+                    "pid": app.get("pid"),
+                }
+            )
+    except json.JSONDecodeError:
+        pass
+
+    return services

sshmcp/security/__init__.py

@@ -0,0 +1,7 @@
+"""Security module for command validation and auditing."""
+
+from sshmcp.security.audit import audit_log
+from sshmcp.security.validator import validate_command, validate_path
+from sshmcp.security.whitelist import CommandWhitelist
+
+__all__ = ["validate_command", "validate_path", "CommandWhitelist", "audit_log"]

sshmcp/security/audit.py

@@ -0,0 +1,314 @@
+"""Audit logging for SSH MCP operations."""
+
+import json
+import os
+import socket
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+import structlog
+
+from sshmcp.security.validator import sanitize_command_for_log
+
+logger = structlog.get_logger()
+
+
+def get_client_info() -> dict[str, str | None]:
+    """
+    Get information about the client making the request.
+
+    Returns:
+        Dictionary with client IP, hostname, and user.
+    """
+    info: dict[str, str | None] = {
+        "client_ip": None,
+        "client_hostname": None,
+        "client_user": None,
+    }
+
+    # Get current user
+    try:
+        info["client_user"] = os.getenv("USER") or os.getenv("USERNAME")
+    except Exception:
+        pass
+
+    # Get local hostname
+    try:
+        info["client_hostname"] = socket.gethostname()
+    except Exception:
+        pass
+
+    # Get SSH client IP if available (from SSH_CLIENT or SSH_CONNECTION env vars)
+    try:
+        ssh_client = os.getenv("SSH_CLIENT")
+        if ssh_client:
+            info["client_ip"] = ssh_client.split()[0]
+        else:
+            ssh_connection = os.getenv("SSH_CONNECTION")
+            if ssh_connection:
+                info["client_ip"] = ssh_connection.split()[0]
+    except Exception:
+        pass
+
+    return info
+
+
+class AuditLogger:
+    """
+    Audit logger for tracking all SSH MCP operations.
+    """
+
+    def __init__(
+        self,
+        log_file: Optional[str] = None,
+        log_to_stdout: bool = True,
+    ) -> None:
+        """
+        Initialize audit logger.
+
+        Args:
+            log_file: Optional path to audit log file.
+            log_to_stdout: Whether to also log to stdout via structlog.
+        """
+        self.log_file = log_file
+        self.log_to_stdout = log_to_stdout
+        self._file_handle = None
+
+        if log_file:
+            path = Path(log_file)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            self._file_handle = open(path, "a", encoding="utf-8")
+
+    def log(
+        self,
+        event: str,
+        host: Optional[str] = None,
+        command: Optional[str] = None,
+        result: Optional[dict[str, Any]] = None,
+        error: Optional[str] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        include_client_info: bool = True,
+    ) -> None:
+        """
+        Log an audit event.
+
+        Args:
+            event: Event type (e.g., 'command_executed', 'file_read').
+            host: Target host name.
+            command: Command that was executed (will be sanitized).
+            result: Result of the operation.
+            error: Error message if operation failed.
+            metadata: Additional metadata to log.
+            include_client_info: Whether to include client IP/user info.
+        """
+        timestamp = datetime.now(timezone.utc).isoformat()
+
+        audit_record: dict[str, Any] = {
+            "timestamp": timestamp,
+            "event": event,
+        }
+
+        # Include client information
+        if include_client_info:
+            client_info = get_client_info()
+            audit_record["client"] = {
+                k: v for k, v in client_info.items() if v is not None
+            }
+
+        # Build log kwargs (without event, as it's passed separately to structlog)
+        log_kwargs: dict[str, Any] = {"timestamp": timestamp}
+
+        if host:
+            audit_record["host"] = host
+            log_kwargs["host"] = host
+
+        if command:
+            # Sanitize command to hide potential secrets
+            sanitized = sanitize_command_for_log(command)
+            audit_record["command"] = sanitized
+            log_kwargs["command"] = sanitized
+
+        if result:
+            # Only include safe result fields
+            safe_result = {
+                k: v
+                for k, v in result.items()
+                if k in ("exit_code", "success", "duration_ms", "size", "path")
+            }
+            audit_record["result"] = safe_result
+            log_kwargs["result"] = safe_result
+
+        if error:
+            audit_record["error"] = error
+            log_kwargs["error"] = error
+
+        if metadata:
+            audit_record["metadata"] = metadata
+            log_kwargs["metadata"] = metadata
+
+        # Log to structlog
+        if self.log_to_stdout:
+            if error:
+                logger.error(event, **log_kwargs)
+            else:
+                logger.info(event, **log_kwargs)
+
+        # Write to file
+        if self._file_handle:
+            self._file_handle.write(json.dumps(audit_record) + "\n")
+            self._file_handle.flush()
+
+    def log_command_executed(
+        self,
+        host: str,
+        command: str,
+        exit_code: int,
+        duration_ms: int,
+    ) -> None:
+        """Log command execution."""
+        self.log(
+            event="command_executed",
+            host=host,
+            command=command,
+            result={
+                "exit_code": exit_code,
+                "duration_ms": duration_ms,
+                "success": exit_code == 0,
+            },
+        )
+
+    def log_command_failed(
+        self,
+        host: str,
+        command: str,
+        error: str,
+    ) -> None:
+        """Log failed command execution."""
+        self.log(
+            event="command_failed",
+            host=host,
+            command=command,
+            error=error,
+        )
+
+    def log_command_rejected(
+        self,
+        host: str,
+        command: str,
+        reason: str,
+    ) -> None:
+        """Log rejected command (security validation failed)."""
+        self.log(
+            event="command_rejected",
+            host=host,
+            command=command,
+            error=reason,
+            metadata={"security_violation": True},
+        )
+
+    def log_file_read(
+        self,
+        host: str,
+        path: str,
+        size: int,
+    ) -> None:
+        """Log file read operation."""
+        self.log(
+            event="file_read",
+            host=host,
+            result={"path": path, "size": size},
+        )
+
+    def log_file_write(
+        self,
+        host: str,
+        path: str,
+        size: int,
+    ) -> None:
+        """Log file write operation."""
+        self.log(
+            event="file_write",
+            host=host,
+            result={"path": path, "size": size},
+        )
+
+    def log_path_rejected(
+        self,
+        host: str,
+        path: str,
+        reason: str,
+    ) -> None:
+        """Log rejected path access."""
+        self.log(
+            event="path_rejected",
+            host=host,
+            error=reason,
+            metadata={"path": path, "security_violation": True},
+        )
+
+    def close(self) -> None:
+        """Close audit log file."""
+        if self._file_handle:
+            self._file_handle.close()
+            self._file_handle = None
+
+
+# Global audit logger instance
+_audit_logger: Optional[AuditLogger] = None
+
+
+def get_audit_logger() -> AuditLogger:
+    """Get or create the global audit logger."""
+    global _audit_logger
+    if _audit_logger is None:
+        _audit_logger = AuditLogger()
+    return _audit_logger
+
+
+def init_audit_logger(
+    log_file: Optional[str] = None,
+    log_to_stdout: bool = True,
+) -> AuditLogger:
+    """
+    Initialize the global audit logger.
+
+    Args:
+        log_file: Optional path to audit log file.
+        log_to_stdout: Whether to also log to stdout.
+
+    Returns:
+        Initialized AuditLogger.
+    """
+    global _audit_logger
+    _audit_logger = AuditLogger(log_file=log_file, log_to_stdout=log_to_stdout)
+    return _audit_logger
+
+
+def audit_log(
+    event: str,
+    host: Optional[str] = None,
+    command: Optional[str] = None,
+    result: Optional[dict[str, Any]] = None,
+    error: Optional[str] = None,
+    metadata: Optional[dict[str, Any]] = None,
+) -> None:
+    """
+    Convenience function to log an audit event.
+
+    Args:
+        event: Event type.
+        host: Target host name.
+        command: Command that was executed.
+        result: Result of the operation.
+        error: Error message if operation failed.
+        metadata: Additional metadata.
+    """
+    get_audit_logger().log(
+        event=event,
+        host=host,
+        command=command,
+        result=result,
+        error=error,
+        metadata=metadata,
+    )

sshmcp/security/rate_limiter.py

@@ -0,0 +1,221 @@
+"""Rate limiting for command execution."""
+
+import threading
+import time
+from collections import defaultdict
+from typing import NamedTuple
+
+import structlog
+
+logger = structlog.get_logger()
+
+
+class RateLimitConfig(NamedTuple):
+    """Rate limit configuration."""
+
+    requests_per_minute: int = 60
+    requests_per_hour: int = 1000
+    burst_size: int = 10
+
+
+class RateLimitExceeded(Exception):
+    """Exception raised when rate limit is exceeded."""
+
+    def __init__(self, message: str, retry_after: float = 0):
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+class TokenBucket:
+    """Token bucket rate limiter implementation."""
+
+    def __init__(self, rate: float, capacity: int):
+        """
+        Initialize token bucket.
+
+        Args:
+            rate: Tokens per second to add.
+            capacity: Maximum tokens in bucket.
+        """
+        self.rate = rate
+        self.capacity = capacity
+        self.tokens = capacity
+        self.last_update = time.monotonic()
+        self._lock = threading.Lock()
+
+    def consume(self, tokens: int = 1) -> bool:
+        """
+        Try to consume tokens from the bucket.
+
+        Args:
+            tokens: Number of tokens to consume.
+
+        Returns:
+            True if tokens were consumed, False if not enough tokens.
+        """
+        with self._lock:
+            now = time.monotonic()
+            elapsed = now - self.last_update
+            self.last_update = now
+
+            # Add tokens based on elapsed time
+            self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
+
+            if self.tokens >= tokens:
+                self.tokens -= tokens
+                return True
+            return False
+
+    def time_until_available(self, tokens: int = 1) -> float:
+        """
+        Calculate time until tokens are available.
+
+        Args:
+            tokens: Number of tokens needed.
+
+        Returns:
+            Seconds until tokens available (0 if available now).
+        """
+        with self._lock:
+            if self.tokens >= tokens:
+                return 0
+            needed = tokens - self.tokens
+            return needed / self.rate
+
+
+class RateLimiter:
+    """Rate limiter for SSH command execution."""
+
+    def __init__(self, config: RateLimitConfig | None = None):
+        """
+        Initialize rate limiter.
+
+        Args:
+            config: Rate limit configuration.
+        """
+        self.config = config or RateLimitConfig()
+        self._buckets: dict[str, dict[str, TokenBucket]] = defaultdict(dict)
+        self._lock = threading.Lock()
+
+    def _get_bucket(self, host: str, limit_type: str) -> TokenBucket:
+        """Get or create a token bucket for host and limit type."""
+        with self._lock:
+            if limit_type not in self._buckets[host]:
+                if limit_type == "minute":
+                    rate = self.config.requests_per_minute / 60.0
+                    capacity = self.config.burst_size
+                elif limit_type == "hour":
+                    rate = self.config.requests_per_hour / 3600.0
+                    capacity = self.config.burst_size * 2
+                else:
+                    rate = 1.0
+                    capacity = 10
+
+                self._buckets[host][limit_type] = TokenBucket(rate, capacity)
+
+            return self._buckets[host][limit_type]
+
+    def check_rate_limit(self, host: str) -> None:
+        """
+        Check if request is within rate limits.
+
+        Args:
+            host: Host name to check rate limit for.
+
+        Raises:
+            RateLimitExceeded: If rate limit is exceeded.
+        """
+        # Check minute limit
+        minute_bucket = self._get_bucket(host, "minute")
+        if not minute_bucket.consume():
+            retry_after = minute_bucket.time_until_available()
+            logger.warning(
+                "rate_limit_exceeded",
+                host=host,
+                limit_type="minute",
+                retry_after=retry_after,
+            )
+            raise RateLimitExceeded(
+                f"Rate limit exceeded for {host}. Try again in {retry_after:.1f}s",
+                retry_after=retry_after,
+            )
+
+        # Check hour limit
+        hour_bucket = self._get_bucket(host, "hour")
+        if not hour_bucket.consume():
+            retry_after = hour_bucket.time_until_available()
+            logger.warning(
+                "rate_limit_exceeded",
+                host=host,
+                limit_type="hour",
+                retry_after=retry_after,
+            )
+            raise RateLimitExceeded(
+                f"Hourly rate limit exceeded for {host}. Try again in {retry_after:.1f}s",
+                retry_after=retry_after,
+            )
+
+    def get_remaining(self, host: str) -> dict:
+        """
+        Get remaining rate limit for host.
+
+        Args:
+            host: Host name.
+
+        Returns:
+            Dictionary with remaining limits.
+        """
+        minute_bucket = self._get_bucket(host, "minute")
+        hour_bucket = self._get_bucket(host, "hour")
+
+        return {
+            "host": host,
+            "minute": {
+                "remaining": int(minute_bucket.tokens),
+                "limit": self.config.requests_per_minute,
+            },
+            "hour": {
+                "remaining": int(hour_bucket.tokens),
+                "limit": self.config.requests_per_hour,
+            },
+        }
+
+    def reset(self, host: str | None = None) -> None:
+        """
+        Reset rate limits.
+
+        Args:
+            host: Host to reset, or None for all hosts.
+        """
+        with self._lock:
+            if host is None:
+                self._buckets.clear()
+            elif host in self._buckets:
+                del self._buckets[host]
+
+
+# Global rate limiter instance
+_rate_limiter: RateLimiter | None = None
+
+
+def get_rate_limiter() -> RateLimiter:
+    """Get or create the global rate limiter."""
+    global _rate_limiter
+    if _rate_limiter is None:
+        _rate_limiter = RateLimiter()
+    return _rate_limiter
+
+
+def init_rate_limiter(config: RateLimitConfig | None = None) -> RateLimiter:
+    """
+    Initialize the global rate limiter with config.
+
+    Args:
+        config: Rate limit configuration.
+
+    Returns:
+        Initialized RateLimiter.
+    """
+    global _rate_limiter
+    _rate_limiter = RateLimiter(config)
+    return _rate_limiter