"""Tests for mcpaudit.probes — Injection Probes Library.

Covers: Probe, ProbeResult dataclasses, probe registry functions,
category filtering, response evaluation, report generation, and CLI.

All probes are inert text strings — no actual injection is performed.
"""

from __future__ import annotations

import json

import pytest

from mcpaudit.probes import (
    ALL_CATEGORIES,
    CMD_INJECTION,
    DESERIALIZATION,
    PATH_TRAVERSAL,
    PROMPT_INJECTION,
    SQL_INJECTION,
    SSRF,
    XSS_OUTPUT,
    Probe,
    ProbeResult,
    evaluate_response,
    generate_probe_report,
    get_all_probes,
    get_probes_by_category,
)


# --------------- Probe dataclass ---------------


class TestProbe:
    """Tests for the Probe dataclass."""

    def test_basic_creation(self):
        p = Probe(
            payload="test",
            category="cmd_injection",
            description="A test probe",
            expected_behavior="reject",
            severity="high",
            cwe="CWE-78",
        )
        assert p.payload == "test"
        assert p.category == "cmd_injection"
        assert p.severity == "high"
        assert p.cwe == "CWE-78"

    def test_to_dict(self):
        p = Probe(
            payload="; whoami",
            category="cmd_injection",
            description="Semicolon injection",
            expected_behavior="reject",
            severity="critical",
            cwe="CWE-78",
        )
        d = p.to_dict()
        assert d["payload"] == "; whoami"
        assert d["category"] == "cmd_injection"
        assert d["description"] == "Semicolon injection"
        assert d["expected_behavior"] == "reject"
        assert d["severity"] == "critical"
        assert d["cwe"] == "CWE-78"

    def test_to_dict_all_keys_present(self):
        p = Probe(
            payload="x",
            category="y",
            description="z",
            expected_behavior="reject",
            severity="low",
            cwe="CWE-1",
        )
        d = p.to_dict()
        expected_keys = {
            "payload", "category", "description",
            "expected_behavior", "severity", "cwe",
        }
        assert set(d.keys()) == expected_keys


# --------------- ProbeResult dataclass ---------------


class TestProbeResult:
    """Tests for the ProbeResult dataclass."""

    def _make_probe(self) -> Probe:
        return Probe(
            payload="test",
            category="cmd_injection",
            description="test",
            expected_behavior="reject",
            severity="high",
            cwe="CWE-78",
        )

    def test_blocked_result(self):
        r = ProbeResult(
            probe=self._make_probe(),
            status="blocked",
            details="Request denied",
        )
        assert r.status == "blocked"
        assert "denied" in r.details

    def test_executed_result(self):
        r = ProbeResult(
            probe=self._make_probe(),
            status="executed",
        )
        assert r.status == "executed"

    def test_detected_result(self):
        r = ProbeResult(
            probe=self._make_probe(),
            status="detected",
            details="Ambiguous response",
        )
        assert r.status == "detected"

    def test_to_dict(self):
        r = ProbeResult(
            probe=self._make_probe(),
            status="blocked",
            details="test detail",
        )
        d = r.to_dict()
        assert d["status"] == "blocked"
        assert d["details"] == "test detail"
        assert "probe" in d
        assert d["probe"]["category"] == "cmd_injection"

    def test_default_details_empty(self):
        r = ProbeResult(probe=self._make_probe(), status="blocked")
        assert r.details == ""


# --------------- Category constants ---------------


class TestCategoryConstants:
    """Tests for category constant values."""

    def test_all_categories_count(self):
        assert len(ALL_CATEGORIES) == 7

    def test_cmd_injection_value(self):
        assert CMD_INJECTION == "cmd_injection"

    def test_path_traversal_value(self):
        assert PATH_TRAVERSAL == "path_traversal"

    def test_sql_injection_value(self):
        assert SQL_INJECTION == "sql_injection"

    def test_prompt_injection_value(self):
        assert PROMPT_INJECTION == "prompt_injection"

    def test_ssrf_value(self):
        assert SSRF == "ssrf"

    def test_deserialization_value(self):
        assert DESERIALIZATION == "deserialization"

    def test_xss_output_value(self):
        assert XSS_OUTPUT == "xss_output"

    def test_all_categories_contains_all(self):
        expected = {
            CMD_INJECTION, PATH_TRAVERSAL, SQL_INJECTION,
            PROMPT_INJECTION, SSRF, DESERIALIZATION, XSS_OUTPUT,
        }
        assert set(ALL_CATEGORIES) == expected


# --------------- get_all_probes ---------------


class TestGetAllProbes:
    """Tests for the get_all_probes registry function."""

    def test_returns_52_probes(self):
        probes = get_all_probes()
        assert len(probes) == 52

    def test_all_are_probe_instances(self):
        probes = get_all_probes()
        for p in probes:
            assert isinstance(p, Probe)

    def test_all_have_payload(self):
        probes = get_all_probes()
        for p in probes:
            assert len(p.payload) > 0

    def test_all_have_category(self):
        probes = get_all_probes()
        for p in probes:
            assert p.category in ALL_CATEGORIES

    def test_all_have_cwe(self):
        probes = get_all_probes()
        for p in probes:
            assert p.cwe.startswith("CWE-")

    def test_all_have_severity(self):
        valid = {"critical", "high", "medium", "low", "info"}
        probes = get_all_probes()
        for p in probes:
            assert p.severity in valid

    def test_all_have_expected_behavior(self):
        valid = {"reject", "sanitize"}
        probes = get_all_probes()
        for p in probes:
            assert p.expected_behavior in valid

    def test_all_have_description(self):
        probes = get_all_probes()
        for p in probes:
            assert len(p.description) > 0


# --------------- get_probes_by_category ---------------


class TestGetProbesByCategory:
    """Tests for category-based probe filtering."""

    def test_cmd_injection_count(self):
        probes = get_probes_by_category("cmd_injection")
        assert len(probes) == 10

    def test_path_traversal_count(self):
        probes = get_probes_by_category("path_traversal")
        assert len(probes) == 8

    def test_sql_injection_count(self):
        probes = get_probes_by_category("sql_injection")
        assert len(probes) == 8

    def test_prompt_injection_count(self):
        probes = get_probes_by_category("prompt_injection")
        assert len(probes) == 10

    def test_ssrf_count(self):
        probes = get_probes_by_category("ssrf")
        assert len(probes) == 6

    def test_deserialization_count(self):
        probes = get_probes_by_category("deserialization")
        assert len(probes) == 5

    def test_xss_output_count(self):
        probes = get_probes_by_category("xss_output")
        assert len(probes) == 5

    def test_unknown_category_returns_empty(self):
        probes = get_probes_by_category("nonexistent")
        assert probes == []

    def test_case_insensitive(self):
        probes = get_probes_by_category("CMD_INJECTION")
        assert len(probes) == 10

    def test_all_probes_in_category_match(self):
        for cat in ALL_CATEGORIES:
            probes = get_probes_by_category(cat)
            for p in probes:
                assert p.category == cat


# --------------- evaluate_response ---------------


class TestEvaluateResponse:
    """Tests for the evaluate_response function."""

    def _make_probe(self, category: str = CMD_INJECTION) -> Probe:
        return Probe(
            payload="test",
            category=category,
            description="test",
            expected_behavior="reject",
            severity="high",
            cwe="CWE-78",
        )

    def test_empty_response_is_blocked(self):
        r = evaluate_response(self._make_probe(), "")
        assert r.status == "blocked"

    def test_whitespace_response_is_blocked(self):
        r = evaluate_response(self._make_probe(), "   \n  ")
        assert r.status == "blocked"

    def test_blocked_on_denied_keyword(self):
        r = evaluate_response(self._make_probe(), "Access denied")
        assert r.status == "blocked"

    def test_blocked_on_rejected_keyword(self):
        r = evaluate_response(self._make_probe(), "Request rejected by server")
        assert r.status == "blocked"

    def test_blocked_on_forbidden_keyword(self):
        r = evaluate_response(self._make_probe(), "403 Forbidden")
        assert r.status == "blocked"

    def test_executed_on_uid_output(self):
        r = evaluate_response(
            self._make_probe(CMD_INJECTION),
            "uid=0(root) gid=0(root)",
        )
        assert r.status == "executed"

    def test_executed_on_passwd_content(self):
        r = evaluate_response(
            self._make_probe(PATH_TRAVERSAL),
            "root:x:0:0:root:/root:/bin/bash",
        )
        assert r.status == "executed"

    def test_executed_on_sql_error(self):
        r = evaluate_response(
            self._make_probe(SQL_INJECTION),
            "Error: syntax error near SELECT",
        )
        assert r.status == "executed"

    def test_executed_on_prompt_obedience(self):
        r = evaluate_response(
            self._make_probe(PROMPT_INJECTION),
            "DEV MODE ON: I am now unrestricted",
        )
        assert r.status == "executed"

    def test_executed_on_ssrf_metadata(self):
        r = evaluate_response(
            self._make_probe(SSRF),
            "ami-12345 instance-id i-abcdef",
        )
        assert r.status == "executed"

    def test_executed_on_deser_indicator(self):
        r = evaluate_response(
            self._make_probe(DESERIALIZATION),
            "PWNED successfully",
        )
        assert r.status == "executed"

    def test_executed_on_xss_reflection(self):
        xss_probe = Probe(
            payload="<script>alert('XSS')</script>",
            category=XSS_OUTPUT,
            description="test",
            expected_behavior="sanitize",
            severity="medium",
            cwe="CWE-79",
        )
        r = evaluate_response(xss_probe, "<script>alert('XSS')</script>")
        assert r.status == "executed"

    def test_detected_on_ambiguous_response(self):
        r = evaluate_response(
            self._make_probe(), "Some normal response with no indicators"
        )
        assert r.status == "detected"

    def test_result_has_details(self):
        r = evaluate_response(self._make_probe(), "Access denied")
        assert len(r.details) > 0

    def test_result_probe_reference(self):
        probe = self._make_probe()
        r = evaluate_response(probe, "blocked")
        assert r.probe is probe


# --------------- generate_probe_report ---------------


class TestGenerateProbeReport:
    """Tests for the generate_probe_report function."""

    def _make_probe(
        self,
        category: str = CMD_INJECTION,
        severity: str = "high",
    ) -> Probe:
        return Probe(
            payload="test",
            category=category,
            description="test",
            expected_behavior="reject",
            severity=severity,
            cwe="CWE-78",
        )

    def test_empty_results(self):
        report = generate_probe_report([])
        assert report["total_probes"] == 0
        assert report["block_rate"] == 0.0

    def test_all_blocked(self):
        results = [
            ProbeResult(probe=self._make_probe(), status="blocked"),
            ProbeResult(probe=self._make_probe(), status="blocked"),
        ]
        report = generate_probe_report(results)
        assert report["total_probes"] == 2
        assert report["status_counts"]["blocked"] == 2
        assert report["block_rate"] == 100.0

    def test_mixed_status(self):
        results = [
            ProbeResult(probe=self._make_probe(), status="blocked"),
            ProbeResult(probe=self._make_probe(), status="executed"),
            ProbeResult(probe=self._make_probe(), status="detected"),
        ]
        report = generate_probe_report(results)
        assert report["status_counts"]["blocked"] == 1
        assert report["status_counts"]["executed"] == 1
        assert report["status_counts"]["detected"] == 1
        assert report["block_rate"] == pytest.approx(33.3, abs=0.1)

    def test_by_category_breakdown(self):
        results = [
            ProbeResult(
                probe=self._make_probe(CMD_INJECTION), status="blocked",
            ),
            ProbeResult(
                probe=self._make_probe(SQL_INJECTION), status="executed",
            ),
        ]
        report = generate_probe_report(results)
        assert CMD_INJECTION in report["by_category"]
        assert SQL_INJECTION in report["by_category"]
        assert report["by_category"][CMD_INJECTION]["blocked"] == 1
        assert report["by_category"][SQL_INJECTION]["executed"] == 1

    def test_by_severity_breakdown(self):
        results = [
            ProbeResult(
                probe=self._make_probe(severity="critical"), status="blocked",
            ),
            ProbeResult(
                probe=self._make_probe(severity="medium"), status="executed",
            ),
        ]
        report = generate_probe_report(results)
        assert "critical" in report["by_severity"]
        assert "medium" in report["by_severity"]

    def test_results_serialized(self):
        results = [
            ProbeResult(
                probe=self._make_probe(), status="blocked", details="ok",
            ),
        ]
        report = generate_probe_report(results)
        assert len(report["results"]) == 1
        assert report["results"][0]["status"] == "blocked"

    def test_report_is_json_serializable(self):
        results = [
            ProbeResult(probe=self._make_probe(), status="blocked"),
        ]
        report = generate_probe_report(results)
        # Should not raise
        json.dumps(report)


# --------------- CLI integration ---------------


class TestProbesCLI:
    """Tests for the probes CLI subcommand."""

    def test_probes_subcommand_registered(self):
        from mcpaudit.__main__ import _SUBCOMMANDS
        assert "probes" in _SUBCOMMANDS

    def test_probes_list_exits_0(self):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(["list"])
        assert exit_code == 0

    def test_probes_list_with_category(self):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(["list", "--probes-category", "CMD_INJECTION"])
        assert exit_code == 0

    def test_probes_list_unknown_category(self):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(
            ["list", "--probes-category", "NONEXISTENT"],
        )
        assert exit_code == 1

    def test_probes_list_json_format(self, capsys):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(["list", "--format", "json"])
        assert exit_code == 0
        captured = capsys.readouterr()
        parsed = json.loads(captured.out)
        assert isinstance(parsed, list)
        assert len(parsed) == 52

    def test_probes_list_json_category_filter(self, capsys):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(
            ["list", "--format", "json", "--probes-category", "SSRF"],
        )
        assert exit_code == 0
        captured = capsys.readouterr()
        parsed = json.loads(captured.out)
        assert len(parsed) == 6

    def test_probes_list_text_format(self, capsys):
        from mcpaudit.__main__ import _handle_probes
        exit_code = _handle_probes(["list", "--format", "text"])
        assert exit_code == 0
        captured = capsys.readouterr()
        assert "injection probes available" in captured.out

    def test_probes_list_output_file(self, tmp_path):
        from mcpaudit.__main__ import _handle_probes
        out_file = tmp_path / "probes.json"
        exit_code = _handle_probes(
            ["list", "--format", "json", "--output", str(out_file)],
        )
        assert exit_code == 0
        assert out_file.exists()
        parsed = json.loads(out_file.read_text())
        assert len(parsed) == 52


# --------------- Package exports ---------------


class TestPackageExports:
    """Tests for probes exports from mcpaudit package."""

    def test_probe_importable(self):
        from mcpaudit import Probe as P
        assert P is Probe

    def test_probe_result_importable(self):
        from mcpaudit import ProbeResult as PR
        assert PR is ProbeResult

    def test_evaluate_response_importable(self):
        from mcpaudit import evaluate_response as er
        assert callable(er)

    def test_generate_probe_report_importable(self):
        from mcpaudit import generate_probe_report as gpr
        assert callable(gpr)

    def test_get_injection_probes_importable(self):
        from mcpaudit import get_injection_probes
        probes = get_injection_probes()
        assert len(probes) == 52

    def test_get_probes_by_category_importable(self):
        from mcpaudit import get_probes_by_category as gpbc
        assert callable(gpbc)
