mcp-proxy-adapter 4.1.1__py3-none-any.whl → 6.0.1__py3-none-any.whl

mcp_proxy_adapter/commands/key_management_command.py

@@ -0,0 +1,506 @@
+"""
+Key Management Command
+
+This module provides commands for key management including generation,
+validation, rotation, backup, and restoration.
+
+Author: MCP Proxy Adapter Team
+Version: 1.0.0
+"""
+
+import logging
+import os
+import shutil
+from typing import Dict, List, Optional, Any
+from pathlib import Path
+from datetime import datetime
+
+from .base import Command
+from .result import CommandResult, SuccessResult, ErrorResult
+from ..core.certificate_utils import CertificateUtils
+
+logger = logging.getLogger(__name__)
+
+
+class KeyResult:
+    """
+    Result class for key operations.
+    
+    Contains key information and operation status.
+    """
+    
+    def __init__(self, key_path: str, key_type: str, key_size: int,
+                 created_date: Optional[str] = None, expiry_date: Optional[str] = None,
+                 status: str = "valid", error: Optional[str] = None):
+        """
+        Initialize key result.
+        
+        Args:
+            key_path: Path to key file
+            key_type: Type of key (RSA, ECDSA, etc.)
+            key_size: Key size in bits
+            created_date: Key creation date
+            expiry_date: Key expiry date
+            status: Key status (valid, expired, error)
+            error: Error message if any
+        """
+        self.key_path = key_path
+        self.key_type = key_type
+        self.key_size = key_size
+        self.created_date = created_date
+        self.expiry_date = expiry_date
+        self.status = status
+        self.error = error
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert to dictionary format.
+        
+        Returns:
+            Dictionary representation
+        """
+        return {
+            "key_path": self.key_path,
+            "key_type": self.key_type,
+            "key_size": self.key_size,
+            "created_date": self.created_date,
+            "expiry_date": self.expiry_date,
+            "status": self.status,
+            "error": self.error
+        }
+    
+    def get_schema(self) -> Dict[str, Any]:
+        """
+        Get JSON schema for this result.
+        
+        Returns:
+            JSON schema dictionary
+        """
+        return {
+            "type": "object",
+            "properties": {
+                "key_path": {"type": "string", "description": "Path to key file"},
+                "key_type": {"type": "string", "description": "Type of key"},
+                "key_size": {"type": "integer", "description": "Key size in bits"},
+                "created_date": {"type": "string", "description": "Key creation date"},
+                "expiry_date": {"type": "string", "description": "Key expiry date"},
+                "status": {"type": "string", "enum": ["valid", "expired", "error"], 
+                          "description": "Key status"},
+                "error": {"type": "string", "description": "Error message if any"}
+            },
+            "required": ["key_path", "key_type", "key_size", "status"]
+        }
+
+
+class KeyManagementCommand(Command):
+    """
+    Command for key management.
+    
+    Provides methods for generating, validating, rotating, backing up, and restoring keys.
+    """
+    
+    # Command metadata
+    name = "key_management"
+    version = "1.0.0"
+    descr = "Private key generation, validation, and management"
+    category = "security"
+    author = "MCP Proxy Adapter Team"
+    email = "team@mcp-proxy-adapter.com"
+    source_url = "https://github.com/mcp-proxy-adapter"
+    result_class = KeyResult
+    
+    def __init__(self):
+        """Initialize key management command."""
+        super().__init__()
+        self.certificate_utils = CertificateUtils()
+    
+    async def execute(self, **kwargs) -> CommandResult:
+        """
+        Execute key management command.
+        
+        Args:
+            **kwargs: Command parameters including:
+                - action: Action to perform (key_generate, key_validate, key_rotate, key_backup, key_restore)
+                - key_type: Type of key to generate (RSA, ECDSA)
+                - key_size: Key size in bits for generation
+                - output_path: Output path for key generation
+                - password: Password for key encryption
+                - key_path: Key file path for validation, rotation, backup
+                - old_key_path: Old key path for rotation
+                - new_key_path: New key path for rotation
+                - cert_path: Certificate path for rotation
+                - backup_old: Whether to backup old key during rotation
+                - backup_path: Backup path for key backup/restore
+                - encrypt_backup: Whether to encrypt backup
+                
+        Returns:
+            CommandResult with key operation status
+        """
+        action = kwargs.get("action", "key_validate")
+        
+        if action == "key_generate":
+            key_type = kwargs.get("key_type")
+            key_size = kwargs.get("key_size")
+            output_path = kwargs.get("output_path")
+            password = kwargs.get("password")
+            return await self.key_generate(key_type, key_size, output_path, password)
+        elif action == "key_validate":
+            key_path = kwargs.get("key_path")
+            password = kwargs.get("password")
+            return await self.key_validate(key_path, password)
+        elif action == "key_rotate":
+            old_key_path = kwargs.get("old_key_path")
+            new_key_path = kwargs.get("new_key_path")
+            cert_path = kwargs.get("cert_path")
+            backup_old = kwargs.get("backup_old", True)
+            return await self.key_rotate(old_key_path, new_key_path, cert_path, backup_old)
+        elif action == "key_backup":
+            key_path = kwargs.get("key_path")
+            backup_path = kwargs.get("backup_path")
+            encrypt_backup = kwargs.get("encrypt_backup", True)
+            password = kwargs.get("password")
+            return await self.key_backup(key_path, backup_path, encrypt_backup, password)
+        elif action == "key_restore":
+            backup_path = kwargs.get("backup_path")
+            key_path = kwargs.get("key_path")
+            password = kwargs.get("password")
+            return await self.key_restore(backup_path, key_path, password)
+        else:
+            return ErrorResult(
+                message=f"Unknown action: {action}. Supported actions: key_generate, key_validate, key_rotate, key_backup, key_restore"
+            )
+    
+    async def key_generate(self, key_type: str, key_size: int, output_path: str,
+                          password: Optional[str] = None) -> CommandResult:
+        """
+        Generate a new private key.
+        
+        Args:
+            key_type: Type of key to generate (RSA, ECDSA)
+            key_size: Key size in bits
+            output_path: Path to save the generated key
+            password: Optional password to encrypt the key
+            
+        Returns:
+            CommandResult with key generation status
+        """
+        try:
+            logger.info(f"Generating {key_type} key with size {key_size} bits")
+            
+            # Validate parameters
+            if key_type not in ["RSA", "ECDSA"]:
+                return ErrorResult(
+                    message="Key type must be RSA or ECDSA"
+                )
+            
+            if key_type == "ECDSA":
+                if key_size not in [256, 384, 521]:
+                    return ErrorResult(
+                        message="ECDSA key size must be 256, 384, or 521 bits"
+                    )
+            elif key_size < 1024:
+                return ErrorResult(
+                    message="Key size must be at least 1024 bits"
+                )
+            
+            # Create output directory if it doesn't exist
+            output_dir = os.path.dirname(output_path)
+            if output_dir:
+                Path(output_dir).mkdir(parents=True, exist_ok=True)
+            
+            # Generate key
+            result = self.certificate_utils.generate_private_key(
+                key_type, key_size, output_path
+            )
+            
+            if not result["success"]:
+                return ErrorResult(
+                    message=f"Key generation failed: {result['error']}"
+                )
+            
+            key_result = KeyResult(
+                key_path=output_path,
+                key_type=key_type,
+                key_size=key_size,
+                created_date=datetime.now().isoformat(),
+                status="valid"
+            )
+            
+            logger.info(f"Key generated successfully: {output_path}")
+            return SuccessResult(
+                data={
+                    "key": key_result.to_dict(),
+                    "details": result
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Key generation failed: {e}")
+            return ErrorResult(
+                message=f"Key generation failed: {str(e)}"
+            )
+    
+    async def key_validate(self, key_path: str, password: Optional[str] = None) -> CommandResult:
+        """
+        Validate a private key.
+        
+        Args:
+            key_path: Path to key file to validate
+            password: Optional password if key is encrypted
+            
+        Returns:
+            CommandResult with key validation status
+        """
+        try:
+            logger.info(f"Validating key: {key_path}")
+            
+            # Validate parameters
+            if not key_path or not os.path.exists(key_path):
+                return ErrorResult(
+                    message=f"Key file not found: {key_path}"
+                )
+            
+            # Validate key
+            result = self.certificate_utils.validate_private_key(key_path)
+            
+            if not result["success"]:
+                return ErrorResult(
+                    message=f"Key validation failed: {result['error']}"
+                )
+            
+            key_result = KeyResult(
+                key_path=key_path,
+                key_type=result.get("key_type", "unknown"),
+                key_size=result.get("key_size", 0),
+                created_date=result.get("created_date"),
+                status="valid"
+            )
+            
+            logger.info(f"Key validation completed: {key_path}")
+            return SuccessResult(
+                data={
+                    "key": key_result.to_dict()
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Key validation failed: {e}")
+            return ErrorResult(
+                message=f"Key validation failed: {str(e)}"
+            )
+    
+    async def key_rotate(self, old_key_path: str, new_key_path: str,
+                        cert_path: Optional[str] = None, backup_old: bool = True) -> CommandResult:
+        """
+        Rotate a private key.
+        
+        Args:
+            old_key_path: Path to old key file
+            new_key_path: Path to new key file
+            cert_path: Optional certificate path to update with new key
+            backup_old: Whether to backup the old key
+            
+        Returns:
+            CommandResult with key rotation status
+        """
+        try:
+            logger.info(f"Rotating key from {old_key_path} to {new_key_path}")
+            
+            # Validate parameters
+            if not old_key_path or not os.path.exists(old_key_path):
+                return ErrorResult(
+                    message=f"Old key file not found: {old_key_path}"
+                )
+            
+            if not new_key_path or not os.path.exists(new_key_path):
+                return ErrorResult(
+                    message=f"New key file not found: {new_key_path}"
+                )
+            
+            # Validate both keys
+            old_key_validation = await self.key_validate(old_key_path)
+            new_key_validation = await self.key_validate(new_key_path)
+            
+            if not old_key_validation.to_dict()["success"]:
+                return ErrorResult(
+                    message=f"Old key validation failed: {old_key_validation.to_dict()['error']['message']}"
+                )
+            
+            if not new_key_validation.to_dict()["success"]:
+                return ErrorResult(
+                    message=f"New key validation failed: {new_key_validation.to_dict()['error']['message']}"
+                )
+            
+            # Backup old key if requested
+            backup_path = None
+            if backup_old:
+                backup_path = f"{old_key_path}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+                shutil.copy2(old_key_path, backup_path)
+                logger.info(f"Old key backed up to: {backup_path}")
+            
+            # Update certificate if provided
+            cert_updated = False
+            if cert_path and os.path.exists(cert_path):
+                try:
+                    # This would require implementing certificate re-signing
+                    # For now, we'll just note that it needs to be done
+                    cert_updated = True
+                    logger.info(f"Certificate {cert_path} needs to be re-signed with new key")
+                except Exception as e:
+                    logger.warning(f"Could not update certificate {cert_path}: {e}")
+            
+            # Replace old key with new key
+            shutil.copy2(new_key_path, old_key_path)
+            
+            logger.info(f"Key rotation completed successfully")
+            return SuccessResult(
+                data={
+                    "old_key_path": old_key_path,
+                    "new_key_path": new_key_path,
+                    "backup_path": backup_path,
+                    "cert_updated": cert_updated,
+                    "message": "Key rotation completed successfully"
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Key rotation failed: {e}")
+            return ErrorResult(
+                message=f"Key rotation failed: {str(e)}"
+            )
+    
+    async def key_backup(self, key_path: str, backup_path: str,
+                        encrypt_backup: bool = True, password: Optional[str] = None) -> CommandResult:
+        """
+        Backup a private key.
+        
+        Args:
+            key_path: Path to key file to backup
+            backup_path: Path to save the backup
+            encrypt_backup: Whether to encrypt the backup
+            password: Password for backup encryption
+            
+        Returns:
+            CommandResult with backup status
+        """
+        try:
+            logger.info(f"Backing up key: {key_path}")
+            
+            # Validate parameters
+            if not key_path or not os.path.exists(key_path):
+                return ErrorResult(
+                    message=f"Key file not found: {key_path}"
+                )
+            
+            # Validate key before backup
+            key_validation = await self.key_validate(key_path)
+            if not key_validation.to_dict()["success"]:
+                return ErrorResult(
+                    message=f"Key validation failed before backup: {key_validation.to_dict()['error']['message']}"
+                )
+            
+            # Create backup directory if it doesn't exist
+            backup_dir = os.path.dirname(backup_path)
+            if backup_dir:
+                Path(backup_dir).mkdir(parents=True, exist_ok=True)
+            
+            # Create backup
+            if encrypt_backup and password:
+                # Encrypted backup
+                result = self.certificate_utils.create_encrypted_backup(
+                    key_path, backup_path, password
+                )
+                if not result["success"]:
+                    return ErrorResult(
+                        message=f"Encrypted backup failed: {result['error']}"
+                    )
+            else:
+                # Simple file copy
+                shutil.copy2(key_path, backup_path)
+            
+            # Verify backup
+            if not os.path.exists(backup_path):
+                return ErrorResult(
+                    message="Backup file was not created"
+                )
+            
+            logger.info(f"Key backup completed successfully: {backup_path}")
+            return SuccessResult(
+                data={
+                    "key_path": key_path,
+                    "backup_path": backup_path,
+                    "encrypted": encrypt_backup,
+                    "backup_size": os.path.getsize(backup_path),
+                    "backup_date": datetime.now().isoformat()
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Key backup failed: {e}")
+            return ErrorResult(
+                message=f"Key backup failed: {str(e)}"
+            )
+    
+    async def key_restore(self, backup_path: str, key_path: str,
+                         password: Optional[str] = None) -> CommandResult:
+        """
+        Restore a private key from backup.
+        
+        Args:
+            backup_path: Path to backup file
+            key_path: Path to restore the key to
+            password: Password if backup is encrypted
+            
+        Returns:
+            CommandResult with restore status
+        """
+        try:
+            logger.info(f"Restoring key from backup: {backup_path}")
+            
+            # Validate parameters
+            if not backup_path or not os.path.exists(backup_path):
+                return ErrorResult(
+                    message=f"Backup file not found: {backup_path}"
+                )
+            
+            # Create target directory if it doesn't exist
+            key_dir = os.path.dirname(key_path)
+            if key_dir:
+                Path(key_dir).mkdir(parents=True, exist_ok=True)
+            
+            # Restore key
+            if password:
+                # Try encrypted restore first
+                result = self.certificate_utils.restore_encrypted_backup(
+                    backup_path, key_path, password
+                )
+                if not result["success"]:
+                    return ErrorResult(
+                        message=f"Encrypted restore failed: {result['error']}"
+                    )
+            else:
+                # Simple file copy
+                shutil.copy2(backup_path, key_path)
+            
+            # Validate restored key
+            key_validation = await self.key_validate(key_path)
+            if not key_validation.to_dict()["success"]:
+                return ErrorResult(
+                    message=f"Restored key validation failed: {key_validation.to_dict()['error']['message']}"
+                )
+            
+            logger.info(f"Key restore completed successfully: {key_path}")
+            return SuccessResult(
+                data={
+                    "backup_path": backup_path,
+                    "key_path": key_path,
+                    "restore_date": datetime.now().isoformat(),
+                    "key_info": key_validation.to_dict().get("data", {}).get("key") if key_validation.to_dict().get("success") else None
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Key restore failed: {e}")
+            return ErrorResult(
+                message=f"Key restore failed: {str(e)}"
+            ) 

mcp_proxy_adapter/commands/load_command.py

@@ -0,0 +1,176 @@
+"""
+Module with load command implementation.
+"""
+
+from typing import Dict, Any, Optional, List
+
+from mcp_proxy_adapter.commands.base import Command
+from mcp_proxy_adapter.commands.result import CommandResult, SuccessResult
+from mcp_proxy_adapter.commands.command_registry import registry
+
+
+class LoadResult(SuccessResult):
+    """
+    Result of the load command execution.
+    """
+    
+    def __init__(self, success: bool, commands_loaded: int, loaded_commands: list, source: str, error: Optional[str] = None):
+        """
+        Initialize load command result.
+        
+        Args:
+            success: Whether loading was successful
+            commands_loaded: Number of commands loaded
+            loaded_commands: List of loaded command names
+            source: Source path or URL
+            error: Error message if loading failed
+        """
+        data = {
+            "success": success,
+            "commands_loaded": commands_loaded,
+            "loaded_commands": loaded_commands,
+            "source": source
+        }
+        if error:
+            data["error"] = error
+            
+        message = f"Loaded {commands_loaded} commands from {source}"
+        if error:
+            message = f"Failed to load commands from {source}: {error}"
+            
+        super().__init__(data=data, message=message)
+    
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """
+        Get JSON schema for result validation.
+        
+        Returns:
+            Dict[str, Any]: JSON schema
+        """
+        return {
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "object",
+                    "properties": {
+                        "success": {"type": "boolean"},
+                        "commands_loaded": {"type": "integer"},
+                        "loaded_commands": {
+                            "type": "array",
+                            "items": {"type": "string"}
+                        },
+                        "source": {"type": "string"},
+                        "error": {"type": "string"}
+                    },
+                    "required": ["success", "commands_loaded", "loaded_commands", "source"]
+                }
+            },
+            "required": ["data"]
+        }
+
+
+class LoadCommand(Command):
+    """
+    Command that loads commands from local path or URL.
+    
+    This command allows dynamic loading of command modules from either local file system
+    or remote HTTP/HTTPS URLs. The command automatically detects whether the source
+    is a local path or URL and handles the loading accordingly.
+    
+    For local paths, the command loads Python modules ending with '_command.py'.
+    For URLs, the command downloads the Python code and loads it as a temporary module.
+    
+    The loaded commands are registered in the command registry and become immediately
+    available for execution. Only commands that inherit from the base Command class
+    and are properly structured will be loaded and registered.
+    
+    Security considerations:
+    - Local paths are validated for existence and proper naming
+    - URLs are downloaded with timeout protection
+    - Temporary files are automatically cleaned up after loading
+    - Only files ending with '_command.py' are accepted
+    
+    Examples:
+    - Load from local file: "./my_command.py"
+    - Load from URL: "https://example.com/remote_command.py"
+    """
+    
+    name = "load"
+    result_class = LoadResult
+    
+    async def execute(self, source: str, **kwargs) -> LoadResult:
+        """
+        Execute load command.
+        
+        Args:
+            source: Source path or URL to load command from
+            **kwargs: Additional parameters
+            
+        Returns:
+            LoadResult: Load command result
+        """
+        # Load command from source
+        result = registry.load_command_from_source(source)
+        
+        return LoadResult(
+            success=result.get("success", False),
+            commands_loaded=result.get("commands_loaded", 0),
+            loaded_commands=result.get("loaded_commands", []),
+            source=result.get("source", source),
+            error=result.get("error")
+        )
+    
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """
+        Get JSON schema for command parameters.
+        
+        Returns:
+            Dict[str, Any]: JSON schema
+        """
+        return {
+            "type": "object",
+            "properties": {
+                "source": {
+                    "type": "string",
+                    "description": "Source path or URL to load command from (must end with '_command.py')",
+                    "examples": [
+                        "./my_command.py",
+                        "https://example.com/remote_command.py"
+                    ]
+                }
+            },
+            "required": ["source"]
+        }
+    
+    @classmethod
+    def _generate_examples(cls, params: Dict[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Generate custom examples for load command.
+        
+        Args:
+            params: Information about command parameters
+            
+        Returns:
+            List of examples
+        """
+        examples = [
+            {
+                "command": cls.name,
+                "params": {"source": "./custom_command.py"},
+                "description": "Load a command from local file system"
+            },
+            {
+                "command": cls.name,
+                "params": {"source": "https://raw.githubusercontent.com/user/repo/main/remote_command.py"},
+                "description": "Load a command from GitHub raw content"
+            },
+            {
+                "command": cls.name,
+                "params": {"source": "https://example.com/api/commands/test_command.py"},
+                "description": "Load a command from remote API endpoint"
+            }
+        ]
+        
+        return examples