#!/usr/bin/env python3
"""sweep.py — secrets and identities in a corpus, including inside archives.

Usage:
  sweep.py <dir> [--report out.json] [--allow-domain example.dev ...]
  sweep.py <dir> --add-literal          read one value from stdin and add its salted hash to the literal list

Three passes over every text file and every member of .gz, .tgz and .zip archives:
  1. hashed literals: values known to be secret or personal for this target, stored only as
     salted sha256 in <target>/.literals.sha256, so the list itself leaks nothing;
  2. shapes: vendor keys, JWTs, private keys, emails, home paths;
  3. key names: JSON keys that carry identities or credentials with a non-redacted value.
Markers of unfinished redaction ("pending remediation", "NOT REDACTED", "TODO redact") also fail.

Exit 0 clean, 1 hit, 2 an archive could not be read. Waivers: run.json "waive" = [{"path": "evidence/dimensions/raw", "classes": ["jwt","home-path"]}] (set with
run_init.py --waive PATH:class,class) drops hits of those classes under that path prefix, for verbatim public pages.
An email waiver holds for one named file only; a directory never waives email, since a public page can carry a person.
Waived hits are counted and printed per path and class, never silently; literals, key names, markers and vendor
secret shapes cannot be waived.

A hit prints path, line, class and a masked
value, never the value. The report carries the manifest hash: the sorted list of files and sizes swept,
so "clean" names exactly what was looked at.
"""
import gzip, hashlib, io, json, os, re, sys, tarfile, zipfile

SHAPES = [
    ("jwt", re.compile(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}")),
    ("private-key", re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----")),
    ("aws-key", re.compile(r"\b(AKIA|ASIA)[A-Z0-9]{16}\b")),
    ("google-key", re.compile(r"\bAIza[0-9A-Za-z_-]{35}\b")),
    ("github-token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{30,}\b")),
    ("slack-token", re.compile(r"\bxox[abposr]-[A-Za-z0-9-]{10,}\b")),
    ("vendor-key", re.compile(r"\b(sk|pk|rk)[-_](live|test|proj|ant)?[-_]?[A-Za-z0-9]{20,}\b")),
    ("bearer", re.compile(r"\b[Bb]earer\s+[A-Za-z0-9._~+/-]{20,}=*")),
    ("home-path", re.compile(r"/Users/[a-z][a-z0-9._-]+/|/home/[a-z][a-z0-9._-]+/")),
    ("email", re.compile(r"\b[A-Za-z0-9._%+-]+@([A-Za-z0-9-]+\.)+[A-Za-z]{2,}\b")),
]
KEY_NAMES = re.compile(r'"(password|passwd|secret|client_secret|api_key|apikey|access_token|refresh_token|id_token|session_token|private_key|authorization|cookie)"\s*:\s*"(?!<REDACTED)([^"]{6,})"', re.I)
MARKERS = re.compile(r"pending remediation|NOT REDACTED|TODO[: ]+redact", re.I)
WAIVABLE = {"jwt", "home-path", "email"}
SKIP_EMAIL_DOMAINS = {"example.com", "example.org", "example.net"}
TEXT_EXT = {".md", ".txt", ".json", ".jsonl", ".js", ".mjs", ".ts", ".html", ".htm", ".csv", ".yaml", ".yml", ".xml", ".har", ".log", ".toml", ".py", ".sh", ".css", ".map", ".svg"}
TOKEN = re.compile(r"[A-Za-z0-9_@.+-]{3,}")

def literal_file(root):
    p = os.path.abspath(root)
    for _ in range(4):
        f = os.path.join(p, ".literals.sha256")
        if os.path.exists(f):
            return f
        p = os.path.dirname(p)
    return os.path.join(os.path.abspath(root), ".literals.sha256")

def load_literals(path):
    if not os.path.exists(path):
        return None, {}
    lines = open(path).read().split("\n")
    salt = lines[0].split("salt:", 1)[1].strip() if lines and lines[0].startswith("salt:") else ""
    by_len = {}
    for line in lines[1:]:
        parts = line.split()
        if len(parts) == 2:
            by_len.setdefault(int(parts[0]), set()).add(parts[1])
    return salt, by_len

def add_literal(root):
    path = literal_file(root)
    value = sys.stdin.readline().rstrip("\n")
    if not value:
        print("no value on stdin")
        return 2
    if not os.path.exists(path):
        with open(path, "w") as f:
            f.write("salt: " + os.urandom(8).hex() + "\n")
    salt, _ = load_literals(path)
    with open(path, "a") as f:
        f.write(f"{len(value)} {hashlib.sha256((salt + value).encode()).hexdigest()}\n")
    print(f"added a literal of length {len(value)} to {path}")
    return 0

def mask(v):
    return v[:2] + "…" + v[-2:] if len(v) > 6 else "…"

def scan_text(name, text, salt, lits, allow, hits):
    for i, line in enumerate(text.splitlines(), 1):
        for cls, rx in SHAPES:
            for m in rx.finditer(line):
                v = m.group(0)
                if cls == "email":
                    dom = v.split("@", 1)[1].lower()
                    if dom in SKIP_EMAIL_DOMAINS or any(dom == a or dom.endswith("." + a) for a in allow):
                        continue
                if "REDACTED" in line[max(0, m.start() - 12):m.end() + 12]:
                    continue
                hits.append((name, i, cls, mask(v)))
        for m in KEY_NAMES.finditer(line):
            hits.append((name, i, "key:" + m.group(1).lower(), mask(m.group(2))))
        if MARKERS.search(line):
            hits.append((name, i, "marker", MARKERS.search(line).group(0)))
        if lits:
            toks = TOKEN.findall(line)
            cands = set(toks) | {a + " " + b for a, b in zip(toks, toks[1:])}
            for c in cands:
                for t in (c, c.lower()):
                    if len(t) in lits and hashlib.sha256((salt + t).encode()).hexdigest() in lits[len(t)]:
                        hits.append((name, i, "literal", "len=" + str(len(t))))

def members(path):
    """(name, bytes) for a plain file, or for every member of an archive."""
    low = path.lower()
    if low.endswith((".tgz", ".tar.gz")):
        with tarfile.open(path, "r:gz") as t:
            for m in t.getmembers():
                if m.isfile():
                    yield path + "!" + m.name, t.extractfile(m).read()
    elif low.endswith(".gz"):
        yield path + "!", gzip.open(path).read()
    elif low.endswith(".zip"):
        with zipfile.ZipFile(path) as z:
            for n in z.namelist():
                if not n.endswith("/"):
                    yield path + "!" + n, z.read(n)
    else:
        yield path, open(path, "rb").read()

def is_text(name, data):
    ext = os.path.splitext(name.split("!")[-1])[1].lower()
    return ext in TEXT_EXT or (b"\x00" not in data[:4096] and ext not in {".png", ".jpg", ".jpeg", ".gif", ".webp", ".woff", ".woff2", ".pdf"})

def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 2
    root = argv[1]
    if "--add-literal" in argv:
        return add_literal(root)
    waive = []
    allow = [argv[i + 1].lower() for i, a in enumerate(argv) if a == "--allow-domain" and i + 1 < len(argv)]
    cfg = os.path.join(os.getcwd(), ".discover", "run.json")
    if os.path.exists(cfg):
        c = json.load(open(cfg))
        allow += [d.lower() for d in c.get("synthetic_domains", []) + c.get("allow_domains", [])]
        waive = [(w["path"].rstrip("/"), set(w["classes"]) & WAIVABLE - ({"email"} if os.path.isdir(os.path.join(root, w["path"])) else set()))
                 for w in c.get("waive", [])]
    salt, lits = load_literals(literal_file(root))
    hits, manifest, bad = [], [], []
    for dp, dns, fns in os.walk(root):
        dns[:] = [d for d in dns if d not in (".git", "node_modules", ".pii-quarantine")]
        for fn in sorted(fns):
            if fn == ".literals.sha256":
                continue
            p = os.path.join(dp, fn)
            try:
                for name, data in members(p):
                    manifest.append(f"{os.path.relpath(name, root)}\t{len(data)}")
                    if is_text(name, data):
                        scan_text(os.path.relpath(name, root), data.decode("utf-8", "replace"), salt, lits, allow, hits)
            except (OSError, tarfile.TarError, zipfile.BadZipFile, EOFError) as e:
                bad.append((p, str(e)))
    waived = {}
    kept = []
    for h in hits:
        w = next((p for p, cl in waive if h[2] in cl and (h[0] == p or h[0].startswith(p + "/"))), None)
        if w:
            waived[(w, h[2])] = waived.get((w, h[2]), 0) + 1
        else:
            kept.append(h)
    hits = kept
    for (w, cls), n in sorted(waived.items()):
        print(f"WAIVED {n} {cls} hit(s) under {w}")
    digest = hashlib.sha256("\n".join(sorted(manifest)).encode()).hexdigest()[:16]
    for h in hits:
        print(f"HIT {h[0]}:{h[1]} {h[2]} {h[3]}")
    for b in bad:
        print(f"UNREADABLE {b[0]}: {b[1]}")
    patterns = [c for c, _ in SHAPES] + ["key-names", "markers"] + (["literals"] if lits else [])
    print(f"swept {len(manifest)} files/members · manifest {digest} · patterns {','.join(patterns)} · {len(hits)} hit(s)")
    if "--report" in argv:
        with open(argv[argv.index("--report") + 1], "w") as f:
            json.dump({"manifest": digest, "files": len(manifest), "patterns": patterns, "hits": [list(h) for h in hits], "waived": [[w, c, n] for (w, c), n in sorted(waived.items())], "unreadable": bad}, f, indent=1)
    return 1 if hits else (2 if bad else 0)

if __name__ == "__main__":
    sys.exit(main(sys.argv))