Megaraptor MCP

"""
Configuration handling for Velociraptor MCP server.

Supports two authentication methods:
1. Config file path (VELOCIRAPTOR_CONFIG_PATH)
2. Environment variables for individual components
"""

import os
import yaml
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Any


@dataclass
class VelociraptorConfig:
    """Configuration for connecting to Velociraptor server."""

    api_url: str
    client_cert: str
    client_key: str
    ca_cert: str
    api_connection_string: Optional[str] = None

    @classmethod
    def from_config_file(cls, config_path: str) -> "VelociraptorConfig":
        """Load configuration from a Velociraptor API client config file.

        The config file is generated by:
        velociraptor --config server.config.yaml config api_client --name <user> api_client.yaml
        """
        path = Path(config_path)
        if not path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        with open(path, "r") as f:
            config = yaml.safe_load(f)

        # Extract the API connection string if present (newer format)
        api_connection_string = config.get("api_connection_string")

        # Extract certificate data
        ca_cert = config.get("ca_certificate", "")
        client_cert = config.get("client_cert", "")
        client_key = config.get("client_private_key", "")

        # Get API URL
        api_url = config.get("api_url", "")
        if not api_url and api_connection_string:
            # Parse from connection string if needed
            api_url = api_connection_string

        return cls(
            api_url=api_url,
            client_cert=client_cert,
            client_key=client_key,
            ca_cert=ca_cert,
            api_connection_string=api_connection_string,
        )

    @classmethod
    def from_env(cls) -> "VelociraptorConfig":
        """Load configuration from environment variables.

        Environment variables:
        - VELOCIRAPTOR_API_URL: The API endpoint URL
        - VELOCIRAPTOR_CLIENT_CERT: Path to client certificate or PEM content
        - VELOCIRAPTOR_CLIENT_KEY: Path to client private key or PEM content
        - VELOCIRAPTOR_CA_CERT: Path to CA certificate or PEM content
        """
        api_url = os.environ.get("VELOCIRAPTOR_API_URL", "")

        # Handle cert/key - can be path or content
        client_cert = cls._load_cert_or_key(
            os.environ.get("VELOCIRAPTOR_CLIENT_CERT", "")
        )
        client_key = cls._load_cert_or_key(
            os.environ.get("VELOCIRAPTOR_CLIENT_KEY", "")
        )
        ca_cert = cls._load_cert_or_key(
            os.environ.get("VELOCIRAPTOR_CA_CERT", "")
        )

        return cls(
            api_url=api_url,
            client_cert=client_cert,
            client_key=client_key,
            ca_cert=ca_cert,
        )

    @staticmethod
    def _load_cert_or_key(value: str) -> str:
        """Load certificate or key from path or return as content."""
        if not value:
            return ""

        # Check if it looks like a file path
        if os.path.exists(value):
            with open(value, "r") as f:
                return f.read()

        # Otherwise treat as PEM content
        return value

    def validate(self) -> None:
        """Validate that required configuration is present."""
        errors = []

        if not self.api_url:
            errors.append("API URL is required")
        if not self.client_cert:
            errors.append("Client certificate is required")
        if not self.client_key:
            errors.append("Client private key is required")
        if not self.ca_cert:
            errors.append("CA certificate is required")

        if errors:
            raise ValueError(f"Invalid configuration: {'; '.join(errors)}")


def load_config() -> VelociraptorConfig:
    """Load Velociraptor configuration from available sources.

    Priority:
    1. VELOCIRAPTOR_CONFIG_PATH environment variable (config file)
    2. Individual environment variables

    Returns:
        VelociraptorConfig: The loaded configuration

    Raises:
        ValueError: If no valid configuration is found
    """
    # Try config file first
    config_path = os.environ.get("VELOCIRAPTOR_CONFIG_PATH")
    if config_path:
        config = VelociraptorConfig.from_config_file(config_path)
        config.validate()
        return config

    # Try environment variables
    config = VelociraptorConfig.from_env()

    # Check if we have any configuration
    if config.api_url or config.client_cert:
        config.validate()
        return config

    raise ValueError(
        "No Velociraptor configuration found. "
        "Set VELOCIRAPTOR_CONFIG_PATH to a config file path, "
        "or set VELOCIRAPTOR_API_URL, VELOCIRAPTOR_CLIENT_CERT, "
        "VELOCIRAPTOR_CLIENT_KEY, and VELOCIRAPTOR_CA_CERT environment variables."
    )


@dataclass
class DeploymentConfig:
    """Configuration for Velociraptor deployment operations.

    Attributes:
        deployment_id: Unique identifier for this deployment
        profile: Deployment profile (rapid, standard, enterprise)
        target: Deployment target (docker, binary, aws, azure)
        server_hostname: Hostname for the Velociraptor server
        bind_address: Address to bind the server to
        gui_port: Port for the GUI/API
        frontend_port: Port for client connections
        data_path: Path for Velociraptor data storage
        log_path: Path for log files
        admin_username: Admin user for initial setup
        enable_monitoring: Enable health monitoring
        ssl_pinning: Enable SSL certificate pinning
        auto_destroy_at: ISO timestamp for auto-destruction (None = never)
        extra_config: Additional configuration options
    """
    deployment_id: str
    profile: str = "standard"
    target: str = "docker"
    server_hostname: str = "localhost"
    bind_address: str = "0.0.0.0"
    gui_port: int = 8889
    frontend_port: int = 8000
    data_path: str = "/opt/velociraptor/data"
    log_path: str = "/opt/velociraptor/logs"
    admin_username: str = "admin"
    enable_monitoring: bool = True
    ssl_pinning: bool = True
    auto_destroy_at: Optional[str] = None
    extra_config: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "deployment_id": self.deployment_id,
            "profile": self.profile,
            "target": self.target,
            "server_hostname": self.server_hostname,
            "bind_address": self.bind_address,
            "gui_port": self.gui_port,
            "frontend_port": self.frontend_port,
            "data_path": self.data_path,
            "log_path": self.log_path,
            "admin_username": self.admin_username,
            "enable_monitoring": self.enable_monitoring,
            "ssl_pinning": self.ssl_pinning,
            "auto_destroy_at": self.auto_destroy_at,
            "extra_config": self.extra_config,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DeploymentConfig":
        """Create from dictionary."""
        return cls(
            deployment_id=data["deployment_id"],
            profile=data.get("profile", "standard"),
            target=data.get("target", "docker"),
            server_hostname=data.get("server_hostname", "localhost"),
            bind_address=data.get("bind_address", "0.0.0.0"),
            gui_port=data.get("gui_port", 8889),
            frontend_port=data.get("frontend_port", 8000),
            data_path=data.get("data_path", "/opt/velociraptor/data"),
            log_path=data.get("log_path", "/opt/velociraptor/logs"),
            admin_username=data.get("admin_username", "admin"),
            enable_monitoring=data.get("enable_monitoring", True),
            ssl_pinning=data.get("ssl_pinning", True),
            auto_destroy_at=data.get("auto_destroy_at"),
            extra_config=data.get("extra_config", {}),
        )


def generate_deployment_id() -> str:
    """Generate a unique deployment identifier."""
    import secrets
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d")
    random_suffix = secrets.token_hex(4)
    return f"vr-{timestamp}-{random_suffix}"