import argparse
import hashlib
import shutil
from pathlib import Path

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

def build_salt(bundle_path: Path) -> bytes:
    name = bundle_path.stem.lower()
    return name[:16].ljust(16, "_").encode("utf-8")

def derive_key(bundle_path: Path, password: bytes) -> bytes:
    return hashlib.pbkdf2_hmac(
        "sha1",
        password,
        build_salt(bundle_path),
        1_000,
        dklen=16,
    )

def transform_prefix(prefix: bytearray, key: bytes) -> None:
    encryptor = Cipher(algorithms.AES(key), modes.ECB()).encryptor()

    for offset in range(0, len(prefix), 16):
        counter = offset // 16 + 1
        counter_block = counter.to_bytes(8, "little") + bytes(8)
        key_stream = encryptor.update(counter_block)
        block = prefix[offset : offset + 16]
        prefix[offset : offset + len(block)] = bytes(
            value ^ key_stream[index] for index, value in enumerate(block)
        )

    encryptor.finalize()

def validate_unityfs(header: bytes, file_size: int) -> None:
    if not header.startswith(b"UnityFS\0"):
        raise ValueError("decryption did not produce a UnityFS header")

    position = 12
    try:
        for _ in range(2):
            position = header.index(0, position) + 1
    except ValueError as error:
        raise ValueError("the UnityFS header is incomplete") from error

    if position + 8 > len(header):
        raise ValueError("the UnityFS header does not contain a file size")

    declared_size = int.from_bytes(header[position : position + 8], "big")
    if declared_size != file_size:
        raise ValueError(
            f"UnityFS size mismatch: header={declared_size}, file={file_size}"
        )

def decrypt_bundle(
    source: Path,
    destination: Path,
    password: bytes,
    overwrite: bool,
) -> None:
    if source.resolve() == destination.resolve():
        raise ValueError("input and output paths must be different")

    file_size = source.stat().st_size

    with source.open("rb") as input_file:
        prefix = bytearray(input_file.read(128))
        transform_prefix(prefix, derive_key(source, password))
        validate_unityfs(prefix, file_size)

        destination.parent.mkdir(parents=True, exist_ok=True)
        with destination.open("wb" if overwrite else "xb") as output_file:
            output_file.write(prefix)
            shutil.copyfileobj(input_file, output_file)

def decrypt_path(
    source: Path,
    destination: Path,
    password: bytes,
    recursive: bool,
    overwrite: bool,
) -> None:
    if source.is_file():
        output_file = (
            destination / source.name if destination.is_dir() else destination
        )
        decrypt_bundle(source, output_file, password, overwrite)
        print(output_file)
        return

    if not source.is_dir():
        raise FileNotFoundError(f"input path does not exist: {source}")

    bundles = sorted(
        source.rglob("*.bundle") if recursive else source.glob("*.bundle")
    )
    if not bundles:
        raise FileNotFoundError(f"no bundle files found in: {source}")

    for bundle in bundles:
        output_file = destination / bundle.relative_to(source)
        decrypt_bundle(bundle, output_file, password, overwrite)
        print(output_file)

def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Decrypt Etrange Overlord Demo Unity asset bundles."
    )
    parser.add_argument(
        "input",
        type=Path,
        help="encrypted bundle or directory",
    )
    parser.add_argument(
        "output",
        type=Path,
        help="output file or directory",
    )
    parser.add_argument(
        "-r",
        "--recursive",
        action="store_true",
        help="search subdirectories when the input is a directory",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="overwrite existing output files",
    )
    return parser

def main() -> None:
    parser = create_parser()
    arguments = parser.parse_args()

    try:
        decrypt_path(
            arguments.input,
            arguments.output,
            b"xKci6s8ZAkbfj85pWbR9BKFfZqKFG3Yr6Urtw9D6",
            arguments.recursive,
            arguments.force,
        )
    except (OSError, ValueError) as error:
        parser.exit(1, f"error: {error}\n")

if __name__ == "__main__":
    main()