# "tkey-gen.py", licensed under the MIT license
# Copyright 2024 NinjaCheetah

import binascii
import hashlib
import libWiiPy
from libWiiPy.types import _ContentRecord


def _secret(start, length):
    ret = b''
    add = start + length
    for _ in range(length):
        unsigned_start = start & 0xFF  # Compensates for how Python handles negative values vs PHP.
        ret += bytes.fromhex(f"{unsigned_start:02x}"[-2:])
        nxt = start + add
        add = start
        start = nxt
    return ret


def _mungetid(tid):
    # Remove leading zeroes from the TID.
    while tid.startswith("00"):
        tid = tid[2:]
    if tid == "":
        tid = "00"
    # In PHP, the last character just gets dropped if you make a hex string from an odd-length input, so this
    # replicates that functionality.
    if len(tid) % 2 != 0:
        tid = tid[:-1]
    return bytes.fromhex(tid)


def _derive_key(tid, passwd):
    key_secret = _secret(-3, 10)
    salt = hashlib.md5(key_secret + _mungetid(tid)).digest()
    # Had to reduce the length here from 32 to 16 when converting to get the same length keys.
    return hashlib.pbkdf2_hmac("sha1", passwd.encode(), salt, 20, 16).hex()


def find_tkey(tid: str, banner_enc: bytes, content_record: _ContentRecord) -> bytes:
    # Find a working Title Key by generating a key with a password, then decrypting content 0 and comparing it to the
    # expected hash. If the hash matches, then we generated the correct key.
    passwds = ["nintendo", "mypass"]
    for passwd in passwds:
        key = binascii.unhexlify(_derive_key(tid, passwd).encode())
        banner_dec = libWiiPy.title.decrypt_content(banner_enc, key, content_record.index, content_record.content_size)
        banner_dec_hash = hashlib.sha1(banner_dec).hexdigest()
        content_record_hash = content_record.content_hash.decode()
        if banner_dec_hash == content_record_hash:
            return key
    raise Exception("Valid Title Key could not be generated")