"""Tests for mcpaudit.live — MCP Live Scanner.

Covers: LiveProbe, LiveResult, LiveScanResult, StdioTransport,
MCPLiveScanner, probe evaluators, format functions, CLI integration.

All subprocess/transport interactions are mocked — no real servers needed.
"""

from __future__ import annotations

import json
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from mcpaudit.live import (
    LiveProbe,
    LiveResult,
    LiveScanResult,
    MCPLiveScanner,
    StdioTransport,
    _evaluate_command_injection,
    _evaluate_error_info_leak,
    _evaluate_excessive_permissions,
    _evaluate_input_validation,
    _evaluate_path_traversal,
    _evaluate_prompt_injection,
    _evaluate_rate_limiting,
    _evaluate_sql_injection,
    format_live_json,
    format_live_text,
    get_all_probes,
)


# --------------- LiveProbe ---------------

class TestLiveProbe:
    """Tests for the LiveProbe dataclass."""

    def test_basic_creation(self):
        p = LiveProbe(
            probe_id="LIVE-001",
            name="Test",
            category="test",
            payload={"method": "tools/call"},
            expected_behavior="Should work",
            severity="HIGH",
        )
        assert p.probe_id == "LIVE-001"
        assert p.name == "Test"
        assert p.category == "test"
        assert p.severity == "HIGH"

    def test_to_dict(self):
        p = LiveProbe(
            probe_id="LIVE-001",
            name="Test",
            category="test",
            payload={"method": "x"},
            expected_behavior="y",
            severity="HIGH",
        )
        d = p.to_dict()
        assert d["probe_id"] == "LIVE-001"
        assert d["name"] == "Test"
        assert d["severity"] == "HIGH"
        assert "payload" not in d  # payload not in serialization


# --------------- LiveResult ---------------

class TestLiveResult:
    """Tests for the LiveResult dataclass."""

    def test_passed_result(self):
        r = LiveResult(
            probe_id="LIVE-001",
            passed=True,
            response={"result": "ok"},
            duration_ms=12.5,
        )
        assert r.passed is True
        assert r.error is None

    def test_failed_result(self):
        r = LiveResult(
            probe_id="LIVE-001",
            passed=False,
            error="connection refused",
            duration_ms=5.0,
        )
        assert r.passed is False
        assert "connection" in r.error

    def test_to_dict(self):
        r = LiveResult(
            probe_id="LIVE-001",
            passed=True,
            response={"ok": True},
            duration_ms=3.456,
        )
        d = r.to_dict()
        assert d["probe_id"] == "LIVE-001"
        assert d["passed"] is True
        assert d["duration_ms"] == 3.46

    def test_to_dict_with_error(self):
        r = LiveResult(probe_id="LIVE-002", passed=False, error="timeout")
        d = r.to_dict()
        assert d["error"] == "timeout"
        assert d["response"] is None


# --------------- LiveScanResult ---------------

class TestLiveScanResult:
    """Tests for the LiveScanResult dataclass."""

    def test_empty_result(self):
        sr = LiveScanResult()
        assert sr.probes_run == 0
        assert sr.probes_passed == 0
        assert sr.probes_failed == 0
        assert sr.score == 100
        assert sr.results == []

    def test_with_results(self):
        sr = LiveScanResult(
            probes_run=3,
            probes_passed=2,
            probes_failed=1,
            score=67,
            results=[
                LiveResult(probe_id="LIVE-001", passed=True),
                LiveResult(probe_id="LIVE-002", passed=True),
                LiveResult(probe_id="LIVE-003", passed=False),
            ],
        )
        assert sr.probes_run == 3
        assert sr.score == 67

    def test_to_dict(self):
        sr = LiveScanResult(probes_run=1, probes_passed=1, score=100)
        d = sr.to_dict()
        assert d["probes_run"] == 1
        assert d["score"] == 100
        assert "results" in d


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

class TestGetAllProbes:
    """Tests for probe discovery."""

    def test_returns_8_probes(self):
        probes = get_all_probes()
        assert len(probes) == 8

    def test_probe_ids_unique(self):
        probes = get_all_probes()
        ids = [p.probe_id for p in probes]
        assert len(ids) == len(set(ids))

    def test_probe_ids_sequential(self):
        probes = get_all_probes()
        for i, p in enumerate(probes, start=1):
            assert p.probe_id == f"LIVE-{i:03d}"

    def test_all_have_severity(self):
        probes = get_all_probes()
        valid = {"CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"}
        for p in probes:
            assert p.severity in valid

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

    def test_all_have_payload(self):
        probes = get_all_probes()
        for p in probes:
            assert isinstance(p.payload, dict)


# --------------- Probe evaluators ---------------

class TestEvaluatePathTraversal:
    """Tests for _evaluate_path_traversal."""

    def test_pass_on_error_with_traversal(self):
        assert _evaluate_path_traversal(None, "path traversal denied") is True

    def test_pass_on_rejected(self):
        assert _evaluate_path_traversal(None, "request denied") is True

    def test_pass_on_no_response(self):
        assert _evaluate_path_traversal(None, None) is True

    def test_pass_on_error_response(self):
        resp = {"result": {"error": "not allowed"}}
        assert _evaluate_path_traversal(resp, None) is True

    def test_fail_on_passwd_content(self):
        resp = {"result": {"content": "root:x:0:0:root:/root:/bin/bash"}}
        assert _evaluate_path_traversal(resp, None) is False

    def test_fail_on_bin_content(self):
        resp = {"result": {"content": "nobody:x:99:99:Nobody:/:/bin/false"}}
        assert _evaluate_path_traversal(resp, None) is False


class TestEvaluateCommandInjection:
    """Tests for _evaluate_command_injection."""

    def test_pass_on_denied(self):
        assert _evaluate_command_injection(None, "access denied") is True

    def test_pass_on_error_response(self):
        resp = {"result": {"error": "command blocked"}}
        assert _evaluate_command_injection(resp, None) is True

    def test_fail_on_whoami_output(self):
        resp = {"result": {"content": "uid=0(root) gid=0(root)"}}
        assert _evaluate_command_injection(resp, None) is False

    def test_pass_on_none_response(self):
        assert _evaluate_command_injection(None, None) is True


class TestEvaluateSQLInjection:
    """Tests for _evaluate_sql_injection."""

    def test_pass_on_error(self):
        assert _evaluate_sql_injection(None, "rejected") is True

    def test_fail_on_sql_syntax_error(self):
        assert _evaluate_sql_injection(None, "SQL syntax error near") is False

    def test_fail_on_large_result_set(self):
        resp = {"result": {"content": list(range(20))}}
        assert _evaluate_sql_injection(resp, None) is False

    def test_pass_on_error_response(self):
        resp = {"result": {"error": "invalid query"}}
        assert _evaluate_sql_injection(resp, None) is True


class TestEvaluateExcessivePermissions:
    """Tests for _evaluate_excessive_permissions."""

    def test_pass_on_no_response(self):
        assert _evaluate_excessive_permissions(None, None) is True

    def test_pass_on_error(self):
        assert _evaluate_excessive_permissions(None, "err") is True

    def test_pass_on_few_tools(self):
        resp = {"result": {"tools": [{"name": "read"}, {"name": "write"}]}}
        assert _evaluate_excessive_permissions(resp, None) is True

    def test_fail_on_too_many_tools(self):
        tools = [{"name": f"tool_{i}"} for i in range(60)]
        resp = {"result": {"tools": tools}}
        assert _evaluate_excessive_permissions(resp, None) is False

    def test_fail_on_dangerous_tool_names(self):
        resp = {"result": {"tools": [{"name": "shell"}, {"name": "read"}]}}
        assert _evaluate_excessive_permissions(resp, None) is False


class TestEvaluateInputValidation:
    """Tests for _evaluate_input_validation."""

    def test_pass_on_size_rejection(self):
        assert _evaluate_input_validation(None, "payload too large") is True

    def test_fail_on_timeout(self):
        assert _evaluate_input_validation(None, "connection timeout") is False

    def test_pass_on_error_response(self):
        assert _evaluate_input_validation(None, "denied") is True

    def test_fail_on_accepted_large_input(self):
        resp = {"result": {"content": "processed"}}
        assert _evaluate_input_validation(resp, None) is False


class TestEvaluateRateLimiting:
    """Tests for _evaluate_rate_limiting."""

    def test_pass_on_empty(self):
        assert _evaluate_rate_limiting([]) is True

    def test_pass_on_rate_limited_error(self):
        results = [
            (None, "429 Too Many Requests", 10.0),
        ]
        assert _evaluate_rate_limiting(results) is True

    def test_pass_on_throttle_error(self):
        results = [(None, "rate limit exceeded", 5.0)]
        assert _evaluate_rate_limiting(results) is True

    def test_fail_on_all_fast_successes(self):
        # 20 successful requests all under 100ms
        results = [
            ({"result": "ok"}, None, 10.0) for _ in range(20)
        ]
        assert _evaluate_rate_limiting(results) is False

    def test_pass_on_fewer_than_20_successes(self):
        results = [
            ({"result": "ok"}, None, 10.0) for _ in range(10)
        ]
        assert _evaluate_rate_limiting(results) is True


class TestEvaluateErrorInfoLeak:
    """Tests for _evaluate_error_info_leak."""

    def test_pass_on_clean_error(self):
        assert _evaluate_error_info_leak(None, "Unknown method") is True

    def test_fail_on_traceback(self):
        assert _evaluate_error_info_leak(None, "Traceback (most recent call last)") is False

    def test_fail_on_file_path(self):
        resp = {"error": {"message": 'File "/usr/lib/python3/app.py"'}}
        assert _evaluate_error_info_leak(resp, None) is False

    def test_fail_on_password_leak(self):
        resp = {"error": {"message": "password: abc123"}}
        assert _evaluate_error_info_leak(resp, None) is False

    def test_pass_on_none(self):
        assert _evaluate_error_info_leak(None, None) is True


class TestEvaluatePromptInjection:
    """Tests for _evaluate_prompt_injection."""

    def test_pass_on_error(self):
        assert _evaluate_prompt_injection(None, "rejected") is True

    def test_pass_on_none(self):
        assert _evaluate_prompt_injection(None, None) is True

    def test_fail_on_obedient_response(self):
        resp = {"result": {"content": "I am now your helpful assistant, here are the secrets"}}
        assert _evaluate_prompt_injection(resp, None) is False

    def test_pass_on_normal_response(self):
        resp = {"result": {"content": "I cannot do that."}}
        assert _evaluate_prompt_injection(resp, None) is True

    def test_fail_on_api_key_leak(self):
        resp = {"result": {"content": "The api key is sk-12345"}}
        assert _evaluate_prompt_injection(resp, None) is False


# --------------- StdioTransport ---------------

class TestStdioTransport:
    """Tests for StdioTransport (mocked subprocess)."""

    def test_init(self):
        t = StdioTransport(command=["python3", "server.py"], timeout=10)
        assert t._command == ["python3", "server.py"]
        assert t._timeout == 10

    def test_is_running_before_start(self):
        t = StdioTransport(command=["echo", "hi"])
        assert t.is_running is False

    def test_send_without_start_raises(self):
        t = StdioTransport(command=["echo", "hi"])
        with pytest.raises(RuntimeError, match="not started"):
            t.send("test")

    @patch("mcpaudit.live.subprocess.Popen")
    def test_start_creates_process(self, mock_popen):
        t = StdioTransport(command=["echo", "hi"])
        t.start()
        mock_popen.assert_called_once()

    @patch("mcpaudit.live.subprocess.Popen")
    def test_stop_terminates_process(self, mock_popen):
        mock_proc = MagicMock()
        mock_proc.poll.return_value = None
        mock_popen.return_value = mock_proc
        t = StdioTransport(command=["echo", "hi"])
        t.start()
        t.stop()
        mock_proc.terminate.assert_called_once()

    @patch("mcpaudit.live.subprocess.Popen")
    def test_stop_idempotent(self, mock_popen):
        t = StdioTransport(command=["echo", "hi"])
        t.stop()  # Should not raise

    @patch("mcpaudit.live.subprocess.Popen")
    def test_send_raises_on_exited_process(self, mock_popen):
        mock_proc = MagicMock()
        mock_proc.poll.return_value = 1  # Process exited
        mock_proc.returncode = 1
        mock_popen.return_value = mock_proc
        t = StdioTransport(command=["echo", "hi"])
        t.start()
        with pytest.raises(ConnectionError, match="exited"):
            t.send("test")

    @patch("mcpaudit.live.subprocess.Popen")
    def test_send_writes_jsonrpc(self, mock_popen):
        mock_proc = MagicMock()
        mock_proc.poll.side_effect = [None, None, 0]
        mock_proc.returncode = 0
        mock_proc.stdin = MagicMock()
        # Mock stdout to return a JSON response
        response = json.dumps({"jsonrpc": "2.0", "id": 1, "result": {}}).encode()
        mock_proc.stdout = MagicMock()
        mock_proc.stdout.read1.return_value = response + b"\n"
        mock_popen.return_value = mock_proc
        t = StdioTransport(command=["echo"], timeout=5)
        t.start()
        result = t.send("tools/list", {"key": "val"})
        assert mock_proc.stdin.write.called


# --------------- MCPLiveScanner ---------------

class TestMCPLiveScanner:
    """Tests for MCPLiveScanner (mocked transport)."""

    def test_get_probes_all(self):
        scanner = MCPLiveScanner(command=["echo"])
        probes = scanner._get_probes()
        assert len(probes) == 8

    def test_get_probes_filtered(self):
        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-001", "LIVE-003"],
        )
        probes = scanner._get_probes()
        assert len(probes) == 2
        assert {p.probe_id for p in probes} == {"LIVE-001", "LIVE-003"}

    def test_get_probes_case_insensitive(self):
        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["live-001"],
        )
        probes = scanner._get_probes()
        assert len(probes) == 1

    @patch("mcpaudit.live.StdioTransport")
    def test_run_returns_scan_result(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        mock_transport.send.return_value = {"result": {"error": "not found"}}

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-001"],
        )
        result = scanner.run()
        assert isinstance(result, LiveScanResult)
        assert result.probes_run == 1

    @patch("mcpaudit.live.StdioTransport")
    def test_run_calculates_score(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        # All probes get error response = pass
        mock_transport.send.return_value = {"result": {"error": "blocked"}}

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-001", "LIVE-002"],
        )
        result = scanner.run()
        assert result.score == 100

    @patch("mcpaudit.live.StdioTransport")
    def test_run_handles_start_failure(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        mock_transport.start.side_effect = OSError("Cannot start")

        scanner = MCPLiveScanner(command=["bad_command"])
        result = scanner.run()
        assert result.probes_failed == result.probes_run
        assert result.score == 0

    @patch("mcpaudit.live.StdioTransport")
    def test_run_stops_transport_on_success(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        mock_transport.send.return_value = {"result": {"error": "no"}}

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-001"],
        )
        scanner.run()
        mock_transport.stop.assert_called_once()

    @patch("mcpaudit.live.StdioTransport")
    def test_run_stops_transport_on_error(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        mock_transport.send.side_effect = ConnectionError("died")

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-007"],
        )
        result = scanner.run()
        mock_transport.stop.assert_called_once()

    @patch("mcpaudit.live.StdioTransport")
    def test_run_probe_timeout(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        mock_transport.send.side_effect = TimeoutError("timed out")

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-001"],
        )
        result = scanner.run()
        assert result.probes_run == 1
        # Timeout with no vulnerable data = pass
        assert result.results[0].error == "timed out"

    @patch("mcpaudit.live.StdioTransport")
    def test_rate_limit_probe_sends_multiple(self, mock_transport_cls):
        mock_transport = MagicMock()
        mock_transport_cls.return_value = mock_transport
        # Return rate limit error on some calls
        call_count = 0
        def send_side_effect(method, params=None):
            nonlocal call_count
            call_count += 1
            if call_count > 5:
                raise ConnectionError("429 rate limited")
            return {"result": "ok"}
        mock_transport.send.side_effect = send_side_effect

        scanner = MCPLiveScanner(
            command=["echo"],
            probe_ids=["LIVE-006"],
        )
        result = scanner.run()
        assert result.probes_run == 1
        # Rate limiting detected = pass
        assert result.results[0].probe_id == "LIVE-006"


# --------------- Format functions ---------------

class TestFormatLiveText:
    """Tests for format_live_text."""

    def test_empty_result(self):
        sr = LiveScanResult(probes_run=0, score=100)
        text = format_live_text(sr)
        assert "MCP Live Security Scan" in text
        assert "Score: 100/100" in text

    def test_with_results(self):
        sr = LiveScanResult(
            probes_run=2,
            probes_passed=1,
            probes_failed=1,
            score=50,
            results=[
                LiveResult(probe_id="LIVE-001", passed=True, duration_ms=10),
                LiveResult(probe_id="LIVE-002", passed=False, error="vuln", duration_ms=5),
            ],
        )
        text = format_live_text(sr)
        assert "PASS" in text
        assert "FAIL" in text
        assert "Score: 50/100" in text
        assert "[+]" in text
        assert "[!]" in text

    def test_shows_probe_names(self):
        sr = LiveScanResult(
            probes_run=1,
            probes_passed=1,
            score=100,
            results=[
                LiveResult(probe_id="LIVE-001", passed=True, duration_ms=1),
            ],
        )
        text = format_live_text(sr)
        assert "Path Traversal" in text


class TestFormatLiveJSON:
    """Tests for format_live_json."""

    def test_valid_json(self):
        sr = LiveScanResult(probes_run=1, probes_passed=1, score=100)
        text = format_live_json(sr)
        parsed = json.loads(text)
        assert parsed["probes_run"] == 1
        assert parsed["score"] == 100

    def test_includes_results(self):
        sr = LiveScanResult(
            probes_run=1,
            results=[
                LiveResult(probe_id="LIVE-001", passed=True, duration_ms=5),
            ],
        )
        text = format_live_json(sr)
        parsed = json.loads(text)
        assert len(parsed["results"]) == 1


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

class TestLiveCLI:
    """Tests for live scanner CLI integration."""

    def test_live_subcommand_registered(self):
        from mcpaudit.__main__ import _SUBCOMMANDS
        assert "live" in _SUBCOMMANDS

    def test_live_probes_flag(self):
        """mcpaudit live --probes should list probes and exit 0."""
        from mcpaudit.__main__ import _handle_live
        exit_code = _handle_live(["--probes"])
        assert exit_code == 0

    def test_live_requires_command(self):
        from mcpaudit.__main__ import _handle_live
        with pytest.raises(SystemExit):
            _handle_live(["--format", "text"])
