#!/usr/bin/env python3
"""Security validation helper for MCP server configurations.

Checks auth patterns, secret hygiene, least-privilege tool design,
and common exposure vectors.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from typing import Any


@dataclass
class ToolDef:
    name: str
    description: str
    is_mutation: bool = False  # write/delete/publish
    requires_auth: bool = False
    scope: str = "global"  # "personal" | "team" | "global"


@dataclass
class SecurityFinding:
    severity: str  # "critical" | "high" | "medium" | "low"
    tool: str
    message: str


def audit_server(tools: list[ToolDef], config: dict[str, Any]) -> list[SecurityFinding]:
    """Run the MCP security checklist against a server's tool surface."""
    findings: list[SecurityFinding] = []

    for tool in tools:
        # Critical: mutation without auth
        if tool.is_mutation and not tool.requires_auth:
            findings.append(SecurityFinding(
                "critical", tool.name,
                "Mutation tool has NO auth requirement. Any caller can invoke it.",
            ))

        # Critical: global-scope mutation by regular agents
        if tool.is_mutation and tool.scope == "global" and "super" not in tool.description.lower():
            findings.append(SecurityFinding(
                "critical", tool.name,
                "Global-scope mutation may be accessible to non-admin agents.",
            ))

        # High: tool that echoes user input back (potential injection vector)
        if any(kw in tool.name.lower() for kw in ("echo", "reflect", "reply")):
            findings.append(SecurityFinding(
                "high", tool.name,
                f"Tool name suggests it echoes input - potential output-trust issue.",
            ))

    # Auth header check
    if not config.get("require_header_auth", False):
        findings.append(SecurityFinding(
            "high", "server",
            "No header-based auth requirement. Agent keys as tool args only is insufficient.",
        ))

    # Transport check
    if config.get("transport") == "http" and not config.get("auth_enabled", False):
        findings.append(SecurityFinding(
            "critical", "server",
            "HTTP transport with no auth: tools are internet-accessible without credentials.",
        ))

    # Secret hygiene
    if config.get("log_args", False) and config.get("log_full_body", False):
        findings.append(SecurityFinding(
            "high", "server",
            "Logging full tool args may leak secrets (API keys, tokens).",
        ))

    return findings


# ── Demo ───────────────────────────────────────────────────────────────────

def main() -> None:
    # A realistic but flawed toolset
    tools = [
        ToolDef("publish_skill", "Publish a skill to the global store", is_mutation=True, scope="global"),
        ToolDef("delete_skill", "Delete a skill permanently", is_mutation=True, requires_auth=True, scope="personal"),
        ToolDef("get_skill", "Read skill content", is_mutation=False),
        ToolDef("search_skills", "Search skill library", is_mutation=False),
    ]

    config = {
        "transport": "http",
        "auth_enabled": False,      # BAD
        "require_header_auth": False, # BAD
        "log_args": True,
        "log_full_body": True,       # BAD
    }

    print("=== MCP Security Audit ===\n")
    for f in audit_server(tools, config):
        print(f"[{f.severity:8s}] {f.tool:20s} - {f.message}")

    print("\nHardening checklist:")
    print("  1. Require auth header ON EVERY mutation tool")
    print("  2. Scope tools: personal < team < global")
    print("  3. Enforce ownership: only owner may update/delete")
    print("  4. Never log raw keys; store only hashes")
    print("  5. HTTP transport ALWAYS requires auth + rate limiting")
    print("  6. Validate tool args server-side (never trust the model)")


if __name__ == "__main__":
    main()