"""MCP-005 through MCP-019: MCP protocol-specific security rules.

Detects security vulnerabilities specific to the Model Context Protocol,
including missing input validation, authentication gaps, transport issues,
and protocol-level misconfigurations.

Note: MCP-001 (tool description), MCP-002 (env access), MCP-003 (confused deputy),
and MCP-004 (prompt injection) are defined in other modules.
"""

from __future__ import annotations

import ast
import re

from . import Rule
from ..models import Finding, Severity


# ---------------------------------------------------------------------------
# Helper: detect if a function is an MCP tool handler
# ---------------------------------------------------------------------------

def _is_tool_handler(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
    """Check if a function looks like an MCP tool handler."""
    name_lower = node.name.lower()
    if any(kw in name_lower for kw in ("tool", "handle", "call_tool")):
        return True
    for dec in node.decorator_list:
        dec_str = ast.dump(dec)
        if "tool" in dec_str.lower():
            return True
    return False


def _is_server_file(source: str) -> bool:
    """Check if source looks like an MCP server file."""
    return "Server(" in source or "mcp" in source.lower()


def _has_pattern(source: str, patterns: list[str]) -> bool:
    """Check if any pattern appears in source (case-insensitive)."""
    source_lower = source.lower()
    return any(p in source_lower for p in patterns)


# ---------------------------------------------------------------------------
# MCP-005: Tool handler missing input validation
# ---------------------------------------------------------------------------

class MCPToolInputValidationRule(Rule):
    """Detects tool handlers that process parameters without validation."""

    @property
    def rule_id(self) -> str:
        return "MCP-005"

    @property
    def name(self) -> str:
        return "Tool Handler Missing Input Validation"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP tool handler processes input parameters without type checking, "
            "schema validation, or sanitization."
        )

    @property
    def cwe(self) -> str:
        return "CWE-20"

    def suggest_fix(self, source: str, finding: Finding) -> str | None:
        """Suggest adding isinstance check for input validation."""
        lines = source.splitlines()
        idx = finding.line - 1
        if idx < 0 or idx >= len(lines):
            return None
        line = lines[idx]
        indent = line[: len(line) - len(line.lstrip())]
        return (
            f"{indent}# TODO: Add isinstance(arguments, dict) "
            f"and validate required keys before processing"
        )

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not _is_tool_handler(node):
                continue

            body_src = ast.get_source_segment(source, node) or ""
            has_params = any(
                kw in body_src
                for kw in ("arguments", "params", "request", "args")
            )
            if not has_params:
                continue

            has_validation = any(
                kw in body_src.lower()
                for kw in (
                    "isinstance", "validate", "schema", "pydantic",
                    "typeguard", "assert ", "raise valueerror",
                    "raise typeerror",
                )
            )
            if has_validation:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Tool handler missing input validation",
                description=(
                    f"Function '{node.name}' processes input parameters "
                    f"without validation or type checking."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Add input validation using isinstance(), Pydantic models, "
                    "or JSON Schema before processing tool parameters."
                ),
                cwe=self.cwe,
                confidence="HIGH",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-006: Tool response size not limited
# ---------------------------------------------------------------------------

class MCPToolResponseSizeRule(Rule):
    """Detects tool handlers returning data without size limits."""

    @property
    def rule_id(self) -> str:
        return "MCP-006"


    @property
    def name(self) -> str:
        return "Tool Response Size Not Limited"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return (
            "MCP tool handler returns data without size limits, "
            "risking memory exhaustion or context overflow."
        )

    @property
    def cwe(self) -> str:
        return "CWE-770"

    def suggest_fix(self, source: str, finding: Finding) -> str | None:
        """Suggest adding [:MAX_SIZE] truncation to response data."""
        lines = source.splitlines()
        idx = finding.line - 1
        if idx < 0 or idx >= len(lines):
            return None
        line = lines[idx]
        indent = line[: len(line) - len(line.lstrip())]
        return (
            f"{indent}# TODO: Add response truncation — "
            f"data = data[:MAX_RESPONSE_SIZE] before returning"
        )

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not _is_tool_handler(node):
                continue

            body_src = ast.get_source_segment(source, node) or ""

            # Reads from files or databases without limiting output
            reads_data = any(
                kw in body_src
                for kw in (".read()", ".readlines()", "fetchall()", "json.load")
            )
            if not reads_data:
                continue

            has_limit = any(
                kw in body_src.lower()
                for kw in (
                    "truncate", "[:max", "limit", "max_size",
                    "max_length", "[:1024", "[:4096",
                )
            )
            if has_limit:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Tool response size not limited",
                description=(
                    f"Function '{node.name}' returns data from I/O "
                    f"operations without size limits."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Limit response size using truncation, pagination, "
                    "or max_size parameters before returning data."
                ),
                cwe=self.cwe,
                confidence="MEDIUM",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-007: Server lacks authentication check
# ---------------------------------------------------------------------------

class MCPServerAuthCheckRule(Rule):
    """Detects MCP servers without authentication on connection handlers."""

    @property
    def rule_id(self) -> str:
        return "MCP-007"


    @property
    def name(self) -> str:
        return "Server Lacks Authentication Check"

    @property
    def severity(self) -> Severity:
        return Severity.CRITICAL

    @property
    def description(self) -> str:
        return (
            "MCP server exposes connection handlers without "
            "authentication or token verification."
        )

    @property
    def cwe(self) -> str:
        return "CWE-306"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        if not _is_server_file(source):
            return findings

        # Look for connection/request handlers
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            name_lower = node.name.lower()
            is_handler = any(
                kw in name_lower
                for kw in ("handle_request", "on_connect", "handle_connection", "serve")
            )
            if not is_handler:
                continue

            body_src = ast.get_source_segment(source, node) or ""
            has_auth = any(
                kw in body_src.lower()
                for kw in (
                    "auth", "token", "api_key", "authenticate",
                    "verify", "credential", "bearer",
                )
            )
            if has_auth:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Server handler lacks authentication",
                description=(
                    f"Connection handler '{node.name}' has no "
                    f"authentication or token verification."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Add authentication middleware or token verification "
                    "before processing requests."
                ),
                cwe=self.cwe,
                confidence="HIGH",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-008: Resource listing unrestricted
# ---------------------------------------------------------------------------

class MCPResourceListUnrestrictedRule(Rule):
    """Detects resource listing endpoints without access control."""

    @property
    def rule_id(self) -> str:
        return "MCP-008"

    @property
    def name(self) -> str:
        return "Resource Listing Unrestricted"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return (
            "MCP server exposes resource listing without access control, "
            "allowing enumeration of all available resources."
        )

    @property
    def cwe(self) -> str:
        return "CWE-284"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            name_lower = node.name.lower()
            is_listing = any(
                kw in name_lower
                for kw in ("list_resource", "list_tool", "get_resources", "get_tools")
            )
            if not is_listing:
                continue

            body_src = ast.get_source_segment(source, node) or ""
            has_access_control = any(
                kw in body_src.lower()
                for kw in (
                    "permission", "authorize", "role", "acl",
                    "allowed", "check_access", "can_access",
                )
            )
            if has_access_control:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Resource listing without access control",
                description=(
                    f"Function '{node.name}' lists resources without "
                    f"checking caller permissions."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Add access control checks to resource listing "
                    "endpoints. Filter results based on caller identity."
                ),
                cwe=self.cwe,
                confidence="MEDIUM",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-009: Server exposes internal paths
# ---------------------------------------------------------------------------

class MCPInternalPathExposureRule(Rule):
    """Detects server responses that expose internal file system paths."""

    @property
    def rule_id(self) -> str:
        return "MCP-009"

    @property
    def name(self) -> str:
        return "Server Exposes Internal Paths"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return (
            "MCP server exposes internal file system paths in responses, "
            "leaking server infrastructure details."
        )

    @property
    def cwe(self) -> str:
        return "CWE-200"

    _PATH_PATTERNS = [
        re.compile(r'''['"]\/(?:home|usr|var|etc|opt|tmp|root)\/[^'"]+['"]'''),
        re.compile(r'''['"]\w:\\\\[^'"]+['"]'''),
        re.compile(r'''__file__|__path__|os\.path\.abspath|os\.getcwd'''),
    ]

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        if not _is_server_file(source):
            return findings

        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            # Skip comments and imports
            stripped = line.strip()
            if stripped.startswith("#") or stripped.startswith("import"):
                continue

            for pattern in self._PATH_PATTERNS:
                if pattern.search(line):
                    # Check if inside a return or response context
                    in_context = any(
                        kw in line.lower()
                        for kw in ("return", "response", "send", "write", "result")
                    )
                    if in_context:
                        findings.append(Finding(
                            rule_id=self.rule_id,
                            severity=self.severity,
                            title="Internal path exposed in response",
                            description=(
                                f"Line {i} exposes an internal file system "
                                f"path in a server response."
                            ),
                            file="",
                            line=i,
                            code_snippet=stripped,
                            recommendation=(
                                "Sanitize responses to remove internal paths. "
                                "Use relative paths or opaque identifiers."
                            ),
                            cwe=self.cwe,
                            confidence="MEDIUM",
                        ))
                    break
        return findings


# ---------------------------------------------------------------------------
# MCP-010: Tool name collision with system tools
# ---------------------------------------------------------------------------

class MCPToolNameCollisionRule(Rule):
    """Detects tool names that collide with common system tool names."""

    @property
    def rule_id(self) -> str:
        return "MCP-010"

    @property
    def name(self) -> str:
        return "Tool Name Collision with System Tools"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP tool uses a name that collides with common system tools, "
            "risking confused deputy attacks or shadowing."
        )

    @property
    def cwe(self) -> str:
        return "CWE-441"

    _SYSTEM_TOOL_NAMES = {
        "exec", "execute", "eval", "run", "shell", "system", "bash",
        "cmd", "command", "spawn", "fork", "kill", "sudo", "admin",
        "rm", "delete", "remove", "drop", "truncate", "format",
        "read_file", "write_file", "list_directory",
    }

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []

        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not _is_tool_handler(node):
                continue

            # Extract the tool name from decorator args or function name
            tool_name = self._extract_tool_name(node)
            if tool_name and tool_name.lower() in self._SYSTEM_TOOL_NAMES:
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Tool name collides with system tool",
                    description=(
                        f"Tool '{tool_name}' uses a name that collides with "
                        f"common system commands, risking confused deputy attacks."
                    ),
                    file="",
                    line=node.lineno,
                    code_snippet=self._get_snippet(source, node.lineno),
                    recommendation=(
                        "Use a namespaced tool name (e.g. 'myapp_exec' instead "
                        "of 'exec') to avoid shadowing system tools."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
        return findings

    def _extract_tool_name(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> str:
        """Extract tool name from decorator arguments or function name."""
        for dec in node.decorator_list:
            # @server.tool(name="exec")
            if isinstance(dec, ast.Call):
                for kw in dec.keywords:
                    if kw.arg == "name" and isinstance(kw.value, ast.Constant):
                        return str(kw.value.value)
                # @server.tool("exec")
                if dec.args and isinstance(dec.args[0], ast.Constant):
                    return str(dec.args[0].value)
        return node.name


# ---------------------------------------------------------------------------
# MCP-011: Missing transport encryption
# ---------------------------------------------------------------------------

class MCPMissingTransportEncryptionRule(Rule):
    """Detects servers using unencrypted transport (HTTP instead of HTTPS)."""

    @property
    def rule_id(self) -> str:
        return "MCP-011"

    @property
    def name(self) -> str:
        return "Missing Transport Encryption"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP server uses unencrypted HTTP transport instead of HTTPS/TLS, "
            "exposing data to interception."
        )

    @property
    def cwe(self) -> str:
        return "CWE-319"

    _HTTP_PATTERNS = [
        re.compile(r'''['"]http://(?!localhost|127\.0\.0\.1|0\.0\.0\.0)[^'"]+['"]'''),
        re.compile(r'''host\s*=\s*['"][^'"]+['"].*ssl\s*=\s*False''', re.IGNORECASE),
    ]

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        if not _is_server_file(source):
            return findings

        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            stripped = line.strip()
            if stripped.startswith("#"):
                continue
            for pattern in self._HTTP_PATTERNS:
                if pattern.search(line):
                    findings.append(Finding(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        title="Unencrypted transport detected",
                        description=(
                            f"Line {i} uses unencrypted HTTP for "
                            f"non-localhost communication."
                        ),
                        file="",
                        line=i,
                        code_snippet=stripped,
                        recommendation=(
                            "Use HTTPS/TLS for all non-localhost connections. "
                            "Configure SSL certificates for production."
                        ),
                        cwe=self.cwe,
                        confidence="HIGH",
                    ))
                    break
        return findings


# ---------------------------------------------------------------------------
# MCP-012: Server version disclosure
# ---------------------------------------------------------------------------

class MCPServerVersionDisclosureRule(Rule):
    """Detects server version strings exposed in responses or headers."""

    @property
    def rule_id(self) -> str:
        return "MCP-012"

    @property
    def name(self) -> str:
        return "Server Version Disclosure"

    @property
    def severity(self) -> Severity:
        return Severity.LOW

    @property
    def description(self) -> str:
        return (
            "MCP server discloses version information in responses, "
            "aiding attackers in targeting known vulnerabilities."
        )

    @property
    def cwe(self) -> str:
        return "CWE-200"

    _VERSION_PATTERNS = [
        re.compile(
            r'''['"](?:server[_-]?version|x[_-]?server|x[_-]?powered[_-]?by)['"]'''
            r'''\s*[:=]\s*['"][^'"]+['"]''',
            re.IGNORECASE,
        ),
        re.compile(
            r'''version\s*[:=]\s*['"][\d]+\.[\d]+''',
            re.IGNORECASE,
        ),
        re.compile(
            r'''['"]version['"]\s*:\s*['"][\d]+\.[\d]+''',
            re.IGNORECASE,
        ),
    ]

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        if not _is_server_file(source):
            return findings

        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            stripped = line.strip()
            if stripped.startswith("#"):
                continue
            for pattern in self._VERSION_PATTERNS:
                if pattern.search(line):
                    # Check if in response context
                    in_response = any(
                        kw in line.lower()
                        for kw in (
                            "response", "header", "return", "send",
                            "metadata", "info", "capabilities",
                        )
                    )
                    if in_response:
                        findings.append(Finding(
                            rule_id=self.rule_id,
                            severity=self.severity,
                            title="Server version disclosed",
                            description=(
                                f"Line {i} exposes server version information "
                                f"in a response or header."
                            ),
                            file="",
                            line=i,
                            code_snippet=stripped,
                            recommendation=(
                                "Remove or obfuscate version information from "
                                "server responses and headers."
                            ),
                            cwe=self.cwe,
                            confidence="LOW",
                        ))
                        break
        return findings


# ---------------------------------------------------------------------------
# MCP-013: Tool schema missing required fields
# ---------------------------------------------------------------------------

class MCPToolSchemaMissingFieldsRule(Rule):
    """Detects tool registrations without proper schema definitions."""

    @property
    def rule_id(self) -> str:
        return "MCP-013"

    @property
    def name(self) -> str:
        return "Tool Schema Missing Required Fields"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return (
            "MCP tool registration lacks required schema fields "
            "(description, inputSchema, or parameter types)."
        )

    @property
    def cwe(self) -> str:
        return "CWE-1059"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        source_lower = source.lower()

        # Look for tool registration patterns
        for node in ast.walk(tree):
            if not isinstance(node, ast.Dict):
                continue

            keys = []
            for key in node.keys:
                if isinstance(key, ast.Constant) and isinstance(key.value, str):
                    keys.append(key.value.lower())

            # Looks like a tool definition dict
            if "name" in keys and ("tool" in source_lower or "handler" in source_lower):
                missing = []
                if "description" not in keys:
                    missing.append("description")
                if "inputschema" not in keys and "input_schema" not in keys and "parameters" not in keys:
                    missing.append("inputSchema")

                if missing:
                    findings.append(Finding(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        title="Tool schema missing required fields",
                        description=(
                            f"Tool definition at line {node.lineno} is missing "
                            f"fields: {', '.join(missing)}."
                        ),
                        file="",
                        line=node.lineno,
                        code_snippet=self._get_snippet(source, node.lineno),
                        recommendation=(
                            "Include 'description' and 'inputSchema' in all "
                            "tool registrations for proper validation."
                        ),
                        cwe=self.cwe,
                        confidence="MEDIUM",
                    ))
        return findings


# ---------------------------------------------------------------------------
# MCP-014: Notification handler accepts arbitrary data
# ---------------------------------------------------------------------------

class MCPNotificationArbitraryDataRule(Rule):
    """Detects notification handlers that accept unvalidated data."""

    @property
    def rule_id(self) -> str:
        return "MCP-014"

    @property
    def name(self) -> str:
        return "Notification Handler Accepts Arbitrary Data"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP notification handler accepts arbitrary data without "
            "schema validation, risking injection attacks."
        )

    @property
    def cwe(self) -> str:
        return "CWE-20"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            name_lower = node.name.lower()
            is_notification = any(
                kw in name_lower
                for kw in ("notification", "on_notify", "handle_notif", "on_event")
            )
            if not is_notification:
                continue

            body_src = ast.get_source_segment(source, node) or ""
            has_validation = any(
                kw in body_src.lower()
                for kw in (
                    "validate", "schema", "isinstance", "pydantic",
                    "raise valueerror", "raise typeerror",
                )
            )
            if has_validation:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Notification handler accepts arbitrary data",
                description=(
                    f"Function '{node.name}' handles notifications "
                    f"without validating incoming data."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Validate notification data against a schema before "
                    "processing. Reject malformed notifications."
                ),
                cwe=self.cwe,
                confidence="HIGH",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-015: Server accepts unauthenticated connections
# ---------------------------------------------------------------------------

class MCPUnauthenticatedConnectionRule(Rule):
    """Detects servers that accept connections without authentication."""

    @property
    def rule_id(self) -> str:
        return "MCP-015"

    @property
    def name(self) -> str:
        return "Server Accepts Unauthenticated Connections"

    @property
    def severity(self) -> Severity:
        return Severity.CRITICAL

    @property
    def description(self) -> str:
        return (
            "MCP server starts listening for connections without "
            "requiring authentication, allowing unauthorized access."
        )

    @property
    def cwe(self) -> str:
        return "CWE-306"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        source_lower = source.lower()

        # Look for server start patterns
        server_start_patterns = [
            re.compile(r'''\.run\s*\('''),
            re.compile(r'''\.serve\s*\('''),
            re.compile(r'''\.start\s*\('''),
            re.compile(r'''run_server\s*\('''),
        ]

        has_server_start = any(p.search(source) for p in server_start_patterns)
        if not has_server_start:
            return findings

        # Check for network transports (not stdio)
        is_network = any(
            kw in source_lower
            for kw in ("sse", "websocket", "http", "tcp", "port", "host=")
        )
        if not is_network:
            return findings

        has_auth = _has_pattern(source, [
            "auth", "token", "api_key", "authenticate", "credential",
            "bearer", "middleware",
        ])
        if has_auth:
            return findings

        # Find the line where the server starts
        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            for pattern in server_start_patterns:
                if pattern.search(line):
                    findings.append(Finding(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        title="Server accepts unauthenticated connections",
                        description=(
                            f"Server starts on a network transport at line {i} "
                            f"without requiring authentication."
                        ),
                        file="",
                        line=i,
                        code_snippet=line.strip(),
                        recommendation=(
                            "Add authentication middleware before starting "
                            "the server on network transports."
                        ),
                        cwe=self.cwe,
                        confidence="HIGH",
                    ))
                    return findings
        return findings


# ---------------------------------------------------------------------------
# MCP-016: Tool call logging insufficient
# ---------------------------------------------------------------------------

class MCPToolCallLoggingRule(Rule):
    """Detects tool handlers without adequate logging."""

    @property
    def rule_id(self) -> str:
        return "MCP-016"

    @property
    def name(self) -> str:
        return "Tool Call Logging Insufficient"

    @property
    def severity(self) -> Severity:
        return Severity.LOW

    @property
    def description(self) -> str:
        return (
            "MCP tool handler lacks logging for security audit trails, "
            "making it difficult to detect or investigate attacks."
        )

    @property
    def cwe(self) -> str:
        return "CWE-778"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not _is_tool_handler(node):
                continue

            body_src = ast.get_source_segment(source, node) or ""
            has_logging = any(
                kw in body_src.lower()
                for kw in (
                    "logging", "logger", "log.", "log(", "print(",
                    "audit", "trace",
                )
            )
            if has_logging:
                continue

            findings.append(Finding(
                rule_id=self.rule_id,
                severity=self.severity,
                title="Tool handler lacks logging",
                description=(
                    f"Function '{node.name}' has no logging statements "
                    f"for security audit trails."
                ),
                file="",
                line=node.lineno,
                code_snippet=self._get_snippet(source, node.lineno),
                recommendation=(
                    "Add logging for tool invocations including caller ID, "
                    "parameters, and result status for audit trails."
                ),
                cwe=self.cwe,
                confidence="LOW",
            ))
        return findings


# ---------------------------------------------------------------------------
# MCP-017: Server-sent events without origin check
# ---------------------------------------------------------------------------

class MCPSSEWithoutOriginCheckRule(Rule):
    """Detects SSE endpoints without origin validation."""

    @property
    def rule_id(self) -> str:
        return "MCP-017"

    @property
    def name(self) -> str:
        return "Server-Sent Events Without Origin Check"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP server uses SSE transport without validating "
            "the client origin, risking cross-site attacks."
        )

    @property
    def cwe(self) -> str:
        return "CWE-346"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        source_lower = source.lower()

        has_sse = any(
            kw in source_lower
            for kw in ("sse", "server-sent", "eventsource", "text/event-stream")
        )
        if not has_sse:
            return findings

        has_origin_check = any(
            kw in source_lower
            for kw in ("origin", "cors", "allowed_origin", "access-control")
        )
        if has_origin_check:
            return findings

        # Find SSE setup line
        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            if any(kw in line.lower() for kw in ("sse", "server-sent", "eventsource", "text/event-stream")):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="SSE without origin validation",
                    description=(
                        f"SSE transport at line {i} lacks origin "
                        f"validation, exposing the server to cross-site attacks."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip(),
                    recommendation=(
                        "Validate the Origin header on SSE connections. "
                        "Configure CORS to restrict allowed origins."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
                return findings
        return findings


# ---------------------------------------------------------------------------
# MCP-018: Batch tool calls without rate limit
# ---------------------------------------------------------------------------

class MCPBatchCallsNoRateLimitRule(Rule):
    """Detects batch tool call support without rate limiting."""

    @property
    def rule_id(self) -> str:
        return "MCP-018"

    @property
    def name(self) -> str:
        return "Batch Tool Calls Without Rate Limit"

    @property
    def severity(self) -> Severity:
        return Severity.MEDIUM

    @property
    def description(self) -> str:
        return (
            "MCP server supports batch tool calls without rate limiting, "
            "enabling abuse through rapid-fire requests."
        )

    @property
    def cwe(self) -> str:
        return "CWE-770"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        source_lower = source.lower()

        has_batch = any(
            kw in source_lower
            for kw in ("batch", "bulk", "multi_call", "parallel_tools")
        )
        if not has_batch:
            return findings

        has_rate_limit = any(
            kw in source_lower
            for kw in (
                "rate_limit", "ratelimit", "throttle", "max_concurrent",
                "semaphore", "max_batch", "batch_limit",
            )
        )
        if has_rate_limit:
            return findings

        # Find batch handler line
        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            if any(kw in line.lower() for kw in ("batch", "bulk", "multi_call")):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Batch calls without rate limiting",
                    description=(
                        f"Batch processing at line {i} lacks rate "
                        f"limiting or concurrency controls."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip(),
                    recommendation=(
                        "Add rate limiting and max batch size controls "
                        "to batch tool call handlers."
                    ),
                    cwe=self.cwe,
                    confidence="MEDIUM",
                ))
                return findings
        return findings


# ---------------------------------------------------------------------------
# MCP-017: Config hot-reload without validation
# ---------------------------------------------------------------------------

class MCPConfigHotReloadRule(Rule):
    """Detects configuration hot-reload without validation."""

    @property
    def rule_id(self) -> str:
        return "MCP-019"

    @property
    def name(self) -> str:
        return "Config Hot-Reload Without Validation"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def description(self) -> str:
        return (
            "MCP server reloads configuration at runtime without "
            "validating the new config, risking injection or misconfiguration."
        )

    @property
    def cwe(self) -> str:
        return "CWE-94"

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        findings: list[Finding] = []
        source_lower = source.lower()

        has_reload = any(
            kw in source_lower
            for kw in ("reload_config", "hot_reload", "config_reload", "watch_config")
        )
        if not has_reload:
            return findings

        has_validation = any(
            kw in source_lower
            for kw in (
                "validate_config", "schema", "pydantic", "jsonschema",
                "validate(", "is_valid",
            )
        )
        if has_validation:
            return findings

        lines = source.splitlines()
        for i, line in enumerate(lines, 1):
            if any(kw in line.lower() for kw in ("reload_config", "hot_reload", "config_reload", "watch_config")):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Config hot-reload without validation",
                    description=(
                        f"Configuration reload at line {i} lacks "
                        f"validation of the new configuration."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip(),
                    recommendation=(
                        "Validate configuration against a schema before "
                        "applying hot-reloaded values. Reject invalid configs."
                    ),
                    cwe=self.cwe,
                    confidence="HIGH",
                ))
                return findings
        return findings


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

def get_all_mcp_rules() -> list[Rule]:
    """Return instances of all MCP protocol-specific rules."""
    return [
        MCPToolInputValidationRule(),
        MCPToolResponseSizeRule(),
        MCPServerAuthCheckRule(),
        MCPResourceListUnrestrictedRule(),
        MCPInternalPathExposureRule(),
        MCPToolNameCollisionRule(),
        MCPMissingTransportEncryptionRule(),
        MCPServerVersionDisclosureRule(),
        MCPToolSchemaMissingFieldsRule(),
        MCPNotificationArbitraryDataRule(),
        MCPUnauthenticatedConnectionRule(),
        MCPToolCallLoggingRule(),
        MCPSSEWithoutOriginCheckRule(),
        MCPBatchCallsNoRateLimitRule(),
        MCPConfigHotReloadRule(),
    ]
