import argparse
from pathlib import Path
from typing import Iterable, List, Tuple

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


def decrypt_xts_sector(ct: bytes, xts_key: bytes, sector_no: int) -> bytes:
    tweak = int(sector_no).to_bytes(16, "little", signed=False)
    cipher = Cipher(algorithms.AES(xts_key), modes.XTS(tweak))
    dec = cipher.decryptor()
    return dec.update(ct) + dec.finalize()


def looks_like_boot_sector(pt: bytes) -> List[str]:
    hits: List[str] = []
    if len(pt) < 512:
        return hits
    if pt[510:512] == b"\x55\xaa":
        hits.append("55aa")
    sig = pt[3:11]
    if sig == b"NTFS    ":
        hits.append("NTFS")
    if sig == b"EXFAT   ":
        hits.append("EXFAT")
    if sig.startswith(b"FAT"):
        hits.append(sig.decode("ascii", errors="ignore"))
    if pt[0] in (0xEB, 0xE9) and pt[2] == 0x90:
        hits.append("jmp")
    return hits


def main() -> int:
    ap = argparse.ArgumentParser(description="Probe VeraCrypt container using AES-XTS keys")
    ap.add_argument("container", type=Path)
    ap.add_argument("--key-a", required=True, help="32-byte hex key A")
    ap.add_argument("--key-b", required=True, help="32-byte hex key B")
    ap.add_argument(
        "--offsets",
        default="0,0x10000,0x20000",
        help="Comma-separated file offsets to try (default: 0,0x10000,0x20000)",
    )
    ap.add_argument(
        "--tweak-bases",
        default="auto",
        help="Comma-separated sector numbers to try as tweak base, or 'auto' for 0 and offset/512",
    )
    args = ap.parse_args()

    key_a = bytes.fromhex(args.key_a)
    key_b = bytes.fromhex(args.key_b)
    if len(key_a) != 32 or len(key_b) != 32:
        raise SystemExit("Keys must be 32 bytes each (64 hex chars)")

    offsets: List[int] = []
    for part in args.offsets.split(","):
        part = part.strip()
        if not part:
            continue
        offsets.append(int(part, 0))

    data = args.container.read_bytes()

    key_orders: List[Tuple[str, bytes]] = [
        ("A||B", key_a + key_b),
        ("B||A", key_b + key_a),
    ]

    for off in offsets:
        if off + 512 > len(data):
            continue
        ct = data[off : off + 512]

        if args.tweak_bases.strip().lower() == "auto":
            tweak_bases = sorted({0, off // 512})
        else:
            tweak_bases = []
            for part in args.tweak_bases.split(","):
                part = part.strip()
                if not part:
                    continue
                tweak_bases.append(int(part, 0))

        for base in tweak_bases:
            for label, xts_key in key_orders:
                pt = decrypt_xts_sector(ct, xts_key, base)
                hits = looks_like_boot_sector(pt)
                if hits:
                    print(
                        f"[+] offset=0x{off:X} tweak={base} order={label} hits={','.join(hits)} sig={pt[3:11]!r}"
                    )
                    print(pt[:64].hex())

    return 0


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