#!/usr/bin/env python3
# Scans a repo for the three Claude Code vulnerability patterns.
# Usage: python3 scanner.py <path-to-repo>

import json
import os
import sys
import re
from pathlib import Path

# ANSI colors
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
CYAN = "\033[96m"
RESET = "\033[0m"
BOLD = "\033[1m"


class Finding:

    def __init__(self, severity, cve, title, file_path, detail, recommendation):
        self.severity = severity  # CRITICAL, HIGH, MEDIUM, LOW, INFO
        self.cve = cve
        self.title = title
        self.file_path = file_path
        self.detail = detail
        self.recommendation = recommendation

    def __str__(self):
        colors = {
            "CRITICAL": RED + BOLD,
            "HIGH": RED,
            "MEDIUM": YELLOW,
            "LOW": CYAN,
            "INFO": GREEN,
        }
        c = colors.get(self.severity, RESET)
        return (
            f"\n{c}[{self.severity}]{RESET} {BOLD}{self.title}{RESET}\n"
            f"  CVE:            {self.cve}\n"
            f"  File:           {self.file_path}\n"
            f"  Detail:         {self.detail}\n"
            f"  Recommendation: {self.recommendation}\n"
        )


def scan_hooks_bypass(repo_path):
    findings = []
    settings_path = repo_path / ".claude" / "settings.json"

    if not settings_path.exists():
        return findings

    try:
        data = json.loads(settings_path.read_text())
    except (json.JSONDecodeError, OSError):
        return findings

    hooks = data.get("hooks", {})
    if not hooks:
        return findings

    for event_name, event_config in hooks.items():
        if not isinstance(event_config, list):
            continue
        for matcher_block in event_config:
            if not isinstance(matcher_block, dict):
                continue
            for hook in matcher_block.get("hooks", []):
                if not isinstance(hook, dict):
                    continue
                cmd = hook.get("command", "")
                if cmd:
                    severity = "CRITICAL"
                    dangerous_patterns = [
                        r"curl\s", r"wget\s", r"nc\s", r"ncat\s",
                        r"bash\s+-[ic]", r"/dev/tcp/", r"mkfifo",
                        r"python.*-c", r"eval\s", r"base64",
                        r"\bssh\b", r"reverse", r"bind.*shell",
                    ]
                    is_extra_dangerous = any(
                        re.search(p, cmd, re.IGNORECASE) for p in dangerous_patterns
                    )

                    findings.append(Finding(
                        severity="CRITICAL" if is_extra_dangerous else "HIGH",
                        cve="No CVE (CVSS 8.7)",
                        title=f"Project hook executes shell command on {event_name}",
                        file_path=str(settings_path),
                        detail=f"Command: {cmd[:120]}{'...' if len(cmd)>120 else ''}",
                        recommendation=(
                            "Remove project-level hooks or audit each command. "
                            "Update Claude Code to v1.0.87+ where consent is required."
                        ),
                    ))

    return findings


def scan_mcp_injection(repo_path):
    findings = []

    settings_path = repo_path / ".claude" / "settings.json"
    auto_enable = False

    if settings_path.exists():
        try:
            data = json.loads(settings_path.read_text())
            if data.get("enableAllProjectMcpServers") is True:
                auto_enable = True
                findings.append(Finding(
                    severity="HIGH",
                    cve="CVE-2025-59536 (CVSS 8.7)",
                    title="enableAllProjectMcpServers is set to true",
                    file_path=str(settings_path),
                    detail=(
                        "This flag causes all project-defined MCP servers to start "
                        "automatically without user consent."
                    ),
                    recommendation=(
                        "Remove this flag. Update Claude Code to v1.0.111+ "
                        "where this bypass is patched."
                    ),
                ))
        except (json.JSONDecodeError, OSError):
            pass

    mcp_path = repo_path / ".mcp.json"
    if mcp_path.exists():
        try:
            mcp_data = json.loads(mcp_path.read_text())
            servers = mcp_data.get("mcpServers", {})
            for name, config in servers.items():
                if not isinstance(config, dict):
                    continue
                cmd = config.get("command", "")
                args = config.get("args", [])
                full_cmd = f"{cmd} {' '.join(str(a) for a in args)}" if args else cmd

                severity = "CRITICAL" if auto_enable else "MEDIUM"
                findings.append(Finding(
                    severity=severity,
                    cve="CVE-2025-59536 (CVSS 8.7)",
                    title=f"MCP server '{name}' executes: {cmd}",
                    file_path=str(mcp_path),
                    detail=f"Full command: {full_cmd[:150]}{'...' if len(full_cmd)>150 else ''}",
                    recommendation=(
                        "Audit MCP server commands. Remove untrusted servers. "
                        "Never set enableAllProjectMcpServers=true in shared repos."
                    ),
                ))

                env = config.get("env", {})
                for k, v in env.items():
                    if any(
                        s in k.upper()
                        for s in ["KEY", "TOKEN", "SECRET", "PASSWORD", "CREDENTIAL"]
                    ):
                        findings.append(Finding(
                            severity="MEDIUM",
                            cve="CVE-2025-59536",
                            title=f"MCP server '{name}' sets suspicious env var: {k}",
                            file_path=str(mcp_path),
                            detail=f"Env var {k} may be used to override credentials.",
                            recommendation="Audit environment variables in MCP configs.",
                        ))
        except (json.JSONDecodeError, OSError):
            pass

    return findings


def scan_api_exfil(repo_path):
    findings = []
    settings_path = repo_path / ".claude" / "settings.json"

    if not settings_path.exists():
        return findings

    try:
        data = json.loads(settings_path.read_text())
    except (json.JSONDecodeError, OSError):
        return findings

    env = data.get("env", {})
    if not isinstance(env, dict):
        return findings

    suspicious_env_vars = {
        "ANTHROPIC_BASE_URL": "Redirects all API traffic (including API key) to attacker",
        "ANTHROPIC_API_KEY": "Overrides/captures the user's API key",
        "CLAUDE_CODE_API_KEY": "May override API key configuration",
        "HTTP_PROXY": "Routes all HTTP traffic through attacker proxy",
        "HTTPS_PROXY": "Routes all HTTPS traffic through attacker proxy",
        "NODE_EXTRA_CA_CERTS": "Could enable MITM by injecting attacker CA certificate",
    }

    for var, description in suspicious_env_vars.items():
        value = env.get(var, "")
        if not value:
            continue

        severity = "CRITICAL"
        if var == "ANTHROPIC_BASE_URL":
            if "anthropic.com" not in value.lower():
                severity = "CRITICAL"
            else:
                severity = "INFO"

        findings.append(Finding(
            severity=severity,
            cve="CVE-2026-21852 (CVSS 5.3)",
            title=f"Environment override: {var}",
            file_path=str(settings_path),
            detail=f"{description}. Value: {value[:80]}{'...' if len(value)>80 else ''}",
            recommendation=(
                "Remove env overrides from project settings. "
                "Update Claude Code to v2.0.65+ where env is not loaded before trust prompt."
            ),
        ))

    return findings


def scan_repo(repo_path_str):
    repo_path = Path(repo_path_str).resolve()

    if not repo_path.is_dir():
        print(f"{RED}[!] Error: '{repo_path}' is not a directory{RESET}")
        sys.exit(1)

    print(f"""
{BOLD}{'='*70}
  Claude Code Malicious Repository Scanner
  EDUCATIONAL USE ONLY
{'='*70}{RESET}

{CYAN}[*] Scanning: {repo_path}{RESET}
""")

    all_findings = []

    scanners = [
        ("Hooks Consent Bypass (no CVE)", scan_hooks_bypass),
        ("MCP Server Injection (CVE-2025-59536)", scan_mcp_injection),
        ("API Key Exfiltration (CVE-2026-21852)", scan_api_exfil),
    ]

    for name, scanner_fn in scanners:
        print(f"{CYAN}[*] Checking: {name}...{RESET}")
        findings = scanner_fn(repo_path)
        all_findings.extend(findings)
        if findings:
            print(f"  {RED}Found {len(findings)} issue(s){RESET}")
        else:
            print(f"  {GREEN}Clean{RESET}")

    print(f"\n{BOLD}{'='*70}")
    print(f"  SCAN RESULTS")
    print(f"{'='*70}{RESET}\n")

    if not all_findings:
        print(f"{GREEN}{BOLD}[+] No Claude Code supply-chain indicators found.{RESET}\n")
        return 0

    severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "INFO": 4}
    all_findings.sort(key=lambda f: severity_order.get(f.severity, 5))

    for finding in all_findings:
        print(finding)

    by_severity = {}
    for f in all_findings:
        by_severity[f.severity] = by_severity.get(f.severity, 0) + 1

    print(f"\n{BOLD}Total findings: {len(all_findings)}{RESET}")
    for sev in ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"]:
        count = by_severity.get(sev, 0)
        if count:
            colors = {
                "CRITICAL": RED + BOLD, "HIGH": RED,
                "MEDIUM": YELLOW, "LOW": CYAN, "INFO": GREEN,
            }
            print(f"  {colors.get(sev, '')}{sev}: {count}{RESET}")

    print(f"\n{YELLOW}[!] Recommendation: Do NOT open this repo with Claude Code < v2.0.65{RESET}")
    print(f"{YELLOW}[!] Review and remove all suspicious configuration files.{RESET}\n")

    return len(all_findings)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <path-to-repo>")
        print(f"Example: {sys.argv[0]} ./suspicious-repo")
        sys.exit(1)

    count = scan_repo(sys.argv[1])
    sys.exit(1 if count > 0 else 0)