import angr
import subprocess
import codec
import sys
import claripy

base = 0x400000
start = 0x1220 + base
cmp_start = 0x1242 + base
cmp_len = 0x8a - 0x42
stop =  0x12a0 + base

project = angr.Project("Stone")

state = project.factory.blank_state(addr=start)
state.regs.rax = 0x10
state.regs.rsi = 0x2
state.regs.rdi = 0x0

word_ir_contracted = []

# A word attempt is expanded into 5 * 16 = 80 bytes, and these bytes
# are contracted again in `check_has_phobos` to 10 bytes.
for i in range(10):
    byte = claripy.BVS('byte_%d' % i, 8)
    word_ir_contracted.append(byte)

@project.hook(cmp_start, length=cmp_len)
def get_byte(state):
    # Shims the decoding part of the loop.
    rdx = state.solver.eval(state.regs.rdx)
    state.regs.al = word_ir_contracted[rdx]

sim = project.factory.simgr(state)
sim.explore(find=stop)

state = sim.one_found
state.add_constraints(state.regs.al == 0) # at the end of the function we want al == 0

word_ir_total = claripy.Concat(*word_ir_contracted)

print("Example valid words:")
sols = state.solver.eval_upto(word_ir_total, 20) # get, say, 20 solutions.
for sol in sols:
    bits = bin(sol)[2:].rjust(80, '0')
    ir = bytes(int(bit) for bit in bits)
    string = codec.contract(ir)
    print(string)
