#!/usr/bin/env python3

import os, sys, struct, argparse
from Crypto.Cipher import AES

KEY_DEFAULT = b"rK7CcATuZk7LAmhqqU4iBLNmAq8QbK3s"
KEY_ALTERNATE = b"x4DZmD6D2HhkzT6qD8HpKeZdgM9HCmXP"


def xor_crypt(data, key):
    cipher = AES.new(key, AES.MODE_ECB)
    out = bytearray(len(data))
    for off in range(0, len(data), 16):
        ctr = off // 16 + 1
        ks = cipher.encrypt(struct.pack('<q', ctr) + b'\x00' * 8)
        for j, b in enumerate(data[off:off+16]):
            out[off+j] = b ^ ks[j]
    return bytes(out)


def process_file(inpath, outpath, key, dry_run=False):
    try:
        with open(inpath, 'rb') as f:
            data = f.read()
        if data[:7] != b'UnityFS':
            print(f"  Warning: '{os.path.basename(inpath)}' does not start with UnityFS magic")
        if dry_run:
            print(f"  Would encrypt: {os.path.basename(inpath)} -> {os.path.basename(outpath)}")
            return True
        encrypted = xor_crypt(data, key)
        os.makedirs(os.path.dirname(outpath) or '.', exist_ok=True)
        with open(outpath, 'wb') as f:
            f.write(encrypted)
        print(f"  {os.path.basename(inpath)} -> {os.path.basename(outpath)}  ({len(data)} bytes)")
        return True
    except Exception as e:
        print(f"  ERROR {os.path.basename(inpath)}: {e}", file=sys.stderr)
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Re-encrypt decrypted Unity bundles for Holo's Hanafuda")
    parser.add_argument('input', help='Input .unity3d file or directory of .unity3d files')
    parser.add_argument('-o', '--output', help='Output file or directory')
    parser.add_argument('--key-alternate', action='store_true',
                        help='Use alternate bundle key candidate (x4DZmD6D2HhkzT6qD8HpKeZdgM9HCmXP)')
    parser.add_argument('--key-hex', metavar='HEX',
                        help='64-char hex AES-256 key (overrides defaults)')
    parser.add_argument('--in-place', action='store_true',
                        help='Write output next to the original .bundle in StandaloneWindows64/')
    parser.add_argument('--dry-run', action='store_true',
                        help='List what would be encrypted without writing')
    args = parser.parse_args()

    if args.key_hex:
        key = bytes.fromhex(args.key_hex)
        if len(key) != 32:
            print("Error: --key-hex must be exactly 64 hex chars (32 bytes)", file=sys.stderr)
            sys.exit(1)
    elif args.key_alternate:
        key = KEY_ALTERNATE
    else:
        key = KEY_DEFAULT

    inpath = args.input

    if os.path.isdir(inpath):
        files = sorted(f for f in os.listdir(inpath) if f.endswith('.unity3d'))
        if not files:
            print(f"No .unity3d files found in '{inpath}'")
            return

        if args.output:
            outdir = args.output
        elif args.in_place:
            outdir = os.path.normpath(os.path.join(inpath, '..', '..', 'StandaloneWindows64'))
        else:
            outdir = os.path.normpath(os.path.join(inpath, '..', 'reencrypted'))

        if not args.dry_run:
            os.makedirs(outdir, exist_ok=True)

        ok = errs = 0
        for fname in files:
            fpath = os.path.join(inpath, fname)
            base = fname.removesuffix('.unity3d')
            outpath = os.path.join(outdir, base + '.bundle')
            if process_file(fpath, outpath, key, args.dry_run):
                ok += 1
            else:
                errs += 1

        total = len(files)
        print(f"\n{ok} OK, {errs} errors of {total}")
        if not args.dry_run:
            print(f"Output directory: {outdir}/")

    else:
        if not os.path.exists(inpath):
            print(f"Error: '{inpath}' not found", file=sys.stderr)
            sys.exit(1)

        if args.output:
            outpath = args.output
        elif args.in_place:
            base = os.path.splitext(os.path.basename(inpath))[0]
            bundle_name = base + '.bundle'
            candidates = [
                os.path.join(os.path.dirname(inpath), bundle_name),
                os.path.normpath(os.path.join(
                    os.path.dirname(inpath), '..', '..', 'StandaloneWindows64', bundle_name)),
            ]
            outpath = None
            for c in candidates:
                if os.path.exists(c):
                    outpath = c
                    break
            if not outpath:
                print(f"Error: --in-place but no existing .bundle found for '{inpath}'",
                      file=sys.stderr)
                sys.exit(1)
        else:
            basedir = os.path.dirname(inpath) or '.'
            basename = os.path.splitext(os.path.basename(inpath))[0] + '.bundle'
            outpath = os.path.join(basedir, basename)

        process_file(inpath, outpath, key, args.dry_run)


if __name__ == '__main__':
    main()