#!/usr/bin/env python3
import os, sys, struct, json, time
from Crypto.Cipher import AES

KEY = b"rK7CcATuZk7LAmhqqU4iBLNmAq8QbK3s"
BUNDLES_DIR = "HolosHanafuda_Data/StreamingAssets/aa/StandaloneWindows64"
OUT = "extracted"

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

def extract_audio(env, dest_dir):
    audio_dirs = {}
    count = 0
    files = getattr(env.file, 'files', {})
    if not files:
        return count
    for obj in env.objects:
        if obj.type.name != 'AudioClip':
            continue
        try:
            data = obj.read()
            res = data.m_Resource
            if not res or res.m_Size == 0:
                continue
            src = res.m_Source.replace('archive:/', '')
            parts = src.split('/')
            res_name = parts[-1] if len(parts) > 1 else parts[0]
            if res_name not in files:
                continue
            reader = files[res_name]
            reader.Position = res.m_Offset
            raw = reader.read(res.m_Size)
            if not raw:
                continue
            # Detect format
            ext = '.fsb'
            if raw[:4] == b'OggS':
                ext = '.ogg'
            elif raw[:4] == b'RIFF':
                ext = '.wav'
            elif raw[:3] == b'FSB' or raw[:4] == b'FSB5':
                ext = '.fsb'
            dest = os.path.join(dest_dir, data.m_Name + ext)
            with open(dest, 'wb') as f:
                f.write(raw)
            count += 1
        except Exception as e:
            pass
    return count

def extract_texture(obj, dest_dir):
    try:
        data = obj.read()
        tname = data.m_Name
        if not tname:
            return None
        data.image.save(os.path.join(dest_dir, tname + '.png'))
        return {'name': tname, 'type': 'Texture2D'}
    except:
        return None

def extract_sprite(obj, dest_dir):
    try:
        data = obj.read()
        tname = data.m_Name
        if not tname:
            return None
        data.image.save(os.path.join(dest_dir, tname + '.png'))
        return {'name': tname, 'type': 'Sprite'}
    except:
        return None

def extract_mesh(obj, dest_dir):
    try:
        data = obj.read()
        tname = data.m_Name
        if not tname:
            return None
        with open(os.path.join(dest_dir, tname + '.obj'), 'w') as f:
            f.write(data.export())
        return {'name': tname, 'type': 'Mesh'}
    except:
        return None

def main():
    os.makedirs(f"{OUT}/bundles", exist_ok=True)
    os.makedirs(f"{OUT}/assets", exist_ok=True)
    os.makedirs(f"{OUT}/audio", exist_ok=True)

    files = sorted(f for f in os.listdir(BUNDLES_DIR) if f.endswith('.bundle'))
    print(f"Found {len(files)} encrypted bundles\n")

    # decrypt all
    print("Decrypting bundles...")
    ok = errs = 0
    for idx, fname in enumerate(files, 1):
        fpath = os.path.join(BUNDLES_DIR, fname)
        try:
            with open(fpath, 'rb') as f:
                enc = f.read()
            dec = decrypt(enc)
        except Exception as e:
            print(f"  [{idx}/{len(files)}] ERROR {fname}: {e}")
            errs += 1
            continue
        base = fname.removesuffix('.bundle')
        with open(f"{OUT}/bundles/{base}.unity3d", 'wb') as f:
            f.write(dec)
        ok += 1
        if ok % 100 == 0 or ok == len(files):
            print(f"  [{idx}/{len(files)}] {fname}")
    print(f"\nDecryption: {ok} OK, {errs} errors\n")

    # extract everything with UnityPy
    print("Extracting assets...")
    try:
        import UnityPy
    except ImportError:
        print("  Missing UnityPy - pip install UnityPy")
        return

    decrypted = sorted(f for f in os.listdir(f"{OUT}/bundles") if f.endswith('.unity3d'))
    type_counts = {}
    asset_index = []
    total_audio = 0
    total_textures = 0
    total_sprites = 0
    total_meshes = 0

    for idx, fname in enumerate(decrypted, 1):
        fpath = os.path.join(f"{OUT}/bundles", fname)
        try:
            env = UnityPy.load(fpath)
        except Exception as e:
            print(f"  [{idx}/{len(decrypted)}] PARSE ERROR {fname}: {e}")
            continue

        # Extract audio from .resource blobs
        audio_count = extract_audio(env, f"{OUT}/audio")
        total_audio += audio_count

        # Extract all container objects
        for path, obj in env.container.items():
            try:
                ttype = obj.type.name
                type_counts[ttype] = type_counts.get(ttype, 0) + 1

                if ttype == 'Texture2D':
                    r = extract_texture(obj, f"{OUT}/assets")
                    if r:
                        total_textures += 1
                        r['bundle'] = fname
                        asset_index.append(r)
                elif ttype == 'Sprite':
                    r = extract_sprite(obj, f"{OUT}/assets")
                    if r:
                        total_sprites += 1
                        r['bundle'] = fname
                        asset_index.append(r)
                elif ttype == 'Mesh':
                    r = extract_mesh(obj, f"{OUT}/assets")
                    if r:
                        total_meshes += 1
                        r['bundle'] = fname
                        asset_index.append(r)
                elif ttype in ('MonoBehaviour', 'TextAsset', 'Shader', 'Material',
                               'AnimatorController', 'VideoClip'):
                    data = obj.read()
                    tname = data.m_Name if hasattr(data, 'm_Name') else path
                    asset_index.append({'name': tname, 'type': ttype, 'bundle': fname})
                elif ttype == 'AudioClip':
                    asset_index.append({'name': ttype, 'type': ttype, 'bundle': fname})
            except Exception as e:
                if idx <= 3:
                    print(f"    SKIP {obj.type.name}: {e}")

        if idx % 200 == 0:
            print(f"  [{idx}/{len(decrypted)}]")
    print(f"  [{len(decrypted)}/{len(decrypted)}] done")

    print(f"\nExtraction Summary")
    print(f"Bundles decrypted: {ok}")
    print(f"Audio clips:        {total_audio} (.fsb/.ogg/.wav)")
    print(f"Textures extracted: {total_textures}")
    print(f"Sprites extracted:  {total_sprites}")
    print(f"Meshes extracted:   {total_meshes}")
    print(f"Asset types found:")
    for t, c in sorted(type_counts.items(), key=lambda x: -x[1]):
        print(f"  {t}: {c}")

    with open(f"{OUT}/asset_index.json", 'w') as f:
        json.dump({'bundles': ok, 'audio': total_audio,
                    'textures': total_textures,
                    'types': type_counts,
                    'assets': asset_index}, f, indent=2)
    print(f"\nIndex:     {OUT}/asset_index.json")
    print(f"Bundles:   {OUT}/bundles/")
    print(f"Images:    {OUT}/assets/")
    print(f"Audio:     {OUT}/audio/")
    print(f"This script extracts the files as a example, for a reliable extraction use something like AssetRipper")

if __name__ == '__main__':
    t0 = time.time()
    main()
    print(f"\nTotal: {time.time()-t0:.1f}s")