import ida_bytes
import ida_kernwin

MASK64 = 0xFFFFFFFFFFFFFFFF
SEED_ADDR = 0x14000E100
SEED_LEN = 120


def m(x: int) -> int:
    return x & MASK64


def rol64(x: int, r: int) -> int:
    r &= 63
    x = m(x)
    return m((x << r) | (x >> (64 - r)))


def ror64(x: int, r: int) -> int:
    r &= 63
    x = m(x)
    return m((x >> r) | (x << (64 - r)))


def mix_byteswap(x: int) -> int:
    # Matches the repeated pattern in sub_1400016C0
    x = m(x)
    x = m(((x << 8) & 0xFF00FF00FF00FF00) ^ x)
    x = m((((x >> 8) & 0x00FF00FF00FF00FF) ^ x))
    return x


def fmix(x: int) -> int:
    x = mix_byteswap(x)
    y = m(0x9E3779B97F4A7C15 * m(ror64(x, 7) ^ rol64(x, 13) ^ x))
    y = m(0xC2B2AE3D27D4EB4F * m((y >> 29) ^ y))
    return m((y >> 32) ^ y)


def load_seed_qwords() -> list[int]:
    blob = ida_bytes.get_bytes(SEED_ADDR, SEED_LEN)
    if not blob or len(blob) != SEED_LEN:
        raise RuntimeError(f"Failed to read seed blob at {SEED_ADDR:#x}")

    # sub_14000AAC0: XOR with 0xA5, pack into 15 little-endian qwords
    qwords: list[int] = []
    for i in range(15):
        v = 0
        for j in range(8):
            v |= (blob[i * 8 + j] ^ 0xA5) << (8 * j)
        qwords.append(m(v))
    return qwords


def derive_key_qwords(q: list[int]) -> tuple[int, int, int, int]:
    # Faithful port of sub_1400016C0 with correct 64-bit overflow behavior
    v0_idx = 0
    v1 = 0
    v2 = 0
    v3 = 0xCBF29CE484222325
    n15 = 0

    while True:
        v5 = n15
        n15 += 1

        a = m(v2 + q[(v5 + 7) % 0xF])
        b = m(q[n15 % 0xF] ^ rol64(q[v0_idx], 5 * (v5 & 0xFF)))
        v6 = (a * b) & ((1 << 128) - 1)

        hi = (v6 >> 64) & MASK64
        lo = v6 & MASK64

        v7 = ror64(hi, 9)
        hi = rol64(hi, 17)
        hi = m(0xFF51AFD7ED558CCD * m(((v7 ^ hi ^ lo) >> 33) ^ v7 ^ hi ^ lo))

        t = m(v3 ^ v1 ^ hi ^ (hi >> 33))
        mixv = mix_byteswap(t)

        hi = rol64(mixv, 13)
        v8 = ror64(mixv, 7)

        v0_idx += 1
        v2 = m(v2 - 0x61C8864680B583EB)
        v1 = m(v1 + 0x100000001B3)

        t2 = m(0x9E3779B97F4A7C15 * m(v8 ^ hi ^ mixv))
        t2 = m(0xC2B2AE3D27D4EB4F * m((t2 >> 29) ^ t2))
        v9 = m((t2 >> 32) ^ t2)
        v3 = v9

        if n15 == 15:
            break

    v10 = q[0]
    v78 = 0
    v11 = m(v9 ^ q[0])

    v12t = m(v9 ^ q[11])
    v12 = m(v12t ^ ((v12t << 8) & 0xFF00FF00FF00FF00))
    v12b = m(((v12 >> 8) & 0x00FF00FF00FF00FF) ^ v12)

    v13 = m(0x9E3779B97F4A7C15 * m(ror64(v12b, 7) ^ rol64(v12b, 13) ^ v12b))
    v13m = m(0xC2B2AE3D27D4EB4F * m((v13 >> 29) ^ v13))
    v14 = m(q[1] + m((v13m >> 32) ^ v13m))

    # These are full 128-bit products in the original code
    p15 = v3 * q[12]
    v15 = (p15 >> 64) & MASK64
    v16 = ror64(v15, 9)
    v17 = rol64(v15, 17)
    prod12 = m(v3 * q[12])
    v18 = m(0xFF51AFD7ED558CCD * m(((prod12 ^ v16 ^ v17) >> 33) ^ prod12 ^ v16 ^ v17))
    v19 = m(q[2] ^ v18 ^ (v18 >> 33))

    p20 = q[14] * m(v3 ^ q[13])
    v20 = rol64((p20 >> 64) & MASK64, 17)
    v21 = ror64((p20 >> 64) & MASK64, 9)

    n600000 = 0
    prod13 = m(q[14] * m(v3 ^ q[13]))
    v23 = m(0xFF51AFD7ED558CCD * m(((prod13 ^ v21 ^ v20) >> 33) ^ prod13 ^ v21 ^ v20))
    v24 = m(q[3] + m((v23 >> 33) ^ v23))

    v79 = 0
    v80 = 0

    while True:
        v25 = mix_byteswap(m(v3 ^ v78))
        v26 = m(ror64(v25, 7) ^ rol64(v25, 13))

        tmp = m(0x9E3779B97F4A7C15 * m(v26 ^ v25))
        v27 = m(0xC2B2AE3D27D4EB4F * m((tmp >> 29) ^ tmp))
        v28 = m((v27 >> 32) ^ v27)
        v29 = m(v28 ^ v11 ^ rol64(v14, 11))

        left = m(v10 ^ m(v28 + v19 + ror64(v24, 17)))
        right = m(v29 + q[(n600000 & 3) + 8])
        v30 = (left * right) & ((1 << 128) - 1)
        v31 = (v30 >> 64) & MASK64

        v30hi = ror64(v31, 9)
        v32 = rol64(v31, 17)
        v30lo = v30 & MASK64
        v33 = m(0xFF51AFD7ED558CCD * m(((v30lo ^ v30hi ^ v32) >> 33) ^ v30lo ^ v30hi ^ v32))
        v34 = m((v33 >> 33) ^ v33)

        left = m(v29 ^ m(v24 + q[2]))
        right = m(v34 ^ (v14 | 1))
        v35 = (left * right) & ((1 << 128) - 1)
        v36 = m(v11 + rol64(m(v3 + (v14 ^ v34)), 23))

        v35hi = (v35 >> 64) & MASK64
        v37 = ror64(v35hi, 9)
        v35hir = rol64(v35hi, 17)
        v38 = m(v19 + rol64(m(v3 + v34 + (v24 ^ v36)), 31))

        v35lo = v35 & MASK64
        v35mix = m(0xFF51AFD7ED558CCD * m(((v35lo ^ v37 ^ v35hir) >> 33) ^ v35lo ^ v37 ^ v35hir))
        v39 = m(v14 ^ ror64(m(m((v35mix >> 33) ^ v35mix) + v19 + v3), 19))
        v40 = m((m(v39 + v38) + q[10] + v3) ^ v24)

        if (n600000 & 0x7FFF) == 0:
            mul = (m(v39 + v40) * m(v36 + v38)) & ((1 << 128) - 1)
            hi = (mul >> 64) & MASK64
            lo = mul & MASK64
            v42 = m(lo ^ ror64(hi, 9) ^ rol64(hi, 17))
            m1 = m(0xFF51AFD7ED558CCD * m(v42 ^ (v42 >> 33)))
            v43 = m(m1 ^ v3 ^ (m1 >> 33) ^ ((m(m1 ^ v3 ^ (m1 >> 33)) << 8) & 0xFF00FF00FF00FF00))
            v44 = m(((v43 >> 8) & 0x00FF00FF00FF00FF) ^ v43)
            v45 = rol64(v44, 13)
            v46 = ror64(v44, 7)
            tmp = m(0x9E3779B97F4A7C15 * m(v46 ^ v45 ^ v44))
            v47 = m(0xC2B2AE3D27D4EB4F * m((tmp >> 29) ^ tmp))
            v3 = m((v47 >> 32) ^ v47)

        v48 = m(v3 + 7 * v40 + 3 * v39 + v36 + 5 * v38)
        v49 = m(v3 + 19 * v40 + 13 * v39 + 11 * v36 + 17 * v38)
        v77 = m(53 * v40 + 43 * v39 + 41 * v36 + 47 * v38 + v3)

        v79 = m(v48 ^ rol64(m(v3 ^ v77), 9)) % q[4]
        v11 = v79
        v80 = m(v49 + ror64(m(v48 ^ v3), 7)) % q[5]
        v14 = v80

        v50 = m(v3 ^ v48 ^ m(v3 + 37 * v40 + 23 * v36 + 31 * v38 + 29 * v39))
        v51 = mix_byteswap(v50)
        v52 = rol64(v51, 13)
        v53 = ror64(v51, 7)
        tmp = m(0x9E3779B97F4A7C15 * m(v53 ^ v52 ^ v51))
        v54 = m(0xC2B2AE3D27D4EB4F * m((tmp >> 29) ^ tmp))

        v55 = mix_byteswap(m(v3 + v77 + v49))
        v19 = m((v54 >> 32) ^ v54)

        v56 = rol64(v55, 13)
        v57 = ror64(v55, 7)

        n600000 += 1
        v78 = m(v78 + q[6])

        tmp = m(0x9E3779B97F4A7C15 * m(v57 ^ v56 ^ v55))
        v58 = m(0xC2B2AE3D27D4EB4F * m((tmp >> 29) ^ tmp))
        v24 = m((v58 >> 32) ^ v58)

        if n600000 == 600000:
            break

        idx = (n600000 >> 3) - 15 * int(((2290649225 * (n600000 >> 3)) >> 32) >> 3)
        v10 = q[idx]

    k0 = fmix(m(v3 ^ v79))
    k1 = fmix(m(v80 + v3))
    k2 = fmix(m(q[9] ^ v19))
    k3 = fmix(m(v24 + q[10]))
    return k0, k1, k2, k3


def main() -> None:
    qwords = load_seed_qwords()
    k0, k1, k2, k3 = derive_key_qwords(qwords)

    key_bytes = b"".join(k.to_bytes(8, "little") for k in (k0, k1, k2, k3))
    key_hex = key_bytes.hex()

    print("[+] Encryption key (hex):", key_hex)
    print("[+] As qwords:")
    print(f"    k0 = 0x{k0:016x}")
    print(f"    k1 = 0x{k1:016x}")
    print(f"    k2 = 0x{k2:016x}")
    print(f"    k3 = 0x{k3:016x}")

    ida_kernwin.info(f"Recovered encryption key:\n{key_hex}")


if __name__ == "__main__":
    main()
