"""
crypto_audit.py - Cryptographic pitfall scanner
Chapter 4, Lab 4.3 (Capstone audit, Section 4.13)

Scans Python source files for the ten recurring cryptographic pitfalls
from the field guide (Section 4.13) and produces a structured report
mapping each finding to its OWASP A02:2021 sub-category and the
chapter section that explains the failure.

Usage:
    python crypto_audit.py sample_vulnerable_code.py
    python crypto_audit.py /path/to/target/

The Juice Shop target (http://juice-shop:3000) hosts a Node.js app;
use the node-grep patterns below instead of Python patterns when auditing it.
See Lab 4.3 Walkthrough in the README for full instructions.

SAFETY NOTE: Only run this against the provided lab target (Juice Shop on
the lab network) or source code you own or are authorised to audit.
"""

import re
import sys
import os
from dataclasses import dataclass, field
from typing import List


@dataclass
class Finding:
    file: str
    line: int
    pitfall: str
    owasp: str
    chapter_section: str
    code_snippet: str
    remediation: str


# Pattern definitions - each tuple: (regex, pitfall_name, owasp, section, fix)
PATTERNS = [
    (
        re.compile(r'\bAES\.MODE_ECB\b|\bECB\b|["\']AES/ECB|AES\.new\([^,]+\)(?!\s*,\s*AES\.MODE_)', re.IGNORECASE),
        "ECB mode",
        "A02:2021 - Cryptographic Failures",
        "4.5.1",
        "Replace with AESGCM or ChaCha20Poly1305 from cryptography.hazmat.primitives.ciphers.aead"
    ),
    (
        re.compile(r'hashlib\.md5\b|hashlib\.sha1\b|MD5\(|SHA1\(', re.IGNORECASE),
        "Deprecated hash (MD5/SHA-1) for passwords or integrity",
        "A02:2021 - Cryptographic Failures",
        "4.6 / 4.7",
        "For passwords: use argon2-cffi (Argon2id). For integrity: use SHA-256 or SHA-3."
    ),
    (
        re.compile(r'\bimport random\b|\bfrom random import\b|random\.randint|random\.random\(\)|Math\.random\(\)'),
        "Statistical PRNG for security-sensitive value",
        "A02:2021 - Cryptographic Failures",
        "4.3.1",
        "Use the 'secrets' module (Python) or crypto.randomBytes() (Node.js)"
    ),
    (
        re.compile(r'(?<![a-z])(?:provided|token|tag|digest|mac)\s*==\s*(?:expected|computed|stored)', re.IGNORECASE),
        "Non-constant-time comparison of MAC/token",
        "A02:2021 - Cryptographic Failures",
        "4.6.3",
        "Use hmac.compare_digest() or secrets.compare_digest()"
    ),
    (
        re.compile(r'verify\s*=\s*False|rejectUnauthorized\s*:\s*false|InsecureSkipVerify\s*:\s*true', re.IGNORECASE),
        "Disabled TLS certificate verification",
        "A02:2021 / A07:2021",
        "4.10.2",
        "Remove verify=False; let the HTTP client use the system CA bundle"
    ),
    (
        re.compile(r'\bSHA256\s*\(\s*(?:secret|key|token|password)\s*\+|\bhashlib\.sha256\s*\(\s*(?:key|secret)\s*\+'),
        "Bare SHA-256 prefix-MAC (length-extension vulnerable)",
        "A02:2021 - Cryptographic Failures",
        "4.6.3",
        "Use hmac.new(key, message, hashlib.sha256) - HMAC closes the length-extension door"
    ),
    (
        re.compile(r'STATIC_IV\s*=\s*b["\'][\\\\x00]*["\']|iv\s*=\s*b["\'][\\\\x00]*["\']|nonce\s*=\s*b["\'][\\\\x00]*["\']', re.IGNORECASE),
        "Static / hardcoded IV or nonce",
        "A02:2021 - Cryptographic Failures",
        "4.5.2",
        "Generate a fresh random IV/nonce with os.urandom(12) or os.urandom(16) per message"
    ),
    (
        re.compile(r'\bMODE_CBC\b|\bCBC\b', re.IGNORECASE),
        "CBC mode (unauthenticated)",
        "A02:2021 - Cryptographic Failures",
        "4.5.2 / 4.5.3",
        "Prefer AES-GCM (AEAD); if CBC must be used, add encrypt-then-HMAC and verify MAC first"
    ),
]


def scan_file(filepath: str) -> List[Finding]:
    findings = []
    try:
        with open(filepath, "r", encoding="utf-8", errors="ignore") as fh:
            lines = fh.readlines()
    except OSError as e:
        print(f"[skip] {filepath}: {e}")
        return findings

    for lineno, line in enumerate(lines, start=1):
        stripped = line.rstrip("\n")
        for pattern, pitfall, owasp, section, fix in PATTERNS:
            if pattern.search(stripped):
                findings.append(Finding(
                    file=filepath,
                    line=lineno,
                    pitfall=pitfall,
                    owasp=owasp,
                    chapter_section=section,
                    code_snippet=stripped.strip(),
                    remediation=fix,
                ))
    return findings


def scan_target(target: str) -> List[Finding]:
    all_findings: List[Finding] = []
    if os.path.isfile(target):
        all_findings.extend(scan_file(target))
    elif os.path.isdir(target):
        for root, _, files in os.walk(target):
            for fname in files:
                if fname.endswith((".py", ".js", ".ts")):
                    all_findings.extend(scan_file(os.path.join(root, fname)))
    return all_findings


def print_report(findings: List[Finding]) -> None:
    if not findings:
        print("No findings. Scan complete.")
        return

    print(f"\n{'='*70}")
    print(f" CRYPTOGRAPHIC AUDIT REPORT  -  {len(findings)} finding(s)")
    print(f"{'='*70}\n")

    for i, f in enumerate(findings, start=1):
        print(f"[{i}] {f.file}:{f.line}")
        print(f"     Pitfall:  {f.pitfall}")
        print(f"     OWASP:    {f.owasp}")
        print(f"     Section:  Chapter {f.chapter_section}")
        print(f"     Code:     {f.code_snippet[:80]}")
        print(f"     Fix:      {f.remediation}")
        print()

    print(f"{'='*70}")
    print(f" Total findings: {len(findings)}")
    print(f" Next step: apply remediations, re-run this tool to verify fixes.")
    print(f"{'='*70}\n")


def main():
    targets = sys.argv[1:] if len(sys.argv) > 1 else ["sample_vulnerable_code.py"]
    all_findings: List[Finding] = []
    for t in targets:
        print(f"[scan] {t}")
        all_findings.extend(scan_target(t))
    print_report(all_findings)


if __name__ == "__main__":
    main()
