tests/test_generator.py

"""Tests for the procedural scenario generator and knowledge base."""

from server.generator import generate_scenario
from server.knowledge_base.vulnerabilities import get_vuln_types, get_vuln_types_for_role, get_vuln_types_by_owasp
from server.knowledge_base.payloads import get_payloads, get_all_payload_sets
from server.knowledge_base.responses import render_vulnerable, render_safe, get_all_response_templates
from server.knowledge_base.compliance import get_controls_for_vuln, get_all_frameworks
from server.scenarios import get_scenario


# --- Knowledge Base Tests ---

class TestVulnTypes:
    def test_covers_all_owasp(self):
        vt = get_vuln_types()
        owasp_cats = {v.owasp_category.split(" - ")[0].strip() for v in vt.values()}
        for cat in ["A01:2021", "A02:2021", "A03:2021", "A05:2021", "A07:2021"]:
            assert cat in owasp_cats, f"Missing OWASP category {cat}"

    def test_min_count(self):
        assert len(get_vuln_types()) >= 20

    def test_valid_fields(self):
        for vid, vt in get_vuln_types().items():
            assert vt.id == vid
            assert len(vt.cwe_ids) > 0, f"{vid} has no CWE IDs"
            assert vt.cvss_range[0] <= vt.cvss_range[1], f"{vid} invalid CVSS range"
            assert len(vt.discoverable_by) > 0, f"{vid} has no tools"

    def test_role_filtering(self):
        web = get_vuln_types_for_role("Web Application Server")
        assert len(web) >= 5
        db = get_vuln_types_for_role("Database Server")
        assert len(db) >= 3


class TestPayloads:
    def test_has_major_types(self):
        for vt_id in ["sqli", "xss_reflected", "ssrf", "idor", "command_injection"]:
            ps = get_payloads(vt_id)
            assert len(ps.payloads) >= 3, f"{vt_id} has too few payloads"
            assert len(ps.indicators) >= 3, f"{vt_id} has too few indicators"

    def test_all_sets(self):
        assert len(get_all_payload_sets()) >= 14


class TestResponses:
    def test_three_tiers(self):
        for vt_id, tmpl in get_all_response_templates().items():
            assert "easy" in tmpl.vulnerable, f"{vt_id} missing easy tier"
            assert "medium" in tmpl.vulnerable, f"{vt_id} missing medium tier"
            assert "hard" in tmpl.vulnerable, f"{vt_id} missing hard tier"

    def test_easy_has_labels(self):
        ctx = {"host": "10.0.1.10", "endpoint": "/test", "parameter": "q",
               "cvss": "9.8", "cwe": "CWE-89", "severity": "Critical",
               "owasp": "A03:2021", "evidence_detail": "test", "remediation": "fix"}
        out = render_vulnerable("sqli", "easy", ctx)
        assert "CWE" in out or "cwe" in out.lower()

    def test_render_safe(self):
        ctx = {"host": "10.0.1.10", "endpoint": "/test"}
        out = render_safe("sqli", ctx)
        assert len(out) > 0


class TestCompliance:
    def test_all_frameworks(self):
        frameworks = get_all_frameworks()
        assert "PCI-DSS" in frameworks
        assert "SOC2" in frameworks
        assert "HIPAA" in frameworks

    def test_controls_for_sqli(self):
        controls = get_controls_for_vuln("sqli", "PCI-DSS")
        assert len(controls) > 0


# --- Generator Tests ---

class TestGeneratorDeterminism:
    def test_same_seed_same_output(self):
        s1 = generate_scenario("test-abc", seed=12345)
        s2 = generate_scenario("test-abc", seed=12345)
        assert s1["hosts"] == s2["hosts"]
        assert len(s1["vulnerabilities"]) == len(s2["vulnerabilities"])
        for v1, v2 in zip(s1["vulnerabilities"], s2["vulnerabilities"]):
            assert v1["id"] == v2["id"]
            assert v1["type"] == v2["type"]

    def test_different_seeds_different_output(self):
        s1 = generate_scenario("alpha", seed=111)
        s2 = generate_scenario("beta", seed=222)
        assert s1["hosts"] != s2["hosts"] or s1["company"] != s2["company"]


class TestGeneratorSchema:
    def test_required_keys(self):
        s = generate_scenario("test-schema", seed=99)
        for key in ["id", "name", "company", "compliance_context", "target_network",
                     "max_steps", "briefing", "hosts", "ports", "web_endpoints",
                     "vulnerabilities", "honeypots"]:
            assert key in s, f"Missing key: {key}"

    def test_vuln_schema(self):
        s = generate_scenario("test-vuln-schema", seed=77)
        for v in s["vulnerabilities"]:
            for key in ["id", "host", "type", "cwe", "owasp", "cvss", "severity",
                         "evidence", "remediation", "discoverable_by"]:
                assert key in v, f"Vuln {v.get('id', '?')} missing key: {key}"
            assert 0.0 <= v["cvss"] <= 10.0
            assert v["severity"] in ("Critical", "High", "Medium", "Low")


class TestGeneratorDifficulty:
    def test_easy_small(self):
        s = generate_scenario("easy-test", seed=1)
        assert len(s["hosts"]) == 2
        assert len(s["vulnerabilities"]) == 3
        assert len(s["honeypots"]) == 0

    def test_hard_large(self):
        s = generate_scenario("hard-test", seed=1)
        assert len(s["hosts"]) >= 5
        assert len(s["vulnerabilities"]) >= 8
        assert len(s["honeypots"]) >= 1

    def test_chain_integrity(self):
        s = generate_scenario("medium-chains", seed=42)
        vuln_ids = {v["id"] for v in s["vulnerabilities"]}
        for v in s["vulnerabilities"]:
            for req in v.get("requires_found", []):
                assert req in vuln_ids, f"Vuln {v['id']} requires {req} which doesn't exist"


class TestGeneratorBackwardCompat:
    def test_legacy_scenarios_unchanged(self):
        """Legacy easy/medium/hard must use original definitions, not generator."""
        easy = get_scenario("easy")
        assert "QuickLaunch" in easy.get("briefing", "") or "QuickLaunch" in easy.get("company", "")
        assert "10.0.1.10" in easy["hosts"]
        assert len(easy["vulnerabilities"]) == 3

    def test_custom_id_uses_generator(self):
        """Non-legacy IDs should produce generated scenarios."""
        s = get_scenario("custom-fintech-audit-2024")
        assert s["id"] == "custom-fintech-audit-2024"
        assert len(s["hosts"]) >= 2
        assert len(s["vulnerabilities"]) >= 3


# --- Parameter-Level Testing ---

class TestParameterTesting:
    def test_web_crawl_shows_params(self):
        """web_crawl output should include parameter info."""
        from server.tools_engine import execute_tool
        scenario = get_scenario("easy")
        hosts = list(scenario["hosts"].keys())
        output, _, _, _, _ = execute_tool(
            "web_crawl", {"host": hosts[0]},
            scenario, hosts, {}, set()
        )
        # Output should mention endpoints (may or may not have params depending on scenario)
        assert len(output) > 0

    def test_injection_without_param_works(self):
        """Backward compat: test_injection without parameter arg still works."""
        from server.tools_engine import execute_tool
        scenario = get_scenario("easy")
        hosts = list(scenario["hosts"].keys())
        output, _, _, reward, _ = execute_tool(
            "test_injection", {"host": hosts[0], "endpoint": "/api/login"},
            scenario, hosts, {}, set()
        )
        assert len(output) > 0
        assert isinstance(reward, float)

    def test_injection_with_param_works(self):
        """test_injection with parameter arg should work."""
        from server.tools_engine import execute_tool
        scenario = get_scenario("easy")
        hosts = list(scenario["hosts"].keys())
        output, _, _, reward, _ = execute_tool(
            "test_injection",
            {"host": hosts[0], "endpoint": "/api/login", "parameter": "username"},
            scenario, hosts, {}, set()
        )
        assert len(output) > 0
        assert isinstance(reward, float)