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

sshmcp/security/totp.py

@@ -0,0 +1,392 @@
+"""TOTP-based two-factor authentication for critical operations."""
+
+import io
+import json
+import os
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable
+
+import structlog
+
+logger = structlog.get_logger()
+
+# Optional imports for TOTP
+try:
+    import pyotp
+
+    PYOTP_AVAILABLE = True
+except ImportError:
+    PYOTP_AVAILABLE = False
+    pyotp = None  # type: ignore
+
+try:
+    import qrcode
+
+    QRCODE_AVAILABLE = True
+except ImportError:
+    QRCODE_AVAILABLE = False
+    qrcode = None  # type: ignore
+
+
+class TOTPError(Exception):
+    """Error in TOTP operations."""
+
+    pass
+
+
+class TOTPNotConfigured(TOTPError):
+    """TOTP is not configured for this host."""
+
+    pass
+
+
+class TOTPVerificationFailed(TOTPError):
+    """TOTP verification failed."""
+
+    pass
+
+
+class TOTPManager:
+    """
+    TOTP-based two-factor authentication manager.
+
+    Provides TOTP setup, verification, and management for critical SSH operations.
+    """
+
+    def __init__(
+        self,
+        secrets_file: str | None = None,
+        issuer: str = "SSH-MCP",
+        digits: int = 6,
+        interval: int = 30,
+    ) -> None:
+        """
+        Initialize TOTP manager.
+
+        Args:
+            secrets_file: Path to encrypted secrets file.
+            issuer: Issuer name for TOTP (shown in authenticator apps).
+            digits: Number of digits in TOTP code.
+            interval: Time interval in seconds.
+        """
+        if not PYOTP_AVAILABLE:
+            raise TOTPError(
+                "pyotp package is required for TOTP. Install with: pip install pyotp"
+            )
+
+        self.secrets_file = secrets_file or str(
+            Path.home() / ".sshmcp" / "totp_secrets.json"
+        )
+        self.issuer = issuer
+        self.digits = digits
+        self.interval = interval
+
+        self._secrets: dict[str, dict] = {}
+        self._lock = threading.Lock()
+        self._critical_commands: list[str] = [
+            r".*rm\s+-rf.*",
+            r".*shutdown.*",
+            r".*reboot.*",
+            r".*systemctl\s+(stop|disable).*",
+            r".*docker\s+(rm|stop|kill).*",
+            r".*DROP\s+DATABASE.*",
+            r".*DROP\s+TABLE.*",
+        ]
+        self._verification_callbacks: list[Callable[[str, bool], None]] = []
+
+        self._load_secrets()
+
+    def setup_totp(self, host: str, account_name: str | None = None) -> dict:
+        """
+        Set up TOTP for a host.
+
+        Args:
+            host: Host name to set up TOTP for.
+            account_name: Account name (default: host name).
+
+        Returns:
+            Dictionary with secret, provisioning URI, and optional QR code.
+        """
+        account = account_name or host
+        secret = pyotp.random_base32()
+
+        totp = pyotp.TOTP(secret, digits=self.digits, interval=self.interval)
+        uri = totp.provisioning_uri(name=account, issuer_name=self.issuer)
+
+        result = {
+            "host": host,
+            "secret": secret,
+            "uri": uri,
+            "account": account,
+            "issuer": self.issuer,
+        }
+
+        # Generate QR code if available
+        if QRCODE_AVAILABLE:
+            qr = qrcode.QRCode(version=1, box_size=10, border=5)
+            qr.add_data(uri)
+            qr.make(fit=True)
+
+            # Generate ASCII QR code
+            buffer = io.StringIO()
+            qr.print_ascii(out=buffer)
+            result["qr_ascii"] = buffer.getvalue()
+
+        # Save secret
+        with self._lock:
+            self._secrets[host] = {
+                "secret": secret,
+                "account": account,
+                "created_at": datetime.now(timezone.utc).isoformat(),
+                "enabled": True,
+            }
+            self._save_secrets()
+
+        logger.info("totp_setup", host=host)
+        return result
+
+    def verify_totp(self, host: str, code: str) -> bool:
+        """
+        Verify a TOTP code for a host.
+
+        Args:
+            host: Host name.
+            code: TOTP code to verify.
+
+        Returns:
+            True if verification successful.
+
+        Raises:
+            TOTPNotConfigured: If TOTP not set up for host.
+            TOTPVerificationFailed: If code is invalid.
+        """
+        with self._lock:
+            if host not in self._secrets:
+                raise TOTPNotConfigured(f"TOTP not configured for host: {host}")
+
+            host_config = self._secrets[host]
+            if not host_config.get("enabled", True):
+                return True  # TOTP disabled for this host
+
+            secret = host_config["secret"]
+
+        totp = pyotp.TOTP(secret, digits=self.digits, interval=self.interval)
+
+        # Verify with 1-step window for clock drift
+        is_valid = totp.verify(code, valid_window=1)
+
+        # Notify callbacks
+        for callback in self._verification_callbacks:
+            try:
+                callback(host, is_valid)
+            except Exception:
+                pass
+
+        if not is_valid:
+            logger.warning("totp_verification_failed", host=host)
+            raise TOTPVerificationFailed("Invalid TOTP code")
+
+        logger.info("totp_verified", host=host)
+        return True
+
+    def is_totp_required(self, host: str) -> bool:
+        """
+        Check if TOTP is required for a host.
+
+        Args:
+            host: Host name.
+
+        Returns:
+            True if TOTP is configured and enabled.
+        """
+        with self._lock:
+            if host not in self._secrets:
+                return False
+            return self._secrets[host].get("enabled", True)
+
+    def is_critical_command(self, command: str) -> bool:
+        """
+        Check if a command is considered critical (requires TOTP).
+
+        Args:
+            command: Command to check.
+
+        Returns:
+            True if command is critical.
+        """
+        import re
+
+        command_lower = command.lower()
+        for pattern in self._critical_commands:
+            if re.match(pattern, command_lower, re.IGNORECASE):
+                return True
+        return False
+
+    def add_critical_pattern(self, pattern: str) -> None:
+        """
+        Add a pattern for critical commands.
+
+        Args:
+            pattern: Regex pattern to add.
+        """
+        self._critical_commands.append(pattern)
+
+    def remove_totp(self, host: str) -> bool:
+        """
+        Remove TOTP configuration for a host.
+
+        Args:
+            host: Host name.
+
+        Returns:
+            True if removed, False if not configured.
+        """
+        with self._lock:
+            if host not in self._secrets:
+                return False
+
+            del self._secrets[host]
+            self._save_secrets()
+
+        logger.info("totp_removed", host=host)
+        return True
+
+    def disable_totp(self, host: str) -> bool:
+        """
+        Temporarily disable TOTP for a host.
+
+        Args:
+            host: Host name.
+
+        Returns:
+            True if disabled.
+        """
+        with self._lock:
+            if host not in self._secrets:
+                return False
+
+            self._secrets[host]["enabled"] = False
+            self._save_secrets()
+
+        logger.info("totp_disabled", host=host)
+        return True
+
+    def enable_totp(self, host: str) -> bool:
+        """
+        Enable TOTP for a host.
+
+        Args:
+            host: Host name.
+
+        Returns:
+            True if enabled.
+        """
+        with self._lock:
+            if host not in self._secrets:
+                return False
+
+            self._secrets[host]["enabled"] = True
+            self._save_secrets()
+
+        logger.info("totp_enabled", host=host)
+        return True
+
+    def get_configured_hosts(self) -> list[str]:
+        """Get list of hosts with TOTP configured."""
+        with self._lock:
+            return list(self._secrets.keys())
+
+    def register_verification_callback(
+        self, callback: Callable[[str, bool], None]
+    ) -> None:
+        """
+        Register a callback for verification events.
+
+        Args:
+            callback: Function called with (host, success).
+        """
+        self._verification_callbacks.append(callback)
+
+    def _load_secrets(self) -> None:
+        """Load secrets from file."""
+        path = Path(self.secrets_file)
+        if not path.exists():
+            return
+
+        try:
+            with open(path) as f:
+                self._secrets = json.load(f)
+            logger.info("totp_secrets_loaded", count=len(self._secrets))
+        except Exception as e:
+            logger.error("totp_secrets_load_error", error=str(e))
+
+    def _save_secrets(self) -> None:
+        """Save secrets to file."""
+        path = Path(self.secrets_file)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            with open(path, "w") as f:
+                json.dump(self._secrets, f, indent=2)
+            # Set restrictive permissions
+            os.chmod(path, 0o600)
+        except Exception as e:
+            logger.error("totp_secrets_save_error", error=str(e))
+
+
+# Global TOTP manager instance
+_totp_manager: TOTPManager | None = None
+
+
+def get_totp_manager() -> TOTPManager:
+    """Get or create the global TOTP manager."""
+    global _totp_manager
+    if _totp_manager is None:
+        _totp_manager = TOTPManager()
+    return _totp_manager
+
+
+def init_totp_manager(
+    secrets_file: str | None = None,
+    issuer: str = "SSH-MCP",
+) -> TOTPManager:
+    """
+    Initialize the global TOTP manager.
+
+    Args:
+        secrets_file: Path to secrets file.
+        issuer: Issuer name for TOTP.
+
+    Returns:
+        Initialized TOTPManager.
+    """
+    global _totp_manager
+    _totp_manager = TOTPManager(secrets_file=secrets_file, issuer=issuer)
+    return _totp_manager
+
+
+def require_totp(host: str, code: str | None = None) -> bool:
+    """
+    Check if TOTP is required and verify if code provided.
+
+    Args:
+        host: Host name.
+        code: Optional TOTP code.
+
+    Returns:
+        True if verified or not required.
+
+    Raises:
+        TOTPVerificationFailed: If code required but invalid/missing.
+    """
+    manager = get_totp_manager()
+
+    if not manager.is_totp_required(host):
+        return True
+
+    if code is None:
+        raise TOTPVerificationFailed(f"TOTP code required for host: {host}")
+
+    return manager.verify_totp(host, code)

sshmcp/security/validator.py

@@ -0,0 +1,234 @@
+"""Security validation for commands and paths."""
+
+import os
+import re
+from typing import List, Tuple
+
+import structlog
+
+from sshmcp.models.machine import SecurityConfig
+
+logger = structlog.get_logger()
+
+
+class ValidationError(Exception):
+    """Command or path validation failed."""
+
+    pass
+
+
+def validate_command(
+    command: str,
+    security: SecurityConfig,
+) -> Tuple[bool, str | None]:
+    """
+    Validate command against security rules.
+
+    Args:
+        command: Command to validate.
+        security: Security configuration with allowed/forbidden patterns.
+
+    Returns:
+        Tuple of (is_valid, error_message).
+    """
+    command = command.strip()
+
+    if not command:
+        return False, "Empty command"
+
+    # Check forbidden commands first (blacklist)
+    for pattern in security.forbidden_commands:
+        try:
+            if re.match(pattern, command, re.IGNORECASE):
+                logger.warning(
+                    "command_forbidden",
+                    command=command,
+                    pattern=pattern,
+                )
+                return False, f"Command matches forbidden pattern: {pattern}"
+        except re.error as e:
+            logger.error("invalid_regex_pattern", pattern=pattern, error=str(e))
+            continue
+
+    # If no allowed commands specified, allow all (except forbidden)
+    if not security.allowed_commands:
+        return True, None
+
+    # Check allowed commands (whitelist)
+    for pattern in security.allowed_commands:
+        try:
+            if re.match(pattern, command, re.IGNORECASE):
+                logger.debug(
+                    "command_allowed",
+                    command=command,
+                    pattern=pattern,
+                )
+                return True, None
+        except re.error as e:
+            logger.error("invalid_regex_pattern", pattern=pattern, error=str(e))
+            continue
+
+    logger.warning(
+        "command_not_allowed",
+        command=command,
+    )
+    return False, "Command not in allowed list"
+
+
+def validate_path(
+    path: str,
+    security: SecurityConfig,
+    check_type: str = "read",
+) -> Tuple[bool, str | None]:
+    """
+    Validate file path against security rules.
+
+    Args:
+        path: Path to validate.
+        security: Security configuration with allowed/forbidden paths.
+        check_type: Type of operation ('read' or 'write').
+
+    Returns:
+        Tuple of (is_valid, error_message).
+    """
+    if not path:
+        return False, "Empty path"
+
+    # Normalize path
+    normalized = normalize_path(path)
+
+    # Check for path traversal attacks
+    if ".." in normalized or normalized.startswith("../"):
+        logger.warning(
+            "path_traversal_attempt",
+            path=path,
+            normalized=normalized,
+        )
+        return False, "Path traversal not allowed"
+
+    # Check forbidden paths
+    for forbidden in security.forbidden_paths:
+        forbidden_normalized = normalize_path(forbidden)
+        if (
+            normalized.startswith(forbidden_normalized)
+            or normalized == forbidden_normalized
+        ):
+            logger.warning(
+                "path_forbidden",
+                path=path,
+                forbidden=forbidden,
+            )
+            return False, f"Path is forbidden: {forbidden}"
+
+    # If no allowed paths specified, allow all (except forbidden)
+    if not security.allowed_paths:
+        return True, None
+
+    # Check allowed paths
+    for allowed in security.allowed_paths:
+        allowed_normalized = normalize_path(allowed)
+        if (
+            normalized.startswith(allowed_normalized)
+            or normalized == allowed_normalized
+        ):
+            logger.debug(
+                "path_allowed",
+                path=path,
+                allowed=allowed,
+            )
+            return True, None
+
+    logger.warning(
+        "path_not_allowed",
+        path=path,
+    )
+    return False, "Path not in allowed list"
+
+
+def normalize_path(path: str) -> str:
+    """
+    Normalize a file path.
+
+    Args:
+        path: Path to normalize.
+
+    Returns:
+        Normalized path.
+    """
+    # Expand user home directory
+    if path.startswith("~"):
+        path = os.path.expanduser(path)
+
+    # Remove trailing slashes
+    path = path.rstrip("/")
+
+    # Collapse multiple slashes
+    while "//" in path:
+        path = path.replace("//", "/")
+
+    return path
+
+
+def check_command_safety(command: str) -> List[str]:
+    """
+    Check command for potential safety issues.
+
+    Returns list of warnings (empty if no issues found).
+
+    Args:
+        command: Command to check.
+
+    Returns:
+        List of warning messages.
+    """
+    warnings = []
+
+    # Dangerous patterns to warn about
+    dangerous_patterns = [
+        (r"rm\s+-rf", "Recursive force delete detected"),
+        (r"rm\s+.*\*", "Wildcard delete detected"),
+        (r">\s*/dev/", "Writing to /dev/ detected"),
+        (r"chmod\s+777", "World-writable permissions detected"),
+        (r"\|\s*sh", "Piping to shell detected"),
+        (r"\|\s*bash", "Piping to bash detected"),
+        (r"curl.*\|\s*", "Curl piped to command detected"),
+        (r"wget.*\|\s*", "Wget piped to command detected"),
+        (r"eval\s+", "Eval command detected"),
+        (r";\s*rm\s+", "Command chaining with rm detected"),
+        (r"&&\s*rm\s+", "Command chaining with rm detected"),
+        (r"\$\(.*\)", "Command substitution detected"),
+        (r"`.*`", "Backtick command substitution detected"),
+    ]
+
+    for pattern, warning in dangerous_patterns:
+        if re.search(pattern, command, re.IGNORECASE):
+            warnings.append(warning)
+
+    return warnings
+
+
+def sanitize_command_for_log(command: str) -> str:
+    """
+    Sanitize command for logging (hide potential secrets).
+
+    Args:
+        command: Command to sanitize.
+
+    Returns:
+        Sanitized command safe for logging.
+    """
+    # Patterns that might contain secrets
+    secret_patterns = [
+        (r"password[=:\s]+\S+", "password=***"),
+        (r"passwd[=:\s]+\S+", "passwd=***"),
+        (r"secret[=:\s]+\S+", "secret=***"),
+        (r"token[=:\s]+\S+", "token=***"),
+        (r"api[_-]?key[=:\s]+\S+", "api_key=***"),
+        (r"AWS_SECRET[=:\s]+\S+", "AWS_SECRET=***"),
+    ]
+
+    sanitized = command
+    for pattern, replacement in secret_patterns:
+        sanitized = re.sub(pattern, replacement, sanitized, flags=re.IGNORECASE)
+
+    return sanitized