from z3 import Solver, BitVec, sat, ULE

known_prefix = "CSCG{"

def rng(x, size):
    return (x*A+B) & ((2**size)-1)

def gen_random(seed, bits, mask):
    state = seed
    while True:
        state = rng(state, bits)
        yield state & mask

A = BitVec("A", 56)
B = BitVec("B", 56)
SEED = BitVec("seed", 56)
BITS = 56 
MASK = 0xFF

state = SEED 

with open("./test_msg.txt", "r") as f:
    nums = [int(n) for n in f.readlines()[1:]]
    print(nums)

flag = []

solver = Solver()

for i, n in enumerate(nums):
    state = rng(state, BITS)
    cur_num = state & MASK
    decrypted_char = cur_num ^ n
    flag.append(decrypted_char)
    if i < len(known_prefix):
        solver.add(decrypted_char == ord(known_prefix[i]))
                     #    else:
                     #        solver.add(ULE(32, decrypted_char))
                     #        solver.add(ULE(decrypted_char, 126))
   
if solver.check() == sat:
    model = solver.model()
    recovered_chars = ""
    for char_expr in flag:
        val = model.evaluate(char_expr).as_long()
        recovered_chars += chr(val)
    print(model)
    print(recovered_chars)
else:
    print("Not solveable")