import z3
import struct

N = 624
L = N

with open("init_state.txt", "r") as f:
    init_mt = [int(x.strip()) for x in f]

def crack():
    s = z3.Solver()

    init_key = [z3.BitVec(f"init_key_{i}", 32) for i in range(L)]
    mt = [z3.BitVecVal(x, 32) for x in init_mt]

    i = 1
    j = 0
    k = N
    while k > 0:
        mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> 30)) * z3.BitVecVal(1664525, 32))) + init_key[j] + z3.BitVecVal(j, 32)
        i += 1
        j += 1
        if i >= N:
            mt[0] = mt[N-1]
            i = 1
        if j >= L:
            j = 0

        k -= 1

    k = N - 1
    while k > 0:
        mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> 30)) * z3.BitVecVal(1566083941, 32))) - z3.BitVecVal(i, 32)
        i += 1
        if i >= N:
            mt[0] = mt[N-1]
            i = 1

        k -= 1

    for idx in range(1, N):
        s.add(mt[idx] == 0)

    s.check()
    m = s.model()

    x = bytearray()
    for ik in init_key:
        ikv = m[ik].as_long()
        x += struct.pack("<I", ikv)

    print(hex(int.from_bytes(x, byteorder='little')))

if __name__ == "__main__":
    crack()
