from z3 import BitVec, Solver, sat

# Replace this with the actual numbers printed by the program
with open("./msg.txt") as f:
    OUTPUTS = [int(line) for line in f.readlines()[1:]]

BITS = 56
MASK = 0xFF

# Symbolic variables for the LCG parameters
A = BitVec('A', BITS)
B = BitVec('B', BITS)
SEED = BitVec('SEED', BITS)

s = Solver()

# Define the state sequence
states = [SEED]
for i in range(len(OUTPUTS)):
    # Next state transition: s_i+1 = (s_i * A + B) & (2**56 - 1)
    # Z3 BitVec arithmetic handles the modulo 2^n automatically
    next_state = (states[i] * A) + B
    states.append(next_state)

# Known prefix: "CSCG{"
known_prefix = "CSCG{"
for i, char in enumerate(known_prefix):
    # output = (state_i+1 & MASK) ^ ord(FLAG[i])
    # Therefore: (state_i+1 & MASK) == output ^ ord(FLAG[i])
    s.add((states[i+1] & MASK) == OUTPUTS[i] ^ ord(char))

# If the flag ends with '}', add that constraint as well
s.add((states[len(OUTPUTS)] & MASK) == OUTPUTS[-1] ^ ord('}'))

if s.check() == sat:
    model = s.model()
    print("Found LCG Parameters:")
    print(f"A: {model[A]}")
    print(f"B: {model[B]}")
    print(f"SEED: {model[SEED]}")
    
    # Recover the flag
    flag = ""
    for i in range(len(OUTPUTS)):
        # Extract the lower 8 bits of the calculated state
        state_val = model.evaluate(states[i+1]).as_long()
        flag_char = chr((state_val & MASK) ^ OUTPUTS[i])
        flag += flag_char
    
    print(f"Flag: {flag}")
else:
    print("Unsatisfiable: Check if the prefix or output data is correct.")