"""Tests for MCP Config Hardener."""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path

import pytest

from mcpaudit.hardener import (
    HardeningChange,
    HardeningResult,
    HardeningRule,
    _apply_dangerous_flags_removal,
    _apply_disabled_removal,
    _apply_duplicate_removal,
    _apply_env_sanitization,
    _apply_logging_config,
    _apply_name_normalization,
    _apply_network_restriction,
    _apply_null_cleanup,
    _apply_path_restriction,
    _apply_resource_limits,
    _apply_retry_config,
    _apply_readonly_enforcement,
    _apply_shell_removal,
    _apply_timeouts,
    _apply_tls_enforcement,
    _apply_unused_env_cleanup,
    _apply_version_pinning,
    _get_all_hardening_rules,
    _has_version_pin,
    _is_dangerous_arg,
    _is_sensitive_env,
    _normalize_server_name,
    format_diff,
    format_hardened_json,
    format_hardening_report_json,
    format_hardening_text,
    get_hardening_rules,
    harden_config,
    harden_file,
)


# --- Data Models ---

class TestHardeningRule:
    def test_creation(self):
        rule = HardeningRule(
            rule_id="HARDEN-001",
            name="Test rule",
            description="A test",
            category="permissions",
            severity="high",
        )
        assert rule.rule_id == "HARDEN-001"
        assert rule.severity == "high"

    def test_to_dict(self):
        rule = HardeningRule(
            rule_id="HARDEN-002",
            name="Version pin",
            description="Pin versions",
            category="version",
        )
        d = rule.to_dict()
        assert d["rule_id"] == "HARDEN-002"
        assert d["category"] == "version"

    def test_default_severity(self):
        rule = HardeningRule(
            rule_id="HARDEN-X",
            name="X",
            description="X",
            category="test",
        )
        assert rule.severity == "medium"


class TestHardeningChange:
    def test_creation(self):
        change = HardeningChange(
            rule_id="HARDEN-001",
            path="mcpServers.test.command",
            description="Replaced shell",
            old_value="bash -c node",
            new_value="node",
        )
        assert change.rule_id == "HARDEN-001"
        assert change.old_value == "bash -c node"

    def test_to_dict(self):
        change = HardeningChange(
            rule_id="HARDEN-003",
            path="mcpServers.test.timeout",
            description="Added timeout",
        )
        d = change.to_dict()
        assert d["rule_id"] == "HARDEN-003"
        assert "old_value" not in d

    def test_to_dict_with_values(self):
        change = HardeningChange(
            rule_id="HARDEN-004",
            path="mcpServers.test.env.API_KEY",
            description="Sanitized secret",
            old_value="<redacted>",
            new_value="${API_KEY}",
        )
        d = change.to_dict()
        assert d["old_value"] == "<redacted>"
        assert d["new_value"] == "${API_KEY}"


class TestHardeningResult:
    def test_empty_result(self):
        result = HardeningResult()
        assert result.total_changes == 0
        assert result.exit_code == 0

    def test_with_changes(self):
        result = HardeningResult(
            changes=[
                HardeningChange("HARDEN-003", "path", "desc"),
                HardeningChange("HARDEN-006", "path2", "desc2"),
            ],
            rules_applied=["HARDEN-003", "HARDEN-006"],
        )
        assert result.total_changes == 2
        assert result.exit_code == 1

    def test_to_dict(self):
        result = HardeningResult(
            original_path="/test/config.json",
            rules_applied=["HARDEN-001"],
            timestamp="2024-01-01T00:00:00Z",
        )
        d = result.to_dict()
        assert d["original_path"] == "/test/config.json"
        assert d["timestamp"] == "2024-01-01T00:00:00Z"


# --- Helper Functions ---

class TestHelpers:
    def test_is_sensitive_env_api_key(self):
        assert _is_sensitive_env("API_KEY") is True

    def test_is_sensitive_env_secret(self):
        assert _is_sensitive_env("MY_SECRET") is True

    def test_is_sensitive_env_token(self):
        assert _is_sensitive_env("GITHUB_TOKEN") is True

    def test_is_sensitive_env_password(self):
        assert _is_sensitive_env("DB_PASSWORD") is True

    def test_is_sensitive_env_safe(self):
        assert _is_sensitive_env("NODE_ENV") is False

    def test_is_sensitive_env_path(self):
        assert _is_sensitive_env("PATH") is False

    def test_is_dangerous_arg_allow_all(self):
        assert _is_dangerous_arg("--allow-all") is True

    def test_is_dangerous_arg_no_sandbox(self):
        assert _is_dangerous_arg("--no-sandbox") is True

    def test_is_dangerous_arg_insecure(self):
        assert _is_dangerous_arg("--insecure") is True

    def test_is_dangerous_arg_safe(self):
        assert _is_dangerous_arg("--output") is False

    def test_has_version_pin_at(self):
        assert _has_version_pin("@modelcontextprotocol/server@1.0.0") is True

    def test_has_version_pin_equals(self):
        assert _has_version_pin("package==1.2.3") is True

    def test_has_version_pin_none(self):
        assert _has_version_pin("@modelcontextprotocol/server") is False

    def test_normalize_server_name_spaces(self):
        assert _normalize_server_name("My Server") == "my-server"

    def test_normalize_server_name_underscores(self):
        assert _normalize_server_name("my_server_name") == "my-server-name"

    def test_normalize_server_name_already_ok(self):
        assert _normalize_server_name("my-server") == "my-server"


# --- Rules Registry ---

class TestRulesRegistry:
    def test_all_rules_count(self):
        rules = _get_all_hardening_rules()
        assert len(rules) >= 15

    def test_public_api(self):
        rules = get_hardening_rules()
        assert len(rules) >= 15

    def test_rules_have_ids(self):
        rules = _get_all_hardening_rules()
        for rule in rules:
            assert rule.rule_id.startswith("HARDEN-")

    def test_rules_have_categories(self):
        rules = _get_all_hardening_rules()
        categories = {r.category for r in rules}
        assert "permissions" in categories
        assert "timeout" in categories
        assert "version" in categories

    def test_unique_rule_ids(self):
        rules = _get_all_hardening_rules()
        ids = [r.rule_id for r in rules]
        assert len(ids) == len(set(ids))


# --- Individual Hardening Actions ---

class TestShellRemoval:
    def test_bash_c_removal(self):
        changes: list[HardeningChange] = []
        config = {"command": "bash", "args": ["-c", "node server.js"]}
        result = _apply_shell_removal("test", config, changes)
        assert result["command"] == "node"
        assert result["args"] == ["server.js"]
        assert len(changes) == 1

    def test_sh_c_removal(self):
        changes: list[HardeningChange] = []
        config = {"command": "sh", "args": ["-c", "python3 app.py"]}
        result = _apply_shell_removal("test", config, changes)
        assert result["command"] == "python3"
        assert len(changes) == 1

    def test_no_shell_no_change(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "args": ["server.js"]}
        result = _apply_shell_removal("test", config, changes)
        assert result["command"] == "node"
        assert len(changes) == 0


class TestVersionPinning:
    def test_pins_unpinned_package(self):
        changes: list[HardeningChange] = []
        config = {"args": ["@modelcontextprotocol/server"]}
        result = _apply_version_pinning("test", config, changes)
        assert result["args"] == ["@modelcontextprotocol/server@latest"]
        assert len(changes) == 1

    def test_keeps_pinned_package(self):
        changes: list[HardeningChange] = []
        config = {"args": ["@modelcontextprotocol/server@1.0.0"]}
        result = _apply_version_pinning("test", config, changes)
        assert result["args"] == ["@modelcontextprotocol/server@1.0.0"]
        assert len(changes) == 0

    def test_no_packages(self):
        changes: list[HardeningChange] = []
        config = {"args": ["--port", "3000"]}
        result = _apply_version_pinning("test", config, changes)
        assert len(changes) == 0


class TestTimeouts:
    def test_adds_timeout(self):
        changes: list[HardeningChange] = []
        config = {"command": "node"}
        result = _apply_timeouts("test", config, changes)
        assert result["timeout"] == 30
        assert len(changes) == 1

    def test_keeps_existing_timeout(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "timeout": 60}
        result = _apply_timeouts("test", config, changes)
        assert result["timeout"] == 60
        assert len(changes) == 0


class TestEnvSanitization:
    def test_sanitizes_api_key(self):
        changes: list[HardeningChange] = []
        config = {"env": {"API_KEY": "sk-1234567890"}}
        result = _apply_env_sanitization("test", config, changes)
        assert result["env"]["API_KEY"] == "${API_KEY}"
        assert len(changes) == 1

    def test_keeps_reference(self):
        changes: list[HardeningChange] = []
        config = {"env": {"API_KEY": "${API_KEY}"}}
        result = _apply_env_sanitization("test", config, changes)
        assert result["env"]["API_KEY"] == "${API_KEY}"
        assert len(changes) == 0

    def test_keeps_safe_env(self):
        changes: list[HardeningChange] = []
        config = {"env": {"NODE_ENV": "production"}}
        result = _apply_env_sanitization("test", config, changes)
        assert result["env"]["NODE_ENV"] == "production"
        assert len(changes) == 0


class TestDangerousFlagsRemoval:
    def test_removes_allow_all(self):
        changes: list[HardeningChange] = []
        config = {"args": ["--port", "3000", "--allow-all"]}
        result = _apply_dangerous_flags_removal("test", config, changes)
        assert "--allow-all" not in result["args"]
        assert "--port" in result["args"]
        assert len(changes) == 1

    def test_removes_multiple(self):
        changes: list[HardeningChange] = []
        config = {"args": ["--no-sandbox", "--insecure", "--port", "80"]}
        result = _apply_dangerous_flags_removal("test", config, changes)
        assert len(result["args"]) == 2
        assert len(changes) == 2

    def test_no_dangerous_flags(self):
        changes: list[HardeningChange] = []
        config = {"args": ["--port", "3000"]}
        result = _apply_dangerous_flags_removal("test", config, changes)
        assert len(changes) == 0


class TestResourceLimits:
    def test_adds_limits(self):
        changes: list[HardeningChange] = []
        config = {"command": "node"}
        result = _apply_resource_limits("test", config, changes)
        assert "resourceLimits" in result
        assert len(changes) == 1

    def test_keeps_existing(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "resourceLimits": {"maxTokens": 50000}}
        result = _apply_resource_limits("test", config, changes)
        assert result["resourceLimits"]["maxTokens"] == 50000
        assert len(changes) == 0


class TestPathRestriction:
    def test_restricts_wildcard(self):
        changes: list[HardeningChange] = []
        config = {"args": ["*"]}
        result = _apply_path_restriction("test", config, changes)
        assert result["args"] == ["./"]
        assert len(changes) == 1

    def test_restricts_glob(self):
        changes: list[HardeningChange] = []
        config = {"args": ["/**"]}
        result = _apply_path_restriction("test", config, changes)
        assert result["args"] == ["./"]

    def test_no_wildcards(self):
        changes: list[HardeningChange] = []
        config = {"args": ["./src"]}
        result = _apply_path_restriction("test", config, changes)
        assert len(changes) == 0


class TestTlsEnforcement:
    def test_enables_tls_verify(self):
        changes: list[HardeningChange] = []
        config = {"tlsVerify": False}
        result = _apply_tls_enforcement("test", config, changes)
        assert result["tlsVerify"] is True
        assert len(changes) == 1

    def test_fixes_node_tls_env(self):
        changes: list[HardeningChange] = []
        config = {"env": {"NODE_TLS_REJECT_UNAUTHORIZED": "0"}}
        result = _apply_tls_enforcement("test", config, changes)
        assert result["env"]["NODE_TLS_REJECT_UNAUTHORIZED"] == "1"
        assert len(changes) == 1

    def test_already_enabled(self):
        changes: list[HardeningChange] = []
        config = {"tlsVerify": True}
        result = _apply_tls_enforcement("test", config, changes)
        assert len(changes) == 0


class TestDuplicateRemoval:
    def test_removes_duplicate(self):
        changes: list[HardeningChange] = []
        servers = {
            "server-a": {"command": "node", "args": ["server.js"]},
            "server-b": {"command": "node", "args": ["server.js"]},
        }
        result = _apply_duplicate_removal(servers, changes)
        assert len(result) == 1
        assert len(changes) == 1

    def test_keeps_unique(self):
        changes: list[HardeningChange] = []
        servers = {
            "server-a": {"command": "node", "args": ["a.js"]},
            "server-b": {"command": "python3", "args": ["b.py"]},
        }
        result = _apply_duplicate_removal(servers, changes)
        assert len(result) == 2
        assert len(changes) == 0


class TestNameNormalization:
    def test_normalizes_underscores(self):
        changes: list[HardeningChange] = []
        servers = {"My_Server": {"command": "node"}}
        result = _apply_name_normalization(servers, changes)
        assert "my-server" in result
        assert len(changes) == 1

    def test_keeps_correct_names(self):
        changes: list[HardeningChange] = []
        servers = {"my-server": {"command": "node"}}
        result = _apply_name_normalization(servers, changes)
        assert "my-server" in result
        assert len(changes) == 0


class TestNullCleanup:
    def test_removes_null(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "description": None, "tags": []}
        result = _apply_null_cleanup("test", config, changes)
        assert "description" not in result
        assert "tags" not in result
        assert len(changes) == 2

    def test_keeps_values(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "timeout": 30}
        result = _apply_null_cleanup("test", config, changes)
        assert len(changes) == 0


class TestRetryConfig:
    def test_adds_retries(self):
        changes: list[HardeningChange] = []
        config = {"command": "node"}
        result = _apply_retry_config("test", config, changes)
        assert result["maxRetries"] == 3
        assert len(changes) == 1

    def test_keeps_existing(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "maxRetries": 5}
        result = _apply_retry_config("test", config, changes)
        assert result["maxRetries"] == 5
        assert len(changes) == 0


class TestNetworkRestriction:
    def test_adds_policy(self):
        changes: list[HardeningChange] = []
        config = {"command": "node"}
        result = _apply_network_restriction("test", config, changes)
        assert "networkPolicy" in result
        assert "localhost" in result["networkPolicy"]["allowedHosts"]
        assert len(changes) == 1

    def test_keeps_existing(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "networkPolicy": {"allowedHosts": ["api.example.com"]}}
        result = _apply_network_restriction("test", config, changes)
        assert len(changes) == 0


class TestDisabledRemoval:
    def test_removes_disabled(self):
        changes: list[HardeningChange] = []
        servers = {
            "active": {"command": "node"},
            "inactive": {"command": "python3", "disabled": True},
        }
        result = _apply_disabled_removal(servers, changes)
        assert "active" in result
        assert "inactive" not in result
        assert len(changes) == 1

    def test_removes_not_enabled(self):
        changes: list[HardeningChange] = []
        servers = {
            "active": {"command": "node"},
            "off": {"command": "python3", "enabled": False},
        }
        result = _apply_disabled_removal(servers, changes)
        assert "off" not in result


class TestLoggingConfig:
    def test_adds_logging(self):
        changes: list[HardeningChange] = []
        config = {"command": "node"}
        result = _apply_logging_config("test", config, changes)
        assert result["logging"]["enabled"] is True
        assert len(changes) == 1

    def test_keeps_existing(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "logging": {"enabled": True, "level": "debug"}}
        result = _apply_logging_config("test", config, changes)
        assert len(changes) == 0


class TestReadonlyEnforcement:
    def test_sets_readonly(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "args": []}
        result = _apply_readonly_enforcement("test", config, changes)
        assert result["readOnly"] is True
        assert len(changes) == 1

    def test_skips_with_write_flag(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "args": ["--write", "output.json"]}
        result = _apply_readonly_enforcement("test", config, changes)
        assert "readOnly" not in result
        assert len(changes) == 0


class TestUnusedEnvCleanup:
    def test_removes_unused(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "args": ["app.js"], "env": {"UNUSED_VAR": "value", "NODE_ENV": "prod"}}
        result = _apply_unused_env_cleanup("test", config, changes)
        assert "UNUSED_VAR" not in result["env"]
        assert len(changes) == 1

    def test_keeps_sensitive(self):
        changes: list[HardeningChange] = []
        config = {"command": "node", "args": [], "env": {"API_KEY": "123"}}
        result = _apply_unused_env_cleanup("test", config, changes)
        assert "API_KEY" in result["env"]


# --- Main Hardening Engine ---

class TestHardenConfig:
    def test_basic_hardening(self):
        config = {
            "mcpServers": {
                "test-server": {
                    "command": "node",
                    "args": ["server.js"],
                },
            },
        }
        result = harden_config(config)
        assert result.total_changes > 0
        assert result.hardened_config["mcpServers"]["test-server"]["timeout"] == 30

    def test_shell_removal(self):
        config = {
            "mcpServers": {
                "my-server": {
                    "command": "bash",
                    "args": ["-c", "node server.js"],
                },
            },
        }
        result = harden_config(config, rules=["HARDEN-001"])
        assert result.hardened_config["mcpServers"]["my-server"]["command"] == "node"

    def test_with_exclude_rules(self):
        config = {
            "mcpServers": {
                "test": {"command": "node"},
            },
        }
        result = harden_config(config, exclude_rules=["HARDEN-003"])
        assert not any(c.rule_id == "HARDEN-003" for c in result.changes)

    def test_empty_config(self):
        config = {"mcpServers": {}}
        result = harden_config(config)
        assert result.total_changes == 0

    def test_preserves_original(self):
        config = {
            "mcpServers": {
                "test": {"command": "node", "args": ["server.js"]},
            },
        }
        import copy
        original = copy.deepcopy(config)
        harden_config(config)
        assert config == original

    def test_multiple_servers(self):
        config = {
            "mcpServers": {
                "server-a": {"command": "node", "args": ["a.js"]},
                "server-b": {"command": "python3", "args": ["b.py"]},
            },
        }
        result = harden_config(config)
        assert "server-a" in result.hardened_config["mcpServers"]
        assert "server-b" in result.hardened_config["mcpServers"]

    def test_specific_rules_only(self):
        config = {
            "mcpServers": {
                "test": {"command": "node"},
            },
        }
        result = harden_config(config, rules=["HARDEN-003"])
        for change in result.changes:
            assert change.rule_id == "HARDEN-003"

    def test_timestamp_set(self):
        config = {"mcpServers": {"test": {"command": "node"}}}
        result = harden_config(config)
        assert result.timestamp != ""


# --- File Operations ---

class TestHardenFile:
    def test_harden_file(self, tmp_path):
        config = {
            "mcpServers": {
                "test": {
                    "command": "bash",
                    "args": ["-c", "node server.js"],
                },
            },
        }
        config_file = tmp_path / "config.json"
        config_file.write_text(json.dumps(config), encoding="utf-8")

        result = harden_file(config_file)
        assert result.original_path == str(config_file)
        assert result.total_changes > 0

    def test_harden_file_with_rules(self, tmp_path):
        config = {"mcpServers": {"test": {"command": "node"}}}
        config_file = tmp_path / "config.json"
        config_file.write_text(json.dumps(config), encoding="utf-8")

        result = harden_file(config_file, rules=["HARDEN-003"])
        assert all(c.rule_id == "HARDEN-003" for c in result.changes)


# --- Formatters ---

class TestFormatters:
    def test_format_diff_empty(self):
        result = HardeningResult()
        output = format_diff(result)
        assert "No changes needed" in output

    def test_format_diff_with_changes(self):
        result = HardeningResult(
            original_path="/test/config.json",
            changes=[
                HardeningChange(
                    rule_id="HARDEN-003",
                    path="mcpServers.test.timeout",
                    description="Added timeout",
                    new_value=30,
                ),
            ],
        )
        output = format_diff(result)
        assert "HARDEN-003" in output
        assert "Added timeout" in output

    def test_format_hardened_json(self):
        result = HardeningResult(
            hardened_config={"mcpServers": {"test": {"command": "node"}}},
        )
        output = format_hardened_json(result)
        data = json.loads(output)
        assert "mcpServers" in data

    def test_format_report_json(self):
        result = HardeningResult(
            original_path="/test/config.json",
            rules_applied=["HARDEN-003"],
            timestamp="2024-01-01T00:00:00Z",
        )
        output = format_hardening_report_json(result)
        data = json.loads(output)
        assert data["original_path"] == "/test/config.json"

    def test_format_text_empty(self):
        result = HardeningResult(timestamp="2024-01-01T00:00:00Z")
        output = format_hardening_text(result)
        assert "already hardened" in output

    def test_format_text_with_changes(self):
        result = HardeningResult(
            timestamp="2024-01-01T00:00:00Z",
            changes=[
                HardeningChange(
                    rule_id="HARDEN-001",
                    path="mcpServers.test.command",
                    description="Replaced shell",
                ),
            ],
            rules_applied=["HARDEN-001"],
        )
        output = format_hardening_text(result)
        assert "HARDEN-001" in output
        assert "Replaced shell" in output

    def test_format_diff_original_path(self):
        result = HardeningResult(
            original_path="/etc/config.json",
            changes=[
                HardeningChange("HARDEN-003", "path", "desc", new_value=30),
            ],
        )
        output = format_diff(result)
        assert "/etc/config.json" in output


# --- Integration / Full Pipeline ---

class TestIntegration:
    def test_full_hardening_pipeline(self, tmp_path):
        config = {
            "mcpServers": {
                "My Server": {
                    "command": "bash",
                    "args": ["-c", "node server.js --allow-all"],
                    "env": {
                        "API_KEY": "sk-12345",
                        "NODE_ENV": "production",
                    },
                    "tlsVerify": False,
                },
                "Duplicate": {
                    "command": "bash",
                    "args": ["-c", "node server.js --allow-all"],
                },
                "disabled-server": {
                    "command": "python3",
                    "disabled": True,
                },
            },
        }
        config_file = tmp_path / "config.json"
        config_file.write_text(json.dumps(config), encoding="utf-8")

        result = harden_file(config_file)
        assert result.total_changes > 0
        assert len(result.rules_applied) > 0

        hardened = result.hardened_config
        servers = hardened["mcpServers"]

        # Disabled should be removed
        assert "disabled-server" not in servers

    def test_roundtrip_json(self, tmp_path):
        config = {
            "mcpServers": {
                "test": {"command": "node", "args": ["server.js"]},
            },
        }
        result = harden_config(config)

        output_file = tmp_path / "hardened.json"
        output_file.write_text(format_hardened_json(result), encoding="utf-8")

        reloaded = json.loads(output_file.read_text(encoding="utf-8"))
        assert "mcpServers" in reloaded
        assert "test" in reloaded["mcpServers"]
