"""
HASH58 implementation for iPod Classic (all gens), Nano 3G, and Nano 4G.
Ported from libgpod's itdb_hash58.c.

This is the complete, working implementation for signing iTunesDB files
for devices that use hashing_scheme=1 (HASH58).

Usage:
    from hash58 import write_hash58

    with open("iTunesDB", "rb") as f:
        itdb_data = bytearray(f.read())

    firewire_id = bytes.fromhex("0011223344556677")  # From SysInfo
    write_hash58(itdb_data, firewire_id)

    with open("iTunesDB", "wb") as f:
        f.write(itdb_data)
"""

import hashlib
from math import gcd

# AES S-Box (from libgpod itdb_hash58.c lines 45-76)
TABLE1 = bytes([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5,
    0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0,
    0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC,
    0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A,
    0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0,
    0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B,
    0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85,
    0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5,
    0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17,
    0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88,
    0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C,
    0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9,
    0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6,
    0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E,
    0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94,
    0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68,
    0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
])

# AES Inverse S-Box (from libgpod itdb_hash58.c lines 78-115)
TABLE2 = bytes([
    0x52, 0x09, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38,
    0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB,
    0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87,
    0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB,
    0x54, 0x7B, 0x94, 0x32, 0xA6, 0xC2, 0x23, 0x3D,
    0xEE, 0x4C, 0x95, 0x0B, 0x42, 0xFA, 0xC3, 0x4E,
    0x08, 0x2E, 0xA1, 0x66, 0x28, 0xD9, 0x24, 0xB2,
    0x76, 0x5B, 0xA2, 0x49, 0x6D, 0x8B, 0xD1, 0x25,
    0x72, 0xF8, 0xF6, 0x64, 0x86, 0x68, 0x98, 0x16,
    0xD4, 0xA4, 0x5C, 0xCC, 0x5D, 0x65, 0xB6, 0x92,
    0x6C, 0x70, 0x48, 0x50, 0xFD, 0xED, 0xB9, 0xDA,
    0x5E, 0x15, 0x46, 0x57, 0xA7, 0x8D, 0x9D, 0x84,
    0x90, 0xD8, 0xAB, 0x00, 0x8C, 0xBC, 0xD3, 0x0A,
    0xF7, 0xE4, 0x58, 0x05, 0xB8, 0xB3, 0x45, 0x06,
    0xD0, 0x2C, 0x1E, 0x8F, 0xCA, 0x3F, 0x0F, 0x02,
    0xC1, 0xAF, 0xBD, 0x03, 0x01, 0x13, 0x8A, 0x6B,
    0x3A, 0x91, 0x11, 0x41, 0x4F, 0x67, 0xDC, 0xEA,
    0x97, 0xF2, 0xCF, 0xCE, 0xF0, 0xB4, 0xE6, 0x73,
    0x96, 0xAC, 0x74, 0x22, 0xE7, 0xAD, 0x35, 0x85,
    0xE2, 0xF9, 0x37, 0xE8, 0x1C, 0x75, 0xDF, 0x6E,
    0x47, 0xF1, 0x1A, 0x71, 0x1D, 0x29, 0xC5, 0x89,
    0x6F, 0xB7, 0x62, 0x0E, 0xAA, 0x18, 0xBE, 0x1B,
    0xFC, 0x56, 0x3E, 0x4B, 0xC6, 0xD2, 0x79, 0x20,
    0x9A, 0xDB, 0xC0, 0xFE, 0x78, 0xCD, 0x5A, 0xF4,
    0x1F, 0xDD, 0xA8, 0x33, 0x88, 0x07, 0xC7, 0x31,
    0xB1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xEC, 0x5F,
    0x60, 0x51, 0x7F, 0xA9, 0x19, 0xB5, 0x4A, 0x0D,
    0x2D, 0xE5, 0x7A, 0x9F, 0x93, 0xC9, 0x9C, 0xEF,
    0xA0, 0xE0, 0x3B, 0x4D, 0xAE, 0x2A, 0xF5, 0xB0,
    0xC8, 0xEB, 0xBB, 0x3C, 0x83, 0x53, 0x99, 0x61,
    0x17, 0x2B, 0x04, 0x7E, 0xBA, 0x77, 0xD6, 0x26,
    0xE1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0C, 0x7D
])

# Fixed bytes for key derivation (from libgpod itdb_hash58.c lines 113-115)
FIXED = bytes([
    0x67, 0x23, 0xFE, 0x30, 0x45, 0x33, 0xF8, 0x90, 0x99,
    0x21, 0x07, 0xC1, 0xD0, 0x12, 0xB2, 0xA1, 0x07, 0x81
])

# Hash scheme identifier for HASH58
ITDB_CHECKSUM_HASH58 = 1

# Header offsets
OFFSET_DB_ID = 0x18       # 8 bytes
OFFSET_UNK_0x32 = 0x32    # 20 bytes
OFFSET_HASHING_SCHEME = 0x30  # 2 bytes (hashing_scheme in MhbdHeader struct)
OFFSET_HASH58 = 0x58      # 20 bytes


def _lcm(a: int, b: int) -> int:
    """Least common multiple."""
    if a == 0 or b == 0:
        return 1
    return (a * b) // gcd(a, b)


def _generate_key(firewire_id: bytes) -> bytes:
    """
    Generate a 64-byte HMAC key from the 8-byte FireWire ID.

    Algorithm:
    1. Take LCM of each pair of consecutive bytes in the FireWire ID
    2. Use the high and low bytes of each LCM to index into TABLE1 and TABLE2
    3. SHA1 hash the FIXED bytes + derived bytes
    4. Pad to 64 bytes
    """
    if len(firewire_id) < 8:
        raise ValueError(f"FireWire ID must be at least 8 bytes, got {len(firewire_id)}")

    y = bytearray(16)
    for i in range(4):
        a = firewire_id[i * 2]
        b = firewire_id[i * 2 + 1]
        cur_lcm = _lcm(a, b)

        hi = (cur_lcm >> 8) & 0xFF
        lo = cur_lcm & 0xFF

        y[i * 4] = TABLE1[hi]
        y[i * 4 + 1] = TABLE2[hi]
        y[i * 4 + 2] = TABLE1[lo]
        y[i * 4 + 3] = TABLE2[lo]

    # SHA1(FIXED + y), then pad to 64 bytes
    h = hashlib.sha1(FIXED + y).digest()
    key = bytearray(64)
    key[:len(h)] = h
    return bytes(key)


def compute_hash58(firewire_id: bytes, itdb_data: bytes) -> bytes:
    """
    Compute HMAC-SHA1 of iTunesDB data using FireWire ID-derived key.

    This is a standard HMAC-SHA1 implementation with a custom key derivation.

    Args:
        firewire_id: 8-20 byte FireWire GUID from SysInfo
        itdb_data: Complete iTunesDB file contents (with hash fields zeroed)

    Returns:
        20-byte SHA1 hash to write at offset 0x58
    """
    key = _generate_key(firewire_id)

    # HMAC-SHA1: H(K XOR opad, H(K XOR ipad, message))
    # Inner hash
    inner_key = bytes(b ^ 0x36 for b in key)
    inner_hash = hashlib.sha1(inner_key + itdb_data).digest()

    # Outer hash
    outer_key = bytes(b ^ 0x5c for b in key)
    return hashlib.sha1(outer_key + inner_hash).digest()


def write_hash58(itdb_data: bytearray, firewire_id: bytes) -> None:
    """
    Compute and write HASH58 checksum to iTunesDB data in-place.

    From libgpod itdb_hash58_write_hash():
    1. Backs up db_id and unk_0x32 fields
    2. Zeros out db_id, unk_0x32, and hash58 fields (required for hash computation)
    3. Sets hashing_scheme to HASH58 (1)
    4. Computes HMAC-SHA1 hash
    5. Writes hash to offset 0x58
    6. Restores backed up fields

    NOTE: Unlike HASH72, HASH58 zeros unk_0x32 before computing the hash!
    This is a key difference between the two algorithms.

    Args:
        itdb_data: Mutable bytearray of complete iTunesDB file
        firewire_id: 8-20 byte FireWire GUID from /iPod_Control/Device/SysInfo

    Raises:
        ValueError: If iTunesDB is too small or FireWire ID is invalid
    """
    if len(itdb_data) < 0x6C:
        raise ValueError(f"iTunesDB file too small ({len(itdb_data)} bytes), need at least 0x6C")

    # Verify this is an mhbd header
    if itdb_data[:4] != b'mhbd':
        raise ValueError("Invalid iTunesDB: expected 'mhbd' header")

    # Backup fields that will be zeroed
    backup_db_id = bytes(itdb_data[OFFSET_DB_ID:OFFSET_DB_ID + 8])
    backup_unk32 = bytes(itdb_data[OFFSET_UNK_0x32:OFFSET_UNK_0x32 + 20])

    # Zero out fields for hash computation
    itdb_data[OFFSET_DB_ID:OFFSET_DB_ID + 8] = b'\x00' * 8
    itdb_data[OFFSET_UNK_0x32:OFFSET_UNK_0x32 + 20] = b'\x00' * 20
    itdb_data[OFFSET_HASH58:OFFSET_HASH58 + 20] = b'\x00' * 20

    # Set hashing scheme to HASH58
    itdb_data[OFFSET_HASHING_SCHEME:OFFSET_HASHING_SCHEME + 2] = \
        ITDB_CHECKSUM_HASH58.to_bytes(2, 'little')

    # Compute and write hash
    hash_val = compute_hash58(firewire_id, bytes(itdb_data))
    if len(hash_val) != 20:
        raise RuntimeError(f"Hash computation failed: expected 20 bytes, got {len(hash_val)}")
    itdb_data[OFFSET_HASH58:OFFSET_HASH58 + 20] = hash_val

    # Restore backed up fields
    itdb_data[OFFSET_DB_ID:OFFSET_DB_ID + 8] = backup_db_id
    itdb_data[OFFSET_UNK_0x32:OFFSET_UNK_0x32 + 20] = backup_unk32


def read_firewire_id(ipod_path: str) -> bytes:
    """Return the FireWire GUID for the connected iPod.

    Reads from the centralised DeviceInfo store.  Raises if not available.
    """
    from clickwheel.ipod.device_info import get_current_device
    device = get_current_device()
    if device is not None:
        fwid = device.firewire_id_bytes
        if fwid:
            return fwid
    raise RuntimeError(
        "FireWire GUID not available. Device info was not populated "
        "by the device scanner."
    )


if __name__ == "__main__":
    # Example usage
    import sys

    if len(sys.argv) < 3:
        print("Usage: python hash58.py <ipod_path> <itunesdb_path>")
        print("Example: python hash58.py E: E:/iPod_Control/iTunes/iTunesDB")
        sys.exit(1)

    ipod_path = sys.argv[1]
    itunesdb_path = sys.argv[2]

    try:
        firewire_id = read_firewire_id(ipod_path)
        print(f"FireWire ID: {firewire_id.hex()}")

        with open(itunesdb_path, 'rb') as f:
            itdb_data = bytearray(f.read())

        print(f"Read {len(itdb_data)} bytes from iTunesDB")

        write_hash58(itdb_data, firewire_id)
        print("Hash computed successfully!")

        # Write back (uncomment to actually write)
        # with open(itunesdb_path, 'wb') as f:
        #     f.write(itdb_data)
        # print("iTunesDB updated!")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
