
import random
import z3


def ROL(X, N):
    if isinstance(X, z3.BitVecRef):
        return z3.RotateLeft(X,N)
    else:
        return ((((X) << (N)) | ((X) >> (64 - (N)))) % (2 ** 64))


def ROR(X, N):
    if isinstance(X, z3.BitVecRef):
        return z3.RotateRight(X,N)
    else:
        return ((((X) >> (N)) | ((X) << (64 - (N)))) % (2 ** 64))

def custom_random(state):
    NEW_STATE = ROL(state,30) ^ ROR(state,12) ^ ROL(state,42) ^ ROL(state,4) ^ ROR(state,5)
    return NEW_STATE % 256, NEW_STATE

def get_seed_for_sequence(bytesequence):
    s = z3.Solver() 
    base_state = z3.BitVec("base_state", 64)
    cur_state = base_state
    s.add(base_state > 0)
    for byte in bytesequence:
        gen_byte, cur_state = custom_random(cur_state)
        s.add(gen_byte == byte)

    if s.check() != z3.sat:
        raise RuntimeError()

    m = s.model()
    return m[base_state].as_long()


def gen_sequence(seed, length):
    r = []
    cs = seed
    for _ in range(length):
        v, cs = custom_random(cs)
        r.append(v)
    return bytes(r)


test = random.getrandbits(32).to_bytes(4, "little")
seed = get_seed_for_sequence(test)
print(seed)
print(test)
assert test == gen_sequence(seed, 4)
