Ludus FastMCP

"""Base scenario builder for Ludus range configurations."""

from typing import Any
import yaml
from .siem_config import (
    SIEMType,
    get_siem_server_config,
    add_siem_agent_to_vm,
    get_siem_network_rules,
)
from ..schemas.scenario_customization import ScenarioCustomization


class BaseScenarioBuilder:
    """Base class for building Ludus range scenarios."""

    def __init__(
        self,
        range_id: str = "{{ range_id }}",
        siem_type: SIEMType = "wazuh",
        resource_profile: str = "minimal",
        customization: ScenarioCustomization | None = None,
        randomize: bool = False,
    ):
        """Initialize the scenario builder.

        Args:
            range_id: Range identifier
            siem_type: SIEM type (wazuh, splunk, elastic, security-onion, none)
            resource_profile: Resource allocation profile - "minimal", "recommended", "maximum"
                - minimal: Lowest resource requirements (may be slow)
                - recommended: Balanced performance and resource usage (default)
                - maximum: High performance with maximum resources
            customization: Optional scenario customization (users, vulnerabilities, etc.)
            randomize: Whether to apply randomization (requires customization with randomization_config)
        """
        self.range_id = range_id
        self.siem_type = siem_type
        self.resource_profile = resource_profile
        self.customization = customization
        self.randomize = randomize
        self.config = {
            "name": None,  # Range name - will be set by scenario builders
            "ludus": [],
            "network": {
                "inter_vlan_default": "REJECT",
                "rules": [],
            },
        }

        # Resource profiles: (ram_gb, cpus)
        self.RESOURCE_PROFILES = {
            "minimal": {
                "dc": (2, 2),
                "workstation": (2, 2),
                "server": (2, 2),
                "sql_server": (4, 2),
                "exchange_server": (4, 2),
                "kali": (4, 2),
                "siem": (4, 2),
                "edr": (4, 2),
                "linux_server": (2, 2),
                "analysis_win": (4, 2),
                "analysis_linux": (4, 2),
            },
            "recommended": {
                "dc": (4, 2),
                "workstation": (4, 2),
                "server": (4, 2),
                "sql_server": (8, 4),
                "exchange_server": (8, 4),
                "kali": (8, 4),
                "siem": (8, 4),
                "edr": (6, 2),
                "linux_server": (4, 2),
                "analysis_win": (8, 4),
                "analysis_linux": (8, 4),
            },
            "maximum": {
                "dc": (8, 4),
                "workstation": (8, 4),
                "server": (8, 4),
                "sql_server": (16, 8),
                "exchange_server": (16, 8),
                "kali": (16, 8),
                "siem": (16, 8),
                "edr": (12, 4),
                "linux_server": (8, 4),
                "analysis_win": (16, 8),
                "analysis_linux": (16, 8),
            },
        }

    def get_resources(self, vm_type: str) -> tuple[int, int]:
        """Get RAM and CPU allocation for a VM type based on resource profile.

        Args:
            vm_type: Type of VM (dc, workstation, server, etc.)

        Returns:
            Tuple of (ram_gb, cpus)
        """
        profile = self.RESOURCE_PROFILES.get(self.resource_profile, self.RESOURCE_PROFILES["minimal"])
        return profile.get(vm_type, (2, 2))  # Default to 2GB/2CPU if type not found

    def add_vm(
        self,
        vm_name: str,
        hostname: str,
        template: str,
        vlan: int,
        ip_last_octet: int,
        ram_gb: int = 8,
        cpus: int = 4,
        **kwargs: Any,
    ) -> None:
        """Add a VM to the configuration."""
        vm_config: dict[str, Any] = {
            "vm_name": vm_name,
            "hostname": hostname,
            "template": template,
            "vlan": vlan,
            "ip_last_octet": ip_last_octet,
            "ram_gb": ram_gb,
            "cpus": cpus,
        }
        vm_config.update(kwargs)
        self.config["ludus"].append(vm_config)

    def add_network_rule(
        self,
        name: str,
        vlan_src: int,
        vlan_dst: int,
        protocol: str = "all",
        ports: str | int = "all",
        action: str = "ACCEPT",
    ) -> None:
        """Add a network rule."""
        rule = {
            "name": name,
            "vlan_src": vlan_src,
            "vlan_dst": vlan_dst,
            "protocol": protocol,
            "ports": ports,
            "action": action,
        }
        self.config["network"]["rules"].append(rule)

    def to_yaml(self) -> str:
        """Convert configuration to YAML string."""
        # Add YAML schema comment
        yaml_str = "# yaml-language-server: $schema=https://docs.ludus.cloud/schemas/range-config.json\n\n"
        yaml_str += yaml.dump(self.config, default_flow_style=False, sort_keys=False)
        return yaml_str

    def add_siem_server(
        self, vlan: int = 10, ip_last_octet: int = 100, siem_type: SIEMType | None = None
    ) -> dict[str, Any]:
        """Add SIEM server to the configuration."""
        if siem_type is None:
            siem_type = self.siem_type
        
        if siem_type == "none":
            return {}
        
        siem_config = get_siem_server_config(siem_type, self.range_id, vlan, ip_last_octet)
        self.add_vm(**siem_config)
        
        # Add SIEM network rules
        for rule in get_siem_network_rules(siem_type):
            # Convert 0 (all VLANs) to actual VLAN numbers for each rule
            if rule["vlan_src"] == 0:
                # Allow from all existing VLANs in the config
                existing_vlans = set()
                for vm in self.config["ludus"]:
                    if "vlan" in vm:
                        existing_vlans.add(vm["vlan"])
                # Add rules for each VLAN to SIEM server VLAN
                for src_vlan in existing_vlans:
                    self.add_network_rule(
                        name=rule["name"] + f" (from VLAN {src_vlan})",
                        vlan_src=src_vlan,
                        vlan_dst=vlan,
                        protocol=rule["protocol"],
                        ports=rule["ports"],
                        action=rule["action"],
                    )
            else:
                self.add_network_rule(**rule)
        
        return siem_config

    def add_siem_agents_to_all_vms(self, siem_server_ip: str | None = None, siem_type: SIEMType | None = None) -> None:
        """Add SIEM agent configuration to all VMs in the range."""
        if siem_type is None:
            siem_type = self.siem_type
        
        if siem_type == "none":
            return
        
        # If no IP provided, calculate from SIEM server config (default: VLAN 10, IP .100)
        if siem_server_ip is None:
            # Default SIEM server IP (will be set during deployment)
            siem_server_ip = "192.168.10.100"  # Default, will be replaced with actual IP
        
        siem_keywords = {
            "wazuh": "wazuh",
            "splunk": "splunk",
            "elastic": "elastic",
            "security-onion": "security-onion",
        }
        
        siem_keyword = siem_keywords.get(siem_type, "")
        
        for vm in self.config["ludus"]:
            # Skip SIEM server itself
            if siem_keyword in vm.get("vm_name", "").lower() and "server" in vm.get("vm_name", "").lower():
                continue
            
            # Add SIEM agent to all other VMs
            add_siem_agent_to_vm(vm, siem_type, siem_server_ip)
    
    # Backward compatibility methods
    def add_wazuh_server(self, vlan: int = 10, ip_last_octet: int = 100) -> dict[str, Any]:
        """Add Wazuh server (backward compatibility)."""
        return self.add_siem_server(vlan, ip_last_octet, "wazuh")
    
    def add_wazuh_agent_to_all_vms(self, wazuh_server_ip: str | None = None) -> None:
        """Add Wazuh agents (backward compatibility)."""
        self.add_siem_agents_to_all_vms(wazuh_server_ip, "wazuh")

    def set_range_name(self, name: str) -> "BaseScenarioBuilder":
        """Set the range name.
        
        Args:
            name: Range name to set
            
        Returns:
            Self for method chaining
        """
        self.config["name"] = name
        return self
    
    def apply_customizations(self) -> None:
        """Apply customizations to the configuration.
        
        This method should be called after the scenario is built but before finalizing.
        """
        if not self.customization:
            return
        
        # Apply network customizations
        if self.customization.network_customization:
            nc = self.customization.network_customization
            if nc.inter_vlan_default:
                self.config["network"]["inter_vlan_default"] = nc.inter_vlan_default
            
            # Apply VLAN changes
            if nc.vlan_changes:
                for vm in self.config.get("ludus", []):
                    vm_name = vm.get("vm_name", "")
                    for pattern, new_vlan in nc.vlan_changes.items():
                        if pattern in vm_name:
                            vm["vlan"] = new_vlan
            
            # Add additional rules
            if nc.additional_rules:
                self.config["network"]["rules"].extend(nc.additional_rules)
            
            # Remove rules
            if nc.remove_rules:
                self.config["network"]["rules"] = [
                    rule for rule in self.config["network"]["rules"]
                    if rule.get("name") not in nc.remove_rules
                ]
        
        # Apply VM customizations
        if self.customization.vm_customization:
            vc = self.customization.vm_customization
            
            # Apply resource overrides
            if vc.resource_overrides:
                for vm in self.config.get("ludus", []):
                    vm_name = vm.get("vm_name", "")
                    if vm_name in vc.resource_overrides:
                        overrides = vc.resource_overrides[vm_name]
                        if "ram_gb" in overrides:
                            vm["ram_gb"] = overrides["ram_gb"]
                        if "cpus" in overrides:
                            vm["cpus"] = overrides["cpus"]
            
            # Remove VMs
            if vc.remove_vms:
                self.config["ludus"] = [
                    vm for vm in self.config["ludus"]
                    if not any(pattern in vm.get("vm_name", "") for pattern in vc.remove_vms)
                ]
            
            # Add additional VMs
            if vc.additional_vms:
                self.config["ludus"].extend(vc.additional_vms)
    
    def to_dict(self) -> dict[str, Any]:
        """Get configuration as dictionary.
        
        Converts ansible_roles format to Ludus roles format:
        - ansible_roles: [{name: "...", vars: {...}}] -> roles: ["..."] and role_vars: {...}
        
        Also applies any customizations before returning.
        """
        # Apply customizations before converting
        self.apply_customizations()
        
        config = self.config.copy()
        
        # Ensure name is set (use default if None)
        if config.get("name") is None:
            config["name"] = "Ludus Range"
        
        # Convert ansible_roles to Ludus format (roles + role_vars)
        for vm in config.get("ludus", []):
            if "ansible_roles" in vm:
                ansible_roles = vm.pop("ansible_roles")
                
                # Extract role names for 'roles' array
                roles_list = []
                role_vars_dict = {}
                
                for role in ansible_roles:
                    if isinstance(role, dict):
                        role_name = role.get("name")
                        if role_name:
                            roles_list.append(role_name)
                            # Merge role-specific vars into role_vars
                            role_vars = role.get("vars", {})
                            if role_vars:
                                # Prefix role vars with role name to avoid conflicts
                                # Or merge directly if no conflicts
                                for key, value in role_vars.items():
                                    if key not in role_vars_dict:
                                        role_vars_dict[key] = value
                                    elif isinstance(role_vars_dict[key], dict) and isinstance(value, dict):
                                        role_vars_dict[key].update(value)
                    elif isinstance(role, str):
                        roles_list.append(role)
                
                if roles_list:
                    vm["roles"] = roles_list
                if role_vars_dict:
                    vm["role_vars"] = role_vars_dict
        
        return config