"""Tests for Attack Path Analyzer."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from mcpaudit.attack_path import (
    AttackEdge,
    AttackGraph,
    AttackNode,
    AttackPath,
    AttackPattern,
    AttackRisk,
    _analyze_config,
    _build_edges_from_paths,
    _label_matches_type,
    _match_pattern_to_nodes,
    _scan_source_for_nodes,
    analyze_attack_paths,
    format_attack_dot,
    format_attack_json,
    format_attack_text,
    get_attack_patterns,
)


# --- Data Models ---

class TestAttackRisk:
    def test_critical(self):
        assert AttackRisk.CRITICAL == "critical"

    def test_high(self):
        assert AttackRisk.HIGH == "high"

    def test_medium(self):
        assert AttackRisk.MEDIUM == "medium"

    def test_low(self):
        assert AttackRisk.LOW == "low"


class TestAttackNode:
    def test_to_dict_minimal(self):
        node = AttackNode("N001", "entry_point", "User Input")
        d = node.to_dict()
        assert d["node_id"] == "N001"
        assert d["node_type"] == "entry_point"
        assert d["label"] == "User Input"
        assert "description" not in d
        assert "file" not in d

    def test_to_dict_full(self):
        node = AttackNode("N001", "entry_point", "User Input",
                          description="test", file="server.py", line=10)
        d = node.to_dict()
        assert d["description"] == "test"
        assert d["file"] == "server.py"
        assert d["line"] == 10


class TestAttackEdge:
    def test_to_dict(self):
        edge = AttackEdge("N001", "N002", "flow", "ATCK-001")
        d = edge.to_dict()
        assert d["source"] == "N001"
        assert d["target"] == "N002"


class TestAttackPath:
    def test_length_no_hops(self):
        p = AttackPath("P1", "src", [], "tgt", AttackRisk.HIGH, "desc", "mit")
        assert p.length == 2

    def test_length_with_hops(self):
        p = AttackPath("P1", "src", ["h1", "h2"], "tgt", AttackRisk.HIGH, "desc", "mit")
        assert p.length == 4

    def test_to_dict(self):
        p = AttackPath("P1", "src", ["h1"], "tgt", AttackRisk.CRITICAL,
                       "description", "mitigation", pattern="test_pattern")
        d = p.to_dict()
        assert d["path_id"] == "P1"
        assert d["source"] == "src"
        assert d["target"] == "tgt"
        assert d["risk_level"] == "critical"
        assert d["pattern"] == "test_pattern"
        assert d["length"] == 3


class TestAttackGraph:
    def test_empty_graph(self):
        g = AttackGraph()
        assert g.total_paths == 0
        assert g.critical_paths == 0
        assert g.high_paths == 0
        assert g.max_risk == "none"
        assert g.exit_code == 0

    def test_with_critical_path(self):
        g = AttackGraph(paths=[
            AttackPath("P1", "s", [], "t", AttackRisk.CRITICAL, "", ""),
        ])
        assert g.critical_paths == 1
        assert g.max_risk == "critical"
        assert g.exit_code == 2

    def test_with_high_path(self):
        g = AttackGraph(paths=[
            AttackPath("P1", "s", [], "t", AttackRisk.HIGH, "", ""),
        ])
        assert g.high_paths == 1
        assert g.max_risk == "high"
        assert g.exit_code == 1

    def test_with_medium_path(self):
        g = AttackGraph(paths=[
            AttackPath("P1", "s", [], "t", AttackRisk.MEDIUM, "", ""),
        ])
        assert g.max_risk == "medium"
        assert g.exit_code == 0

    def test_risk_summary(self):
        g = AttackGraph(paths=[
            AttackPath("P1", "s", [], "t", AttackRisk.CRITICAL, "", ""),
            AttackPath("P2", "s", [], "t", AttackRisk.HIGH, "", ""),
            AttackPath("P3", "s", [], "t", AttackRisk.MEDIUM, "", ""),
            AttackPath("P4", "s", [], "t", AttackRisk.LOW, "", ""),
        ])
        rs = g.risk_summary
        assert rs["critical"] == 1
        assert rs["high"] == 1
        assert rs["medium"] == 1
        assert rs["low"] == 1

    def test_to_dict(self):
        g = AttackGraph(target_path="/test")
        d = g.to_dict()
        assert d["target_path"] == "/test"
        assert "nodes" in d
        assert "edges" in d
        assert "paths" in d


# --- Pattern Matching ---

class TestLabelMatchesType:
    def test_user_input_match(self):
        assert _label_matches_type("tool parameter (mcp tool input)", "user_input") is True

    def test_file_system_match(self):
        assert _label_matches_type("file system access", "file_system") is True

    def test_no_match(self):
        assert _label_matches_type("completely unrelated", "user_input") is False

    def test_command_execution(self):
        assert _label_matches_type("command execution", "command_execution") is True

    def test_secret_match(self):
        assert _label_matches_type("secret/credential exposure", "secret_leak") is True


class TestMatchPatternToNodes:
    def test_simple_match(self):
        pat = AttackPattern(
            "TEST-001", "test", "test",
            ["user_input", "file_system"],
            AttackRisk.HIGH, "test mitigation",
        )
        nodes = [
            AttackNode("N1", "entry_point", "Tool parameter (MCP tool input)"),
            AttackNode("N2", "critical_asset", "File system access"),
        ]
        paths = _match_pattern_to_nodes(pat, nodes)
        assert len(paths) >= 1
        assert paths[0].risk_level == AttackRisk.HIGH

    def test_no_match_missing_type(self):
        pat = AttackPattern(
            "TEST-002", "test", "test",
            ["user_input", "database_query"],
            AttackRisk.HIGH, "test",
        )
        nodes = [
            AttackNode("N1", "entry_point", "Tool parameter (MCP tool input)"),
            AttackNode("N2", "critical_asset", "File system access"),
        ]
        paths = _match_pattern_to_nodes(pat, nodes)
        assert len(paths) == 0

    def test_skip_same_node(self):
        pat = AttackPattern(
            "TEST-003", "test", "test",
            ["user_input", "user_input"],
            AttackRisk.LOW, "test",
        )
        nodes = [
            AttackNode("N1", "entry_point", "Tool parameter (MCP tool input)"),
        ]
        paths = _match_pattern_to_nodes(pat, nodes)
        assert len(paths) == 0


# --- Source Code Analysis ---

class TestScanSourceForNodes:
    def test_tool_decorator(self):
        code = '@server.tool\ndef read_file(path: str):\n    pass\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        entry_nodes = [n for n in nodes if n.node_type == "entry_point"]
        assert len(entry_nodes) >= 1

    def test_subprocess_detection(self):
        code = 'import subprocess\nsubprocess.run(cmd)\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        asset_nodes = [n for n in nodes if n.node_type == "critical_asset"]
        assert len(asset_nodes) >= 1

    def test_env_var_access(self):
        code = 'api_key = os.environ["API_KEY"]\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        intermediate = [n for n in nodes if n.node_type == "intermediate"]
        assert len(intermediate) >= 1

    def test_file_operations(self):
        code = 'with open("data.txt") as f:\n    data = f.read()\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        assert len(nodes) >= 1

    def test_empty_source(self):
        nodes = _scan_source_for_nodes("", "test.py")
        assert len(nodes) == 0

    def test_prompt_construction(self):
        code = 'prompt = f"Tell me about {user_input}"\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        intermediate = [n for n in nodes if n.node_type == "intermediate"]
        assert len(intermediate) >= 1

    def test_request_handler(self):
        code = '@app.route("/api")\ndef handle_request():\n    pass\n'
        nodes = _scan_source_for_nodes(code, "test.py")
        entry_nodes = [n for n in nodes if n.node_type == "entry_point"]
        assert len(entry_nodes) >= 1


# --- Config Analysis ---

class TestAnalyzeConfig:
    def test_server_entry_point(self):
        config = {
            "mcpServers": {
                "my-server": {"command": "python", "args": ["server.py"]},
            }
        }
        nodes = _analyze_config(config)
        entry = [n for n in nodes if n.node_type == "entry_point"]
        assert len(entry) == 1
        assert "my-server" in entry[0].label

    def test_tools_as_intermediates(self):
        config = {
            "mcpServers": {
                "s1": {"command": "python", "tools": ["read_file", "exec_cmd"]},
            }
        }
        nodes = _analyze_config(config)
        intermediate = [n for n in nodes if n.node_type == "intermediate"]
        assert len(intermediate) == 2

    def test_secrets_as_critical_assets(self):
        config = {
            "mcpServers": {
                "s1": {"command": "python", "env": {"API_KEY": "secret123"}},
            }
        }
        nodes = _analyze_config(config)
        assets = [n for n in nodes if n.node_type == "critical_asset"]
        assert len(assets) >= 1

    def test_permissions_as_critical_assets(self):
        config = {
            "mcpServers": {
                "s1": {
                    "command": "python",
                    "permissions": {"read": True},
                },
            }
        }
        nodes = _analyze_config(config)
        assets = [n for n in nodes if n.node_type == "critical_asset"]
        assert len(assets) >= 1

    def test_empty_config(self):
        nodes = _analyze_config({})
        assert len(nodes) == 0

    def test_non_dict_server(self):
        config = {"mcpServers": {"s1": "invalid"}}
        nodes = _analyze_config(config)
        assert len(nodes) == 0


# --- Full Analysis ---

class TestAnalyzeAttackPaths:
    def test_source_file_with_tool_and_subprocess(self, tmp_path):
        f = tmp_path / "server.py"
        f.write_text(
            '@server.tool\n'
            'def handle_request(query):\n'
            '    result = subprocess.run(query, shell=True)\n'
            '    return result.stdout\n'
        )
        graph = analyze_attack_paths(f)
        assert graph.total_paths >= 1
        assert len(graph.nodes) >= 2

    def test_json_config_analysis(self, tmp_path):
        config = {
            "mcpServers": {
                "s1": {
                    "command": "python",
                    "tools": ["read_file"],
                    "env": {"SECRET_KEY": "abc"},
                    "permissions": {"read": True},
                },
            }
        }
        f = tmp_path / "config.json"
        f.write_text(json.dumps(config))
        graph = analyze_attack_paths(f)
        assert len(graph.nodes) >= 3

    def test_directory_scan(self, tmp_path):
        f1 = tmp_path / "server.py"
        f1.write_text(
            '@server.tool\ndef read_file(path):\n'
            '    return open(path).read()\n'
        )
        f2 = tmp_path / "config.json"
        f2.write_text(json.dumps({"mcpServers": {"s1": {"command": "python"}}}))
        graph = analyze_attack_paths(tmp_path)
        assert len(graph.nodes) >= 2

    def test_nonexistent_path(self):
        graph = analyze_attack_paths(Path("/nonexistent"))
        assert graph.total_paths == 0

    def test_empty_directory(self, tmp_path):
        graph = analyze_attack_paths(tmp_path)
        assert graph.total_paths == 0

    def test_invalid_json_file(self, tmp_path):
        f = tmp_path / "bad.json"
        f.write_text("not json {{{")
        graph = analyze_attack_paths(f)
        assert graph.total_paths == 0

    def test_skips_node_modules(self, tmp_path):
        nm = tmp_path / "node_modules" / "pkg"
        nm.mkdir(parents=True)
        f = nm / "code.js"
        f.write_text('var x = require("fs");\n')
        graph = analyze_attack_paths(tmp_path)
        assert len(graph.nodes) == 0


# --- Edge Building ---

class TestBuildEdgesFromPaths:
    def test_builds_edges(self):
        n1 = AttackNode("N1", "entry_point", "Source")
        n2 = AttackNode("N2", "critical_asset", "Target")
        graph = AttackGraph(
            nodes=[n1, n2],
            paths=[
                AttackPath("P1", "Source", [], "Target", AttackRisk.HIGH,
                           "desc", "mit", nodes=[n1, n2]),
            ],
        )
        _build_edges_from_paths(graph)
        assert len(graph.edges) == 1
        assert graph.edges[0].source == "N1"
        assert graph.edges[0].target == "N2"

    def test_no_duplicate_edges(self):
        n1 = AttackNode("N1", "entry_point", "Source")
        n2 = AttackNode("N2", "critical_asset", "Target")
        graph = AttackGraph(
            nodes=[n1, n2],
            paths=[
                AttackPath("P1", "Source", [], "Target", AttackRisk.HIGH,
                           "d", "m", nodes=[n1, n2]),
                AttackPath("P2", "Source", [], "Target", AttackRisk.HIGH,
                           "d", "m", nodes=[n1, n2]),
            ],
        )
        _build_edges_from_paths(graph)
        assert len(graph.edges) == 1

    def test_empty_path_nodes(self):
        graph = AttackGraph(
            paths=[AttackPath("P1", "s", [], "t", AttackRisk.LOW, "d", "m")],
        )
        _build_edges_from_paths(graph)
        assert len(graph.edges) == 0


# --- Formatters ---

class TestFormatAttackText:
    def test_empty_graph(self):
        g = AttackGraph(target_path="/test")
        text = format_attack_text(g)
        assert "Attack Path Analysis" in text
        assert "No attack paths detected" in text

    def test_with_paths(self):
        g = AttackGraph(
            target_path="/test",
            paths=[
                AttackPath("P1", "User Input", ["Tool Call"], "File System",
                           AttackRisk.CRITICAL, "desc", "Use allowlists for paths",
                           pattern="user_input->tool_call->file_system"),
            ],
        )
        text = format_attack_text(g)
        assert "[!!]" in text
        assert "User Input" in text
        assert "File System" in text
        assert "critical" in text

    def test_shows_risk_breakdown(self):
        g = AttackGraph(
            target_path="/test",
            paths=[
                AttackPath("P1", "s", [], "t", AttackRisk.CRITICAL, "", "mit"),
                AttackPath("P2", "s", [], "t", AttackRisk.HIGH, "", "mit"),
            ],
        )
        text = format_attack_text(g)
        assert "critical=1" in text
        assert "high=1" in text


class TestFormatAttackJson:
    def test_valid_json(self):
        g = AttackGraph(
            target_path="/test",
            nodes=[AttackNode("N1", "entry_point", "Input")],
            paths=[
                AttackPath("P1", "Input", [], "Target", AttackRisk.HIGH,
                           "desc", "mitigation", pattern="test"),
            ],
        )
        text = format_attack_json(g)
        data = json.loads(text)
        assert data["total_paths"] == 1
        assert data["max_risk"] == "high"
        assert len(data["paths"]) == 1


class TestFormatAttackDot:
    def test_basic_dot(self):
        g = AttackGraph(
            nodes=[
                AttackNode("N1", "entry_point", "User Input"),
                AttackNode("N2", "intermediate", "Tool Call"),
                AttackNode("N3", "critical_asset", "File System"),
            ],
            edges=[
                AttackEdge("N1", "N2", "flow"),
                AttackEdge("N2", "N3", "access"),
            ],
        )
        dot = format_attack_dot(g)
        assert "digraph" in dot
        assert "N1" in dot
        assert "N2" in dot
        assert "->" in dot

    def test_empty_graph_dot(self):
        g = AttackGraph()
        dot = format_attack_dot(g)
        assert "digraph" in dot

    def test_node_colors_in_dot(self):
        g = AttackGraph(
            nodes=[
                AttackNode("N1", "entry_point", "Entry"),
                AttackNode("N2", "intermediate", "Mid"),
                AttackNode("N3", "critical_asset", "Asset"),
            ],
        )
        dot = format_attack_dot(g)
        assert "#90EE90" in dot
        assert "#FFD700" in dot
        assert "#FF6347" in dot


# --- Public API ---

class TestGetAttackPatterns:
    def test_returns_patterns(self):
        patterns = get_attack_patterns()
        assert len(patterns) == 10

    def test_pattern_fields(self):
        patterns = get_attack_patterns()
        for p in patterns:
            assert p.pattern_id.startswith("ATCK-")
            assert p.name
            assert p.description
            assert len(p.chain) >= 2
            assert p.risk_level in ("critical", "high", "medium", "low")
            assert p.mitigation

    def test_has_critical_patterns(self):
        patterns = get_attack_patterns()
        critical = [p for p in patterns if p.risk_level == AttackRisk.CRITICAL]
        assert len(critical) >= 3
