import argparse
import hashlib
import os
import sys

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


def password_derive_bytes(password, salt):
    """Match the old .NET PasswordDeriveBytes defaults used by the game."""
    value = hashlib.sha1(password + salt).digest()
    for _ in range(98):
        value = hashlib.sha1(value).digest()
    return hashlib.sha1(value).digest()[:16]


def transform(data, key, file_offset=0):
    first_block = file_offset // 16
    key_position = file_offset % 16
    block_count = (key_position + len(data) + 15) // 16
    counters = bytearray(block_count * 16)

    for block in range(block_count):
        counter = first_block + block + 1
        start = block * 16
        counters[start:start + 8] = counter.to_bytes(8, "little")

    encryptor = Cipher(algorithms.AES(key), modes.ECB()).encryptor()
    key_stream = encryptor.update(bytes(counters)) + encryptor.finalize()
    return bytes(
        value ^ key_stream[key_position + i]
        for i, value in enumerate(data)
    )


def read_unityfs_size(header):
    if not header.startswith(b"UnityFS\0"):
        return None

    pos = 12
    for _ in range(2):
        end = header.find(b"\0", pos)
        if end < 0:
            return None
        pos = end + 1

    if pos + 8 > len(header):
        return None

    return int.from_bytes(header[pos:pos + 8], "big")


def decrypt_bundle(source_name, output_name, password):
    if os.path.normcase(os.path.abspath(source_name)) == os.path.normcase(
        os.path.abspath(output_name)
    ):
        raise RuntimeError("input and output are the same file")

    salt = os.path.splitext(os.path.basename(source_name))[0]
    key = password_derive_bytes(password, salt.encode("utf-8"))

    with open(source_name, "rb") as source:
        header = transform(source.read(128), key)

        if read_unityfs_size(header) != os.path.getsize(source_name):
            marker = header[:8].hex(" ")
            raise RuntimeError(f"not a matching bundle ({marker})")

        os.makedirs(
            os.path.dirname(os.path.abspath(output_name)),
            exist_ok=True,
        )

        source.seek(0)
        position = 0
        with open(output_name, "wb") as output:
            while True:
                chunk = source.read(1024 * 1024)
                if not chunk:
                    break
                output.write(transform(chunk, key, position))
                position += len(chunk)


def cached_files(root):
    for current, _, names in os.walk(root):
        for name in sorted(names):
            yield os.path.join(current, name)


def main():
    password = b"M3kR9/aq9W"

    parser = argparse.ArgumentParser(
        description="Decrypt Wizardry Variants Daphne asset-cache bundles."
    )
    parser.add_argument("input", help="one cache file or the assetcache directory")
    parser.add_argument("output", help="output file or directory")
    args = parser.parse_args()

    source = os.path.abspath(args.input)
    if not os.path.exists(source):
        parser.error("input was not found")

    if os.path.isfile(source):
        files = [source]
        output_root = None
    else:
        files = list(cached_files(source))
        output_root = os.path.abspath(args.output)

    passed = 0
    failed = 0

    for filename in files:
        try:
            if output_root:
                relative = os.path.relpath(filename, source) + ".bundle"
                output_name = os.path.join(output_root, relative)
            else:
                output_name = os.path.abspath(args.output)

            decrypt_bundle(filename, output_name, password)

            passed += 1
            print(output_name)
        except (OSError, RuntimeError) as error:
            failed += 1
            print(f"Failed: {filename}: {error}", file=sys.stderr)

    print(f"\n{passed} bundle(s) decrypted; {failed} failed")
    return 1 if failed else 0


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