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

sshmcp/core/container.py

@@ -0,0 +1,291 @@
+"""Dependency injection container for sshmcp."""
+
+import threading
+from dataclasses import dataclass
+from typing import Any, Callable, Generic, TypeVar
+
+import structlog
+
+from sshmcp.models.machine import MachinesConfig
+from sshmcp.security.audit import AuditLogger
+from sshmcp.security.rate_limiter import RateLimiter
+from sshmcp.ssh.pool import SSHConnectionPool
+from sshmcp.ssh.shell import ShellManager
+
+logger = structlog.get_logger()
+
+T = TypeVar("T")
+
+
+class Provider(Generic[T]):
+    """Lazy provider for dependency instances."""
+
+    def __init__(self, factory: Callable[[], T]) -> None:
+        """
+        Initialize provider.
+
+        Args:
+            factory: Factory function to create instance.
+        """
+        self._factory = factory
+        self._instance: T | None = None
+        self._lock = threading.Lock()
+
+    def get(self) -> T:
+        """Get or create the instance."""
+        if self._instance is None:
+            with self._lock:
+                if self._instance is None:
+                    self._instance = self._factory()
+        return self._instance
+
+    def reset(self) -> None:
+        """Reset the instance."""
+        with self._lock:
+            self._instance = None
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if instance is initialized."""
+        return self._instance is not None
+
+
+@dataclass
+class ContainerConfig:
+    """Configuration for the container."""
+
+    # Pool settings
+    max_connections_per_host: int = 3
+    idle_timeout: int = 300
+    cleanup_interval: int = 60
+    health_check_interval: int = 30
+    enable_background_cleanup: bool = True
+
+    # Rate limiter settings
+    requests_per_minute: int = 60
+    requests_per_hour: int = 1000
+
+    # Audit settings
+    audit_log_path: str | None = None
+    audit_retention_days: int = 30
+
+
+class Container:
+    """
+    Dependency injection container.
+
+    Manages lifecycle of shared services like connection pool,
+    shell manager, audit logger, rate limiter, etc.
+    """
+
+    def __init__(self, config: ContainerConfig | None = None) -> None:
+        """
+        Initialize container.
+
+        Args:
+            config: Optional configuration.
+        """
+        self.config = config or ContainerConfig()
+        self._providers: dict[str, Provider[Any]] = {}
+        self._machines_config: MachinesConfig | None = None
+        self._lock = threading.Lock()
+
+        self._register_defaults()
+
+    def _register_defaults(self) -> None:
+        """Register default providers."""
+        self._providers["pool"] = Provider(self._create_pool)
+        self._providers["shell_manager"] = Provider(self._create_shell_manager)
+        self._providers["audit_logger"] = Provider(self._create_audit_logger)
+        self._providers["rate_limiter"] = Provider(self._create_rate_limiter)
+
+    def _create_pool(self) -> SSHConnectionPool:
+        """Create connection pool."""
+        pool = SSHConnectionPool(
+            max_connections_per_host=self.config.max_connections_per_host,
+            idle_timeout=self.config.idle_timeout,
+            cleanup_interval=self.config.cleanup_interval,
+            health_check_interval=self.config.health_check_interval,
+            enable_background_cleanup=self.config.enable_background_cleanup,
+        )
+
+        # Register machines if config available
+        if self._machines_config:
+            for machine in self._machines_config.machines:
+                pool.register_machine(machine)
+
+        logger.info("container_pool_created")
+        return pool
+
+    def _create_shell_manager(self) -> ShellManager:
+        """Create shell manager."""
+        logger.info("container_shell_manager_created")
+        return ShellManager()
+
+    def _create_audit_logger(self) -> AuditLogger:
+        """Create audit logger."""
+        logger.info("container_audit_logger_created")
+        return AuditLogger(log_file=self.config.audit_log_path)
+
+    def _create_rate_limiter(self) -> RateLimiter:
+        """Create rate limiter."""
+        from sshmcp.security.rate_limiter import RateLimitConfig
+
+        logger.info("container_rate_limiter_created")
+        config = RateLimitConfig(
+            requests_per_minute=self.config.requests_per_minute,
+            requests_per_hour=self.config.requests_per_hour,
+        )
+        return RateLimiter(config=config)
+
+    def set_machines_config(self, config: MachinesConfig) -> None:
+        """
+        Set machines configuration.
+
+        Args:
+            config: Machines configuration.
+        """
+        self._machines_config = config
+
+        # If pool already exists, register machines
+        if self._providers["pool"].is_initialized:
+            pool = self.pool
+            for machine in config.machines:
+                pool.register_machine(machine)
+
+    @property
+    def pool(self) -> SSHConnectionPool:
+        """Get connection pool."""
+        return self._providers["pool"].get()
+
+    @property
+    def shell_manager(self) -> ShellManager:
+        """Get shell manager."""
+        return self._providers["shell_manager"].get()
+
+    @property
+    def audit_logger(self) -> AuditLogger:
+        """Get audit logger."""
+        return self._providers["audit_logger"].get()
+
+    @property
+    def rate_limiter(self) -> RateLimiter:
+        """Get rate limiter."""
+        return self._providers["rate_limiter"].get()
+
+    def register(self, name: str, factory: Callable[[], Any]) -> None:
+        """
+        Register a custom provider.
+
+        Args:
+            name: Provider name.
+            factory: Factory function.
+        """
+        with self._lock:
+            self._providers[name] = Provider(factory)
+
+    def get(self, name: str) -> Any:
+        """
+        Get a service by name.
+
+        Args:
+            name: Service name.
+
+        Returns:
+            Service instance.
+
+        Raises:
+            KeyError: If service not registered.
+        """
+        if name not in self._providers:
+            raise KeyError(f"Service not registered: {name}")
+        return self._providers[name].get()
+
+    def reset(self, name: str | None = None) -> None:
+        """
+        Reset service(s).
+
+        Args:
+            name: Service name to reset, or None to reset all.
+        """
+        with self._lock:
+            if name:
+                if name in self._providers:
+                    self._providers[name].reset()
+            else:
+                for provider in self._providers.values():
+                    provider.reset()
+
+    def shutdown(self) -> None:
+        """Shutdown all services."""
+        logger.info("container_shutting_down")
+
+        # Shutdown pool
+        if self._providers["pool"].is_initialized:
+            self.pool.shutdown()
+
+        # Close shell sessions
+        if self._providers["shell_manager"].is_initialized:
+            self.shell_manager.close_all()
+
+        # Reset all
+        self.reset()
+        logger.info("container_shutdown_complete")
+
+    def __enter__(self) -> "Container":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Context manager exit."""
+        self.shutdown()
+
+
+# Global container instance
+_container: Container | None = None
+_container_lock = threading.Lock()
+
+
+def get_container() -> Container:
+    """Get or create the global container."""
+    global _container
+    with _container_lock:
+        if _container is None:
+            _container = Container()
+        return _container
+
+
+def init_container(
+    config: ContainerConfig | None = None,
+    machines: MachinesConfig | None = None,
+) -> Container:
+    """
+    Initialize the global container.
+
+    Args:
+        config: Container configuration.
+        machines: Machines configuration.
+
+    Returns:
+        Initialized Container.
+    """
+    global _container
+    with _container_lock:
+        if _container is not None:
+            _container.shutdown()
+
+        _container = Container(config=config)
+
+        if machines:
+            _container.set_machines_config(machines)
+
+        return _container
+
+
+def reset_container() -> None:
+    """Reset the global container."""
+    global _container
+    with _container_lock:
+        if _container is not None:
+            _container.shutdown()
+            _container = None

sshmcp/models/__init__.py

@@ -0,0 +1,15 @@
+"""Pydantic models for SSH MCP Server."""
+
+from sshmcp.models.command import CommandError, CommandResult
+from sshmcp.models.file import FileContent, FileInfo
+from sshmcp.models.machine import AuthConfig, MachineConfig, SecurityConfig
+
+__all__ = [
+    "MachineConfig",
+    "AuthConfig",
+    "SecurityConfig",
+    "CommandResult",
+    "CommandError",
+    "FileInfo",
+    "FileContent",
+]

sshmcp/models/command.py

@@ -0,0 +1,69 @@
+"""Pydantic models for command execution results."""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class CommandResult(BaseModel):
+    """Result of command execution on remote server."""
+
+    exit_code: int = Field(description="Command exit code (0 = success)")
+    stdout: str = Field(default="", description="Standard output")
+    stderr: str = Field(default="", description="Standard error output")
+    duration_ms: int = Field(ge=0, description="Execution duration in milliseconds")
+    host: str = Field(description="Host where command was executed")
+    command: str = Field(description="Executed command")
+    timestamp: datetime = Field(
+        default_factory=datetime.utcnow, description="Execution timestamp"
+    )
+
+    @property
+    def success(self) -> bool:
+        """Check if command executed successfully."""
+        return self.exit_code == 0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        return {
+            "exit_code": self.exit_code,
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+            "duration_ms": self.duration_ms,
+            "host": self.host,
+            "command": self.command,
+            "success": self.success,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+
+class CommandError(BaseModel):
+    """Error details for failed command execution."""
+
+    error_type: str = Field(description="Type of error")
+    message: str = Field(description="Error message")
+    host: str | None = Field(default=None, description="Host where error occurred")
+    command: str | None = Field(default=None, description="Command that caused error")
+    details: dict[str, Any] | None = Field(
+        default=None, description="Additional error details"
+    )
+    timestamp: datetime = Field(
+        default_factory=datetime.utcnow, description="Error timestamp"
+    )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        result = {
+            "error": True,
+            "error_type": self.error_type,
+            "message": self.message,
+            "timestamp": self.timestamp.isoformat(),
+        }
+        if self.host:
+            result["host"] = self.host
+        if self.command:
+            result["command"] = self.command
+        if self.details:
+            result["details"] = self.details
+        return result

sshmcp/models/file.py

@@ -0,0 +1,102 @@
+"""Pydantic models for file operations."""
+
+from datetime import datetime
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class FileInfo(BaseModel):
+    """Information about a file on remote server."""
+
+    name: str = Field(description="File name")
+    path: str = Field(description="Full path to file")
+    type: Literal["file", "directory", "link", "other"] = Field(description="File type")
+    size: int = Field(ge=0, description="File size in bytes")
+    modified: datetime = Field(description="Last modification time")
+    permissions: str | None = Field(
+        default=None, description="File permissions (e.g., '0644')"
+    )
+    owner: str | None = Field(default=None, description="File owner")
+    group: str | None = Field(default=None, description="File group")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        result = {
+            "name": self.name,
+            "path": self.path,
+            "type": self.type,
+            "size": self.size,
+            "modified": self.modified.isoformat(),
+        }
+        if self.permissions:
+            result["permissions"] = self.permissions
+        if self.owner:
+            result["owner"] = self.owner
+        if self.group:
+            result["group"] = self.group
+        return result
+
+
+class FileContent(BaseModel):
+    """Content of a file from remote server."""
+
+    content: str = Field(description="File content")
+    path: str = Field(description="Full path to file")
+    size: int = Field(ge=0, description="File size in bytes")
+    encoding: str = Field(default="utf-8", description="File encoding")
+    truncated: bool = Field(
+        default=False, description="Whether content was truncated due to size limit"
+    )
+    host: str = Field(description="Host where file is located")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        return {
+            "content": self.content,
+            "path": self.path,
+            "size": self.size,
+            "encoding": self.encoding,
+            "truncated": self.truncated,
+            "host": self.host,
+        }
+
+
+class FileUploadResult(BaseModel):
+    """Result of file upload operation."""
+
+    success: bool = Field(description="Whether upload was successful")
+    path: str = Field(description="Path where file was uploaded")
+    size: int = Field(ge=0, description="Uploaded file size in bytes")
+    host: str = Field(description="Host where file was uploaded")
+    message: str | None = Field(default=None, description="Additional message")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        result = {
+            "success": self.success,
+            "path": self.path,
+            "size": self.size,
+            "host": self.host,
+        }
+        if self.message:
+            result["message"] = self.message
+        return result
+
+
+class FileListResult(BaseModel):
+    """Result of listing files in directory."""
+
+    files: list[FileInfo] = Field(description="List of files")
+    directory: str = Field(description="Listed directory path")
+    host: str = Field(description="Host where files are located")
+    total_count: int = Field(ge=0, description="Total number of files")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for MCP response."""
+        return {
+            "files": [f.to_dict() for f in self.files],
+            "directory": self.directory,
+            "host": self.host,
+            "total_count": self.total_count,
+        }

sshmcp/models/machine.py

@@ -0,0 +1,139 @@
+"""Pydantic models for machine configuration."""
+
+import os
+from typing import Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class AuthConfig(BaseModel):
+    """Authentication configuration for SSH connection."""
+
+    type: Literal["key", "password", "agent"] = Field(
+        description="Authentication type: 'key' for SSH key, 'password' for password, 'agent' for SSH agent"
+    )
+    key_path: str | None = Field(
+        default=None, description="Path to SSH private key file"
+    )
+    passphrase: str | None = Field(
+        default=None, description="Passphrase for encrypted SSH key"
+    )
+    password: str | None = Field(
+        default=None, description="Password for password authentication"
+    )
+    agent_forwarding: bool = Field(
+        default=False, description="Enable SSH agent forwarding"
+    )
+
+    @field_validator("key_path")
+    @classmethod
+    def expand_key_path(cls, v: str | None) -> str | None:
+        """Expand ~ in key path."""
+        if v:
+            return os.path.expanduser(v)
+        return v
+
+    def model_post_init(self, __context: object) -> None:
+        """Validate that required fields are present based on auth type."""
+        if self.type == "key" and not self.key_path:
+            raise ValueError("key_path is required when auth type is 'key'")
+        if self.type == "password" and not self.password:
+            raise ValueError("password is required when auth type is 'password'")
+        # 'agent' type doesn't require additional fields
+
+
+class SecurityConfig(BaseModel):
+    """Security configuration for command execution."""
+
+    allowed_commands: list[str] = Field(
+        default_factory=list,
+        description="List of regex patterns for allowed commands",
+    )
+    forbidden_commands: list[str] = Field(
+        default_factory=lambda: [
+            r".*rm\s+-rf\s+/.*",
+            r".*sudo\s+.*",
+            r".*su\s+-.*",
+            r".*dd\s+if=.*",
+            r".*mkfs\..*",
+            r".*:\(\)\{.*",  # Fork bomb
+        ],
+        description="List of regex patterns for forbidden commands",
+    )
+    timeout_seconds: int = Field(
+        default=30, ge=1, le=3600, description="Command timeout in seconds"
+    )
+    max_concurrent_commands: int = Field(
+        default=3, ge=1, le=10, description="Maximum concurrent commands per machine"
+    )
+    allowed_paths: list[str] = Field(
+        default_factory=list,
+        description="List of allowed file paths for read/write operations",
+    )
+    forbidden_paths: list[str] = Field(
+        default_factory=lambda: [
+            "/etc/passwd",
+            "/etc/shadow",
+            "/etc/sudoers",
+            "/root",
+            "~/.ssh",
+        ],
+        description="List of forbidden file paths",
+    )
+
+
+class MachineConfig(BaseModel):
+    """Configuration for a VPS machine."""
+
+    name: str = Field(
+        min_length=1, max_length=64, description="Unique name for the machine"
+    )
+    host: str = Field(min_length=1, description="Hostname or IP address")
+    port: int = Field(default=22, ge=1, le=65535, description="SSH port")
+    user: str = Field(min_length=1, description="SSH username")
+    auth: AuthConfig = Field(description="Authentication configuration")
+    security: SecurityConfig = Field(
+        default_factory=SecurityConfig, description="Security configuration"
+    )
+    description: str | None = Field(
+        default=None, description="Optional description of the machine"
+    )
+    tags: list[str] = Field(
+        default_factory=list,
+        description="Tags for grouping and filtering servers (e.g., ['production', 'web'])",
+    )
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        """Validate machine name contains only safe characters."""
+        import re
+
+        if not re.match(r"^[a-zA-Z0-9_-]+$", v):
+            raise ValueError(
+                "Machine name must contain only alphanumeric characters, underscores, and hyphens"
+            )
+        return v
+
+
+class MachinesConfig(BaseModel):
+    """Root configuration containing all machines."""
+
+    machines: list[MachineConfig] = Field(
+        default_factory=list, description="List of configured machines"
+    )
+
+    def get_machine(self, name: str) -> MachineConfig | None:
+        """Get machine by name."""
+        for machine in self.machines:
+            if machine.name == name:
+                return machine
+        return None
+
+    def has_machine(self, name: str) -> bool:
+        """Check if machine exists."""
+        return self.get_machine(name) is not None
+
+    def get_machine_names(self) -> list[str]:
+        """Get list of all machine names."""
+        return [m.name for m in self.machines]