import argparse
from hashlib import blake2b
from pathlib import Path
from typing import cast

from UnityPy import Environment
from UnityPy.classes import TextAsset

# MU3.Sys.System.keyIvDigestFixed
KEY_IV_DIGEST_FIXED = bytes([
    179, 10, 98, 130, 17, 166, 184, 233, 246, 211,
    46, 229, 236, 79, 78, 83, 107, 151, 195, 172,
    57, 72, 120, 103, 17, 124, 18, 64, 15, 225,
    169, 39
])
FILE_MAP = {
    "noise0": "Key",
    "noise1": "IV",
    "noise2": "Endpoint salt",
}

parser = argparse.ArgumentParser()
_ = parser.add_argument("mu3_data_path", type=Path)
_ = parser.add_argument("key_iv_digest_fixed", type=bytes.fromhex, nargs="?", default=KEY_IV_DIGEST_FIXED)

args = parser.parse_args()

env = Environment()

with (cast(Path, args.mu3_data_path) / "resources.assets").open("rb") as f:
    env.load_file(f)

    for object in env.objects:
        if object.type.name != "TextAsset":
            continue

        data = cast(TextAsset, object.read())

        if data.name in FILE_MAP:
            key = blake2b(data.m_Script, digest_size=32, key=args.key_iv_digest_fixed).hexdigest()

            print(f"{FILE_MAP[data.name]}: {key}")