#!/usr/bin/env python3
"""
pqc-crypto-scan — a lightweight cryptographic-inventory scanner for PQC readiness.

Scans a source tree for cryptographic usage (algorithm names, crypto libraries,
key/cert material), classifies each finding by quantum risk, and emits a draft
CBOM (Cryptography Bill of Materials, CycloneDX-style) plus a console summary.

Heuristic and dependency-free (Python 3 standard library only). It is a starting
point for a cryptographic inventory, NOT an authoritative audit — review findings
by hand and confirm against source.

Usage:
    python3 pqc-crypto-scan.py <path> [--out cbom.json] [--quiet]

Part of the QianHeng PQC docs — https://qusecurelabs.com/pqc-docs/discovery.html
"""
import os, re, json, sys, argparse, datetime

# (regex, canonical name, primitive, family, status)
ALGOS = [
    (r'AES-?256|aes_?256',            'AES-256',   'block-cipher',  'symmetric', 'symmetric-ok'),
    (r'\bAES\b|aes-?(128|192)',       'AES',       'block-cipher',  'symmetric', 'quantum-weakened'),
    (r'ChaCha20',                     'ChaCha20',  'stream-cipher', 'symmetric', 'quantum-weakened'),
    (r'\bSM4\b',                      'SM4',       'block-cipher',  'symmetric', 'quantum-weakened'),
    (r'\bZUC\b',                      'ZUC',       'stream-cipher', 'symmetric', 'quantum-weakened'),
    (r'\bSM3\b',                      'SM3',       'hash',          'hash',      'quantum-weakened'),
    (r'SHA-?(256|384|512)|SHA-?3|SHA3', 'SHA-2/3', 'hash',          'hash',      'quantum-weakened'),
    (r'\bRSA\b|RSA-?\d{3,4}',         'RSA',       'pke/signature', 'rsa',       'quantum-vulnerable'),
    (r'\bECDSA\b',                    'ECDSA',     'signature',     'ecc',       'quantum-vulnerable'),
    (r'\bECDH\b|ECDHE',              'ECDH(E)',   'key-agreement', 'ecc',       'quantum-vulnerable'),
    (r'Ed25519',                      'Ed25519',   'signature',     'ecc',       'quantum-vulnerable'),
    (r'X25519',                       'X25519',    'key-agreement', 'ecc',       'quantum-vulnerable'),
    (r'secp256r1|prime256v1|secp384r1|P-256|P-384', 'EC curve', 'ecc', 'ecc',   'quantum-vulnerable'),
    (r'\bDSA\b',                      'DSA',       'signature',     'dlog',      'quantum-vulnerable'),
    (r'Diffie-?Hellman',             'DH',        'key-agreement', 'dlog',      'quantum-vulnerable'),
    (r'\bSM2\b',                      'SM2',       'pke/signature', 'ecc',       'quantum-vulnerable'),
    (r'\bSM9\b',                      'SM9',       'ibc',           'pairing',   'quantum-vulnerable'),
    (r'\b3DES\b|TripleDES|\bDES\b',  '3DES/DES',  'block-cipher',  'symmetric', 'legacy-weak'),
    (r'\bRC4\b',                      'RC4',       'stream-cipher', 'symmetric', 'legacy-weak'),
    (r'\bMD5\b',                      'MD5',       'hash',          'hash',      'legacy-weak'),
    (r'SHA-?1\b',                     'SHA-1',     'hash',          'hash',      'legacy-weak'),
    (r'ML-?KEM|\bKyber\b',           'ML-KEM',    'kem',           'lattice',   'pqc'),
    (r'ML-?DSA|Dilithium',           'ML-DSA',    'signature',     'lattice',   'pqc'),
    (r'SLH-?DSA|SPHINCS',            'SLH-DSA',   'signature',     'hash',      'pqc'),
    (r'\bFalcon\b|FN-?DSA',          'FN-DSA',    'signature',     'lattice',   'pqc'),
    (r'\bHQC\b',                      'HQC',       'kem',           'code',      'pqc'),
    (r'\bXMSS\b|\bLMS\b',            'XMSS/LMS',  'signature',     'hash',      'pqc'),
]
LIBS = r'(openssl|libcrypto|boringssl|libsodium|liboqs|oqs-provider|bouncycastle|bcprov|cryptography\.hazmat|pyca|gmssl|tongsuo|wolfssl|mbedtls|forge|jsrsasign)'
KEY_EXT = {'.pem', '.crt', '.cer', '.der', '.key', '.p12', '.pfx', '.jks', '.keystore', '.pub'}
SCAN_EXT = {'.py', '.js', '.ts', '.java', '.go', '.c', '.h', '.cc', '.cpp', '.rs', '.rb', '.php',
            '.cs', '.kt', '.swift', '.scala', '.conf', '.cnf', '.cfg', '.ini', '.toml', '.yaml',
            '.yml', '.json', '.xml', '.properties', '.gradle', '.sh', '.env', '.tf', '.txt', '.md'}
SKIP_DIR = {'.git', 'node_modules', 'vendor', 'dist', 'build', '__pycache__', '.venv', 'venv', 'target'}
STATUS_RISK = {'quantum-vulnerable': 'HIGH', 'legacy-weak': 'HIGH', 'quantum-weakened': 'MEDIUM',
               'symmetric-ok': 'LOW', 'pqc': 'NONE'}
COMPILED = [(re.compile(p, re.I), n, prim, fam, st) for p, n, prim, fam, st in ALGOS]
LIB_RE = re.compile(LIBS, re.I)

def scan(root):
    findings = {}           # name -> {meta, occurrences:[(file,line)], libs:set}
    keymat = []
    for dirpath, dirs, files in os.walk(root):
        dirs[:] = [d for d in dirs if d not in SKIP_DIR]
        for fn in files:
            ext = os.path.splitext(fn)[1].lower()
            full = os.path.join(dirpath, fn)
            rel = os.path.relpath(full, root)
            if ext in KEY_EXT:
                keymat.append(rel)
                continue
            if ext not in SCAN_EXT and fn not in ('Dockerfile', 'Makefile'):
                continue
            try:
                with open(full, encoding='utf-8', errors='ignore') as fh:
                    lines = fh.readlines()
            except Exception:
                continue
            for i, line in enumerate(lines, 1):
                if len(line) > 4000:
                    continue
                for rx, name, prim, fam, st in COMPILED:
                    if rx.search(line):
                        f = findings.setdefault(name, {'primitive': prim, 'family': fam,
                                                        'status': st, 'occ': [], 'libs': set()})
                        if len(f['occ']) < 25:
                            f['occ'].append('%s:%d' % (rel, i))
                        m = LIB_RE.search(line)
                        if m:
                            f['libs'].add(m.group(1).lower())
    return findings, keymat

def to_cbom(findings, keymat, target):
    comps = []
    for name, f in sorted(findings.items()):
        comps.append({
            "type": "cryptographic-asset",
            "name": name,
            "cryptoProperties": {
                "assetType": "algorithm",
                "algorithmProperties": {"primitive": f['primitive'], "cryptoFamily": f['family']},
            },
            "properties": [
                {"name": "quantumRisk", "value": STATUS_RISK[f['status']]},
                {"name": "status", "value": f['status']},
                {"name": "libraries", "value": ", ".join(sorted(f['libs'])) or "n/a"},
                {"name": "occurrences", "value": str(len(f['occ']))},
            ],
            "evidence": {"occurrences": [{"location": o} for o in f['occ']]},
        })
    for k in keymat:
        comps.append({"type": "cryptographic-asset", "name": os.path.basename(k),
                      "cryptoProperties": {"assetType": "relatedCryptoMaterial"},
                      "properties": [{"name": "quantumRisk", "value": "REVIEW"},
                                     {"name": "location", "value": k}]})
    return {"bomFormat": "CycloneDX", "specVersion": "1.6",
            "metadata": {"timestamp": datetime.datetime.now().isoformat(timespec='seconds'),
                         "tools": [{"name": "pqc-crypto-scan", "version": "1.0"}],
                         "component": {"type": "application", "name": os.path.basename(os.path.abspath(target)) or "target"}},
            "components": comps}

def summarize(findings, keymat):
    buckets = {}
    for name, f in findings.items():
        buckets.setdefault(STATUS_RISK[f['status']], []).append(name)
    print("\n=== PQC readiness summary ===")
    for risk in ('HIGH', 'MEDIUM', 'LOW', 'NONE'):
        if buckets.get(risk):
            print("  %-7s %s" % (risk, ", ".join(sorted(buckets[risk]))))
    if keymat:
        print("  REVIEW  %d key/cert file(s) found" % len(keymat))
    hi = len(buckets.get('HIGH', []))
    print("\n  -> %d quantum-vulnerable/legacy algorithm class(es) need migration." % hi
          if hi else "\n  -> No quantum-vulnerable algorithms detected (verify manually).")

def main():
    ap = argparse.ArgumentParser(description="PQC crypto-inventory scanner -> draft CBOM")
    ap.add_argument("path")
    ap.add_argument("--out", default="cbom.json", help="output CBOM JSON path")
    ap.add_argument("--quiet", action="store_true")
    a = ap.parse_args()
    if not os.path.exists(a.path):
        sys.exit("path not found: " + a.path)
    findings, keymat = scan(a.path)
    cbom = to_cbom(findings, keymat, a.path)
    with open(a.out, "w", encoding="utf-8") as fh:
        json.dump(cbom, fh, ensure_ascii=False, indent=2)
    if not a.quiet:
        summarize(findings, keymat)
        print("\nwrote", a.out, "(%d components)" % len(cbom["components"]))

if __name__ == "__main__":
    main()