import argparse
import os
import struct
from pathlib import Path
from typing import Iterable, Iterator, Tuple


def iter_input_files(paths: Iterable[str]) -> Iterator[Path]:
    for p in paths:
        path = Path(p)
        if path.is_dir():
            for child in sorted(path.rglob("*")):
                if child.is_file():
                    yield child
        elif path.is_file():
            yield path


def scan_password_structs(
    data: bytes, *, min_len: int, max_len: int
) -> Iterator[Tuple[int, int, str]]:
    """
    Heuristic scan for the (TrueCrypt/VeraCrypt) Password struct:
      uint32 Length; char Text[...];

    We look for:
      <len:uint32le> <len bytes printable ASCII> <any byte> 0x00 0x00 0x00

    This mirrors volatility3's truecrypt passphrase finder which validates the
    3 bytes *after* the presumed NUL terminator but doesn't explicitly check
    the terminator byte itself.
    """
    mv = memoryview(data)
    n = len(data)
    min_total = 4 + min_len + 4
    if n < min_total:
        return


    unpack_from = struct.unpack_from
    for i in range(0, n - min_total + 1):
        (length,) = unpack_from("<I", mv, i)
        if length < min_len or length > max_len:
            continue

        start = i + 4
        end = start + length
        tail = end + 4  
        if tail > n:
            continue

        pw = mv[start:end]

        if any((c < 0x20 or c >= 0x7F) for c in pw):
            continue

        if data[end + 1 : tail] != b"\x00\x00\x00":
            continue

        try:
            pw_str = pw.tobytes().decode("ascii")
        except UnicodeDecodeError:
            continue

        yield i, length, pw_str


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("paths", nargs="+", help="Files and/or directories to scan")
    ap.add_argument("--min-len", type=int, default=5)
    ap.add_argument("--max-len", type=int, default=64)
    args = ap.parse_args()

    seen = set()
    hits = 0
    for fp in iter_input_files(args.paths):
        try:
            data = fp.read_bytes()
        except Exception as exc:
            print(f"[!] Failed to read {fp}: {exc}")
            continue

        for off, length, pw in scan_password_structs(
            data, min_len=args.min_len, max_len=args.max_len
        ):
            key = (pw,)
            if key in seen:
                continue
            seen.add(key)
            hits += 1
            print(f"{fp}\t0x{off:X}\t{length}\t{pw}")

    if hits == 0:
        print("[!] No candidates found")
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())