"""
AIPTX Beast Mode - Exploit Chain Templates
==========================================

Pre-defined multi-step attack chain patterns for common exploitation scenarios.
These templates are used by the chain builder to plan optimal attack paths.

Chain Categories:
- Web Application: SQLi→RCE, XSS→Session, Upload→Shell
- Network: SSRF→Internal, LFI→Creds→Lateral
- Post-Exploit: Shell→PrivEsc→Persistence
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


class ChainCategory(str, Enum):
    """Categories of exploit chains."""
    WEB_APP = "web_application"
    NETWORK = "network"
    POST_EXPLOIT = "post_exploitation"
    DATA_EXFIL = "data_exfiltration"
    PERSISTENCE = "persistence"
    LATERAL = "lateral_movement"


class StepType(str, Enum):
    """Types of steps in an exploit chain."""
    RECONNAISSANCE = "recon"
    VULNERABILITY = "vuln"
    EXPLOITATION = "exploit"
    POST_EXPLOIT = "post_exploit"
    ESCALATION = "escalation"
    EXFILTRATION = "exfil"
    PERSISTENCE = "persist"
    CLEANUP = "cleanup"


@dataclass
class ChainStepTemplate:
    """Template for a single step in an exploit chain."""
    name: str
    step_type: StepType
    description: str
    techniques: list[str]
    required_conditions: list[str] = field(default_factory=list)
    produces: list[str] = field(default_factory=list)  # What this step produces
    tools: list[str] = field(default_factory=list)
    fallback_techniques: list[str] = field(default_factory=list)


@dataclass
class ChainTemplate:
    """Template for a complete exploit chain."""
    name: str
    category: ChainCategory
    description: str
    entry_vuln_types: list[str]  # Vulnerability types that trigger this chain
    steps: list[ChainStepTemplate]
    success_indicators: list[str]
    estimated_impact: str
    stealth_rating: int  # 1-10, higher = stealthier
    complexity: int  # 1-10, higher = more complex

    def get_step_names(self) -> list[str]:
        """Get list of step names in order."""
        return [step.name for step in self.steps]

    def validate_conditions(self, available_conditions: list[str]) -> bool:
        """Check if first step conditions are met."""
        if not self.steps:
            return False
        first_step = self.steps[0]
        return all(cond in available_conditions for cond in first_step.required_conditions)


# =============================================================================
# SQL INJECTION CHAINS
# =============================================================================

SQLI_TO_RCE_CHAIN = ChainTemplate(
    name="sqli_to_rce",
    category=ChainCategory.WEB_APP,
    description="SQL Injection → File Write → Remote Code Execution",
    entry_vuln_types=["sql_injection", "sqli", "blind_sqli"],
    steps=[
        ChainStepTemplate(
            name="confirm_sqli",
            step_type=StepType.VULNERABILITY,
            description="Confirm SQL injection and identify DBMS",
            techniques=["error_based", "union_based", "blind_boolean", "blind_time"],
            required_conditions=["sqli_parameter"],
            produces=["dbms_type", "injection_type", "privileges"],
            tools=["sqlmap", "custom_sqli"],
        ),
        ChainStepTemplate(
            name="enumerate_privileges",
            step_type=StepType.RECONNAISSANCE,
            description="Check FILE privilege and DBA status",
            techniques=["select_user()", "select_privilege", "is_dba"],
            required_conditions=["dbms_type"],
            produces=["file_privilege", "dba_status", "current_user"],
            tools=["sqlmap", "custom_sqli"],
        ),
        ChainStepTemplate(
            name="write_webshell",
            step_type=StepType.EXPLOITATION,
            description="Write webshell using INTO OUTFILE/DUMPFILE",
            techniques=[
                "into_outfile",
                "into_dumpfile",
                "load_file_write",
                "mysql_file_priv",
                "mssql_xp_cmdshell",
                "postgres_copy",
            ],
            required_conditions=["file_privilege"],
            produces=["webshell_path", "webshell_url"],
            tools=["sqlmap", "custom_sqli"],
            fallback_techniques=["udf_injection", "stored_procedure_injection"],
        ),
        ChainStepTemplate(
            name="execute_commands",
            step_type=StepType.POST_EXPLOIT,
            description="Execute OS commands via webshell",
            techniques=["http_request", "command_injection"],
            required_conditions=["webshell_url"],
            produces=["command_output", "os_type", "current_privileges"],
            tools=["curl", "custom_http"],
        ),
    ],
    success_indicators=["webshell_accessible", "command_execution"],
    estimated_impact="Critical - Full server compromise",
    stealth_rating=3,
    complexity=6,
)

SQLI_TO_DATA_EXFIL_CHAIN = ChainTemplate(
    name="sqli_to_data_exfil",
    category=ChainCategory.DATA_EXFIL,
    description="SQL Injection → Database Enumeration → Data Exfiltration",
    entry_vuln_types=["sql_injection", "sqli", "blind_sqli"],
    steps=[
        ChainStepTemplate(
            name="confirm_sqli",
            step_type=StepType.VULNERABILITY,
            description="Confirm SQL injection",
            techniques=["error_based", "union_based", "blind_boolean"],
            required_conditions=["sqli_parameter"],
            produces=["dbms_type", "injection_type"],
            tools=["sqlmap"],
        ),
        ChainStepTemplate(
            name="enumerate_databases",
            step_type=StepType.RECONNAISSANCE,
            description="List all databases",
            techniques=["information_schema", "sysdatabases"],
            required_conditions=["dbms_type"],
            produces=["database_list"],
            tools=["sqlmap"],
        ),
        ChainStepTemplate(
            name="enumerate_tables",
            step_type=StepType.RECONNAISSANCE,
            description="List tables in target database",
            techniques=["information_schema.tables", "sysobjects"],
            required_conditions=["database_list"],
            produces=["table_list", "column_list"],
            tools=["sqlmap"],
        ),
        ChainStepTemplate(
            name="dump_sensitive_data",
            step_type=StepType.EXFILTRATION,
            description="Extract sensitive data (users, passwords, PII)",
            techniques=["union_dump", "blind_extraction", "out_of_band"],
            required_conditions=["table_list"],
            produces=["credentials", "pii_data", "sensitive_records"],
            tools=["sqlmap", "custom_sqli"],
        ),
    ],
    success_indicators=["credentials_extracted", "data_dumped"],
    estimated_impact="High - Data breach",
    stealth_rating=5,
    complexity=4,
)

# =============================================================================
# SSRF CHAINS
# =============================================================================

SSRF_TO_INTERNAL_CHAIN = ChainTemplate(
    name="ssrf_to_internal",
    category=ChainCategory.NETWORK,
    description="SSRF → Internal Network Access → Service Exploitation",
    entry_vuln_types=["ssrf", "server_side_request_forgery"],
    steps=[
        ChainStepTemplate(
            name="confirm_ssrf",
            step_type=StepType.VULNERABILITY,
            description="Confirm SSRF with external callback",
            techniques=["dns_callback", "http_callback", "time_based"],
            required_conditions=["ssrf_parameter"],
            produces=["ssrf_confirmed", "protocol_support"],
            tools=["burp_collaborator", "interactsh"],
        ),
        ChainStepTemplate(
            name="enumerate_protocols",
            step_type=StepType.RECONNAISSANCE,
            description="Test supported protocols (file://, gopher://, dict://)",
            techniques=["protocol_bruteforce", "schema_test"],
            required_conditions=["ssrf_confirmed"],
            produces=["supported_protocols", "filter_bypasses"],
            tools=["custom_ssrf"],
        ),
        ChainStepTemplate(
            name="scan_internal_network",
            step_type=StepType.RECONNAISSANCE,
            description="Scan internal network ranges",
            techniques=["port_scan", "service_detection", "host_discovery"],
            required_conditions=["ssrf_confirmed"],
            produces=["internal_hosts", "open_ports", "services"],
            tools=["custom_ssrf"],
        ),
        ChainStepTemplate(
            name="access_cloud_metadata",
            step_type=StepType.EXPLOITATION,
            description="Access cloud metadata services",
            techniques=[
                "aws_metadata",  # 169.254.169.254
                "gcp_metadata",
                "azure_metadata",
                "kubernetes_api",
            ],
            required_conditions=["ssrf_confirmed"],
            produces=["cloud_credentials", "iam_role", "tokens"],
            tools=["custom_ssrf"],
        ),
        ChainStepTemplate(
            name="exploit_internal_service",
            step_type=StepType.EXPLOITATION,
            description="Exploit discovered internal services",
            techniques=["redis_rce", "memcached_injection", "elasticsearch_rce"],
            required_conditions=["internal_hosts", "open_ports"],
            produces=["internal_access", "credentials"],
            tools=["custom_exploit"],
        ),
    ],
    success_indicators=["internal_access", "cloud_credentials", "service_compromised"],
    estimated_impact="Critical - Internal network access",
    stealth_rating=6,
    complexity=7,
)

# =============================================================================
# FILE UPLOAD CHAINS
# =============================================================================

UPLOAD_TO_SHELL_CHAIN = ChainTemplate(
    name="upload_to_shell",
    category=ChainCategory.WEB_APP,
    description="File Upload → Webshell → System Access",
    entry_vuln_types=["file_upload", "unrestricted_upload", "upload_bypass"],
    steps=[
        ChainStepTemplate(
            name="identify_upload",
            step_type=StepType.VULNERABILITY,
            description="Identify file upload functionality and restrictions",
            techniques=["extension_test", "mime_test", "content_test"],
            required_conditions=["upload_endpoint"],
            produces=["allowed_extensions", "validation_type", "storage_path"],
            tools=["burp", "custom_upload"],
        ),
        ChainStepTemplate(
            name="bypass_restrictions",
            step_type=StepType.EXPLOITATION,
            description="Bypass upload restrictions",
            techniques=[
                "double_extension",  # shell.php.jpg
                "null_byte",  # shell.php%00.jpg
                "mime_spoof",  # image/jpeg with PHP content
                "case_variation",  # shell.PhP
                "htaccess_upload",  # .htaccess to enable PHP
                "polyglot",  # Valid image + PHP code
            ],
            required_conditions=["allowed_extensions"],
            produces=["bypass_method", "uploaded_file"],
            tools=["custom_upload"],
        ),
        ChainStepTemplate(
            name="locate_uploaded_file",
            step_type=StepType.RECONNAISSANCE,
            description="Find uploaded file path/URL",
            techniques=["response_parse", "path_bruteforce", "predictable_path"],
            required_conditions=["uploaded_file"],
            produces=["file_url", "file_path"],
            tools=["custom_http"],
        ),
        ChainStepTemplate(
            name="execute_webshell",
            step_type=StepType.POST_EXPLOIT,
            description="Access and execute commands via webshell",
            techniques=["direct_access", "include_trigger", "wrapper_access"],
            required_conditions=["file_url"],
            produces=["command_execution", "os_access"],
            tools=["curl", "custom_http"],
        ),
    ],
    success_indicators=["webshell_accessible", "command_execution"],
    estimated_impact="Critical - Full server compromise",
    stealth_rating=4,
    complexity=5,
)

# =============================================================================
# XSS CHAINS
# =============================================================================

XSS_TO_SESSION_CHAIN = ChainTemplate(
    name="xss_to_session",
    category=ChainCategory.WEB_APP,
    description="XSS → Session Hijacking → Account Takeover",
    entry_vuln_types=["xss", "cross_site_scripting", "stored_xss", "reflected_xss"],
    steps=[
        ChainStepTemplate(
            name="confirm_xss",
            step_type=StepType.VULNERABILITY,
            description="Confirm XSS and identify context",
            techniques=["alert_test", "console_log", "dom_manipulation"],
            required_conditions=["xss_parameter"],
            produces=["xss_type", "context", "encoding_required"],
            tools=["burp", "custom_xss"],
        ),
        ChainStepTemplate(
            name="bypass_filters",
            step_type=StepType.EXPLOITATION,
            description="Bypass XSS filters and CSP",
            techniques=[
                "encoding_bypass",
                "tag_mutation",
                "event_handler_bypass",
                "csp_bypass",
                "dom_clobbering",
            ],
            required_conditions=["xss_type"],
            produces=["working_payload", "bypass_method"],
            tools=["custom_xss"],
        ),
        ChainStepTemplate(
            name="steal_session",
            step_type=StepType.EXPLOITATION,
            description="Exfiltrate session tokens",
            techniques=[
                "cookie_theft",
                "localstorage_theft",
                "sessionstorage_theft",
                "token_extraction",
            ],
            required_conditions=["working_payload"],
            produces=["session_token", "auth_cookies"],
            tools=["custom_xss", "burp_collaborator"],
        ),
        ChainStepTemplate(
            name="hijack_session",
            step_type=StepType.POST_EXPLOIT,
            description="Use stolen tokens for account access",
            techniques=["cookie_injection", "token_replay", "session_fixation"],
            required_conditions=["session_token"],
            produces=["authenticated_session", "user_access"],
            tools=["burp", "custom_http"],
        ),
    ],
    success_indicators=["session_stolen", "account_accessed"],
    estimated_impact="High - Account takeover",
    stealth_rating=7,
    complexity=5,
)

# =============================================================================
# LFI/RFI CHAINS
# =============================================================================

LFI_TO_RCE_CHAIN = ChainTemplate(
    name="lfi_to_rce",
    category=ChainCategory.WEB_APP,
    description="Local File Inclusion → Log Poisoning/Wrapper → RCE",
    entry_vuln_types=["lfi", "local_file_inclusion", "path_traversal"],
    steps=[
        ChainStepTemplate(
            name="confirm_lfi",
            step_type=StepType.VULNERABILITY,
            description="Confirm LFI with known file read",
            techniques=["etc_passwd", "win_ini", "proc_self"],
            required_conditions=["lfi_parameter"],
            produces=["lfi_confirmed", "os_type", "traversal_depth"],
            tools=["burp", "custom_lfi"],
        ),
        ChainStepTemplate(
            name="identify_rce_vector",
            step_type=StepType.RECONNAISSANCE,
            description="Find path to RCE (logs, wrappers, sessions)",
            techniques=[
                "log_location",
                "session_file",
                "php_wrappers",
                "proc_environ",
            ],
            required_conditions=["lfi_confirmed"],
            produces=["rce_vector", "writable_path"],
            tools=["custom_lfi"],
        ),
        ChainStepTemplate(
            name="poison_or_inject",
            step_type=StepType.EXPLOITATION,
            description="Inject PHP code into includable location",
            techniques=[
                "log_poisoning",  # User-Agent → access.log
                "session_poisoning",  # Session data → /tmp/sess_X
                "php_filter_chain",  # data://text/plain;base64
                "expect_wrapper",  # expect://id
                "zip_wrapper",  # zip://shell.zip#shell.php
            ],
            required_conditions=["rce_vector"],
            produces=["injected_code", "include_path"],
            tools=["custom_lfi"],
        ),
        ChainStepTemplate(
            name="trigger_rce",
            step_type=StepType.POST_EXPLOIT,
            description="Include poisoned file to execute code",
            techniques=["lfi_include", "wrapper_execute"],
            required_conditions=["include_path"],
            produces=["command_execution", "webshell"],
            tools=["custom_lfi"],
        ),
    ],
    success_indicators=["code_execution", "shell_access"],
    estimated_impact="Critical - Remote code execution",
    stealth_rating=4,
    complexity=6,
)

# =============================================================================
# POST-EXPLOITATION CHAINS
# =============================================================================

SHELL_TO_ROOT_CHAIN = ChainTemplate(
    name="shell_to_root",
    category=ChainCategory.POST_EXPLOIT,
    description="Low-Privilege Shell → Privilege Escalation → Root/SYSTEM",
    entry_vuln_types=["rce", "webshell", "reverse_shell"],
    steps=[
        ChainStepTemplate(
            name="enumerate_system",
            step_type=StepType.RECONNAISSANCE,
            description="Gather system information for privesc",
            techniques=[
                "kernel_version",
                "os_release",
                "running_processes",
                "network_config",
                "user_groups",
            ],
            required_conditions=["shell_access"],
            produces=["os_info", "kernel_version", "current_user", "groups"],
            tools=["linpeas", "winpeas", "custom_enum"],
        ),
        ChainStepTemplate(
            name="find_privesc_vectors",
            step_type=StepType.RECONNAISSANCE,
            description="Identify privilege escalation opportunities",
            techniques=[
                "suid_check",
                "sudo_misconfig",
                "cron_abuse",
                "writable_services",
                "kernel_exploits",
                "docker_escape",
            ],
            required_conditions=["os_info"],
            produces=["privesc_vectors", "recommended_exploit"],
            tools=["linpeas", "winpeas", "linux_exploit_suggester"],
        ),
        ChainStepTemplate(
            name="exploit_privesc",
            step_type=StepType.ESCALATION,
            description="Execute privilege escalation exploit",
            techniques=[
                "suid_exploit",
                "sudo_exploit",
                "kernel_exploit",
                "service_hijack",
                "token_impersonation",
            ],
            required_conditions=["privesc_vectors"],
            produces=["elevated_shell", "root_access"],
            tools=["custom_exploit", "metasploit"],
        ),
        ChainStepTemplate(
            name="establish_persistence",
            step_type=StepType.PERSISTENCE,
            description="Maintain access after reboot",
            techniques=[
                "ssh_key",
                "cron_backdoor",
                "service_backdoor",
                "rootkit",
                "registry_run_key",
            ],
            required_conditions=["root_access"],
            produces=["persistent_access", "backdoor"],
            tools=["custom_persist"],
        ),
    ],
    success_indicators=["root_shell", "persistent_access"],
    estimated_impact="Critical - Full system compromise",
    stealth_rating=3,
    complexity=7,
)

# =============================================================================
# TEMPLATE REGISTRY
# =============================================================================

CHAIN_TEMPLATES: dict[str, ChainTemplate] = {
    # SQLi chains
    "sqli_to_rce": SQLI_TO_RCE_CHAIN,
    "sqli_to_data_exfil": SQLI_TO_DATA_EXFIL_CHAIN,
    # SSRF chains
    "ssrf_to_internal": SSRF_TO_INTERNAL_CHAIN,
    # Upload chains
    "upload_to_shell": UPLOAD_TO_SHELL_CHAIN,
    # XSS chains
    "xss_to_session": XSS_TO_SESSION_CHAIN,
    # LFI chains
    "lfi_to_rce": LFI_TO_RCE_CHAIN,
    # Post-exploit chains
    "shell_to_root": SHELL_TO_ROOT_CHAIN,
}


def get_chain_template(name: str) -> ChainTemplate | None:
    """Get a chain template by name."""
    return CHAIN_TEMPLATES.get(name)


def get_chains_for_vuln_type(vuln_type: str) -> list[ChainTemplate]:
    """Get all chain templates that can start from a given vulnerability type."""
    vuln_type_lower = vuln_type.lower()
    matching_chains = []

    for chain in CHAIN_TEMPLATES.values():
        if any(vuln_type_lower in entry.lower() for entry in chain.entry_vuln_types):
            matching_chains.append(chain)

    # Sort by complexity (simpler first) then by impact
    matching_chains.sort(key=lambda c: (c.complexity, -c.stealth_rating))
    return matching_chains


def get_chains_by_category(category: ChainCategory) -> list[ChainTemplate]:
    """Get all chain templates in a category."""
    return [chain for chain in CHAIN_TEMPLATES.values() if chain.category == category]


__all__ = [
    "ChainCategory",
    "StepType",
    "ChainStepTemplate",
    "ChainTemplate",
    "CHAIN_TEMPLATES",
    "get_chain_template",
    "get_chains_for_vuln_type",
    "get_chains_by_category",
    # Individual templates
    "SQLI_TO_RCE_CHAIN",
    "SQLI_TO_DATA_EXFIL_CHAIN",
    "SSRF_TO_INTERNAL_CHAIN",
    "UPLOAD_TO_SHELL_CHAIN",
    "XSS_TO_SESSION_CHAIN",
    "LFI_TO_RCE_CHAIN",
    "SHELL_TO_ROOT_CHAIN",
]
