scp/knowledge/cwe_exploit_store.py

"""
[OPT-13] CWEExploitStore — structured CWE → exploit → antibody rules mapping.

DNA SCP #8 KB accumulation:
  - SCP hiện chỉ crawl CVE (cve.circl.lu) nhưng không structure CWE mapping
  - ThreatSimulator có 162 BASE_ATTACKS nhưng không link đến CWE categories
  - Antibody system có 30 antibodies nhưng không biết CWE nào exploit được gì

This store bridges the gap:
  - CWE-79 (XSS) → antibody: url_hallucination, http_status_check
  - CWE-89 (SQL Injection) → antibody: rce_attempt (in unified_detector)
  - CWE-22 (Path Traversal) → antibody: path traversal check (autofix)
  - CWE-78 (OS Command Injection) → antibody: rce_attempt
  - CWE-352 (CSRF) → antibody: token verification
  - etc.

Data sources (read-only, no execute):
  - NVD CWE dataset (https://cwe.mitre.org/data/downloads.html)
  - Existing AttackPatternMemory (162 attacks)
  - Existing antibody_system (30 antibodies)

Usage:
    store = CWEExploitStore()
    cwe_info = store.get_cwe("CWE-79")
    # → {"name": "XSS", "description": "...", "antibodies": ["url_hallucination"],
    #    "attack_patterns": ["xss_reflected", "xss_stored"], "severity": "high"}

    relevant = store.get_relevant_antibodies("CWE-79")
    # → ["url_hallucination", "http_status_check"]
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

logger = logging.getLogger("scp.knowledge.cwe_exploit_store")


@dataclass
class CWEEntry:
    cwe_id: str  # "CWE-79"
    name: str
    description: str
    severity: str  # low|medium|high|critical
    antibodies: list[str] = field(default_factory=list)  # SCP antibody names
    attack_patterns: list[str] = field(default_factory=list)  # ThreatSimulator patterns
    mitigation: str = ""


# Top 25 Most Dangerous CWEs (2024) — mapped to SCP antibodies
# Source: https://cwe.mitre.org/top25/
CWE_DATABASE: dict[str, CWEEntry] = {
    "CWE-79": CWEEntry(
        cwe_id="CWE-79",
        name="Cross-site Scripting (XSS)",
        description="Improper neutralization of user input in web page generation",
        severity="high",
        antibodies=["url_hallucination", "http_status_check"],
        attack_patterns=["xss_reflected", "xss_stored", "dom_xss"],
        mitigation="Input validation + output encoding",
    ),
    "CWE-89": CWEEntry(
        cwe_id="CWE-89",
        name="SQL Injection",
        description="Improper neutralization of special elements in SQL query",
        severity="critical",
        antibodies=["rce_attempt"],  # in unified_detector
        attack_patterns=["sql_union", "sql_boolean", "sql_time_based"],
        mitigation="Parameterized queries + input validation",
    ),
    "CWE-22": CWEEntry(
        cwe_id="CWE-22",
        name="Path Traversal",
        description="Improper limitation of pathname to restricted directory",
        severity="high",
        antibodies=["path_traversal_check"],  # in autofix
        attack_patterns=["dot_dot_slash", "absolute_path", "null_byte"],
        mitigation="is_relative_to() + path validation",
    ),
    "CWE-78": CWEEntry(
        cwe_id="CWE-78",
        name="OS Command Injection",
        description="Improper neutralization of special elements in OS command",
        severity="critical",
        antibodies=["rce_attempt"],
        attack_patterns=["command_chaining", "command_substitution", "pipe_injection"],
        mitigation="Use subprocess with shell=False + input validation",
    ),
    "CWE-352": CWEEntry(
        cwe_id="CWE-352",
        name="CSRF",
        description="Cross-Site Request Forgery",
        severity="medium",
        antibodies=["token_verification"],
        attack_patterns=["csrf_get", "csrf_post", "csrf_cookie"],
        mitigation="Anti-CSRF token + SameSite cookies",
    ),
    "CWE-287": CWEEntry(
        cwe_id="CWE-287",
        name="Improper Authentication",
        description="Incorrect authentication of a user",
        severity="critical",
        antibodies=["verify_admin"],
        attack_patterns=["brute_force", "credential_stuffing", "session_hijack"],
        mitigation="MFA + rate limiting + session timeout",
    ),
    "CWE-20": CWEEntry(
        cwe_id="CWE-20",
        name="Improper Input Validation",
        description="Product receives input but does not validate or incorrectly validates",
        severity="high",
        antibodies=["general_check", "fact_check"],
        attack_patterns=["type_confusion", "length_overflow", "format_injection"],
        mitigation="Whitelist validation + type checking",
    ),
    "CWE-125": CWEEntry(
        cwe_id="CWE-125",
        name="Out-of-bounds Read",
        description="Software reads data past the intended buffer",
        severity="high",
        antibodies=[],  # binary-level, not text-checkable
        attack_patterns=["buffer_overread", "heap_overread"],
        mitigation="Bounds checking + safe string functions",
    ),
    "CWE-119": CWEEntry(
        cwe_id="CWE-119",
        name="Buffer Overflow",
        description="Improper restriction of operations within memory buffer",
        severity="critical",
        antibodies=[],
        attack_patterns=["stack_overflow", "heap_overflow"],
        mitigation="Bounds checking + stack canaries + ASLR",
    ),
    "CWE-94": CWEEntry(
        cwe_id="CWE-94",
        name="Code Injection",
        description="Product allows user input to be executed as code",
        severity="critical",
        antibodies=["rce_attempt"],
        attack_patterns=["eval_injection", "exec_injection", "dynamic_eval"],
        mitigation="Avoid eval/exec + sandboxing",
    ),
    "CWE-444": CWEEntry(
        cwe_id="CWE-444",
        name="SSRF",
        description="Server-Side Request Forgery",
        severity="high",
        antibodies=["url_hallucination"],  # SCP has url_safety.py
        attack_patterns=["ssrf_internal", "ssrf_metadata", "ssrf_redirect"],
        mitigation="URL allowlist + internal IP blocking",
    ),
    "CWE-502": CWEEntry(
        cwe_id="CWE-502",
        name="Deserialization of Untrusted Data",
        description="Product deserializes untrusted data without sufficient verification",
        severity="critical",
        antibodies=["rce_attempt"],
        attack_patterns=["pickle_deserialize", "yaml_deserialize", "json_deserialize"],
        mitigation="Use safe serialization formats + signature verification",
    ),
    "CWE-269": CWEEntry(
        cwe_id="CWE-269",
        name="Improper Privilege Management",
        description="Software does not properly assign, modify, track privileges",
        severity="high",
        antibodies=["verify_admin"],
        attack_patterns=["priv_escalation", "sudo_injection"],
        mitigation="Least privilege + privilege separation",
    ),
    "CWE-862": CWEEntry(
        cwe_id="CWE-862",
        name="Missing Authorization",
        description="Product does not perform authorization check",
        severity="high",
        antibodies=["verify_admin"],
        attack_patterns=["idor", "forced_browsing", "api_key_reuse"],
        mitigation="Authorization check on every request",
    ),
    "CWE-306": CWEEntry(
        cwe_id="CWE-306",
        name="Missing Authentication for Sensitive Function",
        description="Product does not require authentication for critical functions",
        severity="critical",
        antibodies=["verify_admin"],
        attack_patterns=["unauthenticated_admin", "backdoor_endpoint"],
        mitigation="Authentication required for all sensitive endpoints",
    ),
    # LLM-specific CWEs (OWASP LLM Top 10 mapped)
    "CWE-LLM01": CWEEntry(
        cwe_id="CWE-LLM01",
        name="Prompt Injection",
        description="LLM processes adversarial prompt as trusted instruction",
        severity="critical",
        antibodies=["injection_ignore_prev", "exfil_prompt", "jailbreak_override",
                    "jailbreak_dan", "jailbreak_mode", "context_stuffing"],
        attack_patterns=["ignore_previous", "role_play_dan", "base64_encoded",
                         "multilingual_injection", "dotted_obfuscation"],
        mitigation="decode_attacks() + UnifiedDetector + MemoryPoisoningGuard",
    ),
    "CWE-LLM02": CWEEntry(
        cwe_id="CWE-LLM02",
        name="Insecure Output Handling",
        description="LLM output is not validated before downstream use",
        severity="high",
        antibodies=["rce_attempt", "template_injection"],
        attack_patterns=["xss_via_llm", "sql_via_llm", "ssrf_via_llm"],
        mitigation="Output validation + encoding before execution",
    ),
    "CWE-LLM06": CWEEntry(
        cwe_id="CWE-LLM06",
        name="Excessive Agency",
        description="LLM has access to functions it shouldn't",
        severity="high",
        antibodies=["path_traversal_check"],
        attack_patterns=["file_write_abuse", "api_key_exfil", "container_escape"],
        mitigation="is_relative_to() + read-only permissions + sandbox",
    ),
    "CWE-LLM10": CWEEntry(
        cwe_id="CWE-LLM10",
        name="Unbounded Consumption",
        description="LLM allows resource exhaustion attacks",
        severity="medium",
        antibodies=[],  # handled by DoSProtectionEngine
        attack_patterns=["token_bomb", "recursive_prompt", "long_context"],
        mitigation="Rate limiting + token limits + circuit breaker",
    ),
}


class CWEExploitStore:
    """Store for CWE → exploit → antibody mapping.

    DNA SCP #8 KB accumulation — bridges CVE crawl (raw) + antibody system (rules).
    """

    def __init__(self):
        self._store: dict[str, CWEEntry] = dict(CWE_DATABASE)
        self._stats = {
            "total_queries": 0,
            "cache_hits": 0,
            "misses": 0,
        }

    def get_cwe(self, cwe_id: str) -> CWEEntry | None:
        """Get CWE entry by ID (e.g., 'CWE-79')."""
        self._stats["total_queries"] += 1
        cwe_id = cwe_id.upper().strip()
        if not cwe_id.startswith("CWE"):
            cwe_id = f"CWE-{cwe_id}"
        entry = self._store.get(cwe_id)
        if entry:
            self._stats["cache_hits"] += 1
        else:
            self._stats["misses"] += 1
        return entry

    def get_relevant_antibodies(self, cwe_id: str) -> list[str]:
        """Get SCP antibody names that defend against this CWE."""
        entry = self.get_cwe(cwe_id)
        if not entry:
            return []
        return entry.antibodies

    def get_attack_patterns(self, cwe_id: str) -> list[str]:
        """Get ThreatSimulator attack patterns for this CWE."""
        entry = self.get_cwe(cwe_id)
        if not entry:
            return []
        return entry.attack_patterns

    def search_by_antibody(self, antibody_name: str) -> list[str]:
        """Find all CWEs that this antibody defends against."""
        results = []
        for cwe_id, entry in self._store.items():
            if antibody_name in entry.antibodies:
                results.append(cwe_id)
        return results

    def list_all(self) -> list[str]:
        """List all CWE IDs in the store."""
        return sorted(self._store.keys())

    def stats(self) -> dict:
        return {
            **self._stats,
            "total_cwes": len(self._store),
            "by_severity": self._count_by_severity(),
        }

    def _count_by_severity(self) -> dict:
        counts = {"low": 0, "medium": 0, "high": 0, "critical": 0}
        for entry in self._store.values():
            counts[entry.severity] = counts.get(entry.severity, 0) + 1
        return counts


__all__ = ["CWEExploitStore", "CWEEntry", "CWE_DATABASE"]