import angr
import subprocess
import codec
import sys
import claripy

def word_for_expanded_ir(exp_ir):
    bits = bin(exp_ir)[2:].rjust(640, '0')
    bytes_ = bytes(int(bits[i*8:(i+1)*8]) for i in range(80))
    word = codec.contract(bytes_)
    return word

def gen_word_ir_vars():
    word_ir_expanded = []
    word_ir_contracted = []
    for i in range(80):
        bit = claripy.BVS('byte_%d' % i, 1)
        rest = claripy.BVV(0, 7)
        b = claripy.Concat(rest, bit)
        word_ir_expanded.append(b)
        word_ir_contracted.append(bit)

    return (word_ir_expanded, word_ir_contracted)

def solve(word_ir_expanded, word_ir_contracted):
    project = angr.Project("Stone")

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

    get_phobos_start = 0x46b0 + base
    word_ir_base = 0xff00ff00

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

    phobos_out_contracted = []

    @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 = claripy.Concat(*word_ir_contracted[rdx*8:(rdx+1)*8])

    @project.hook(0x46f0 + base, length=3)
    def store_bit(state):
        # Instead of storing the nth bit of the phobos into memory,
        # we just save the symbolic bits into an array here.
        # Later, we can solve for the contents of this array
        # to get possible phobi.
        phobos_out_contracted.append(state.regs.al[0].to_claripy())

    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_expanded)
    state.memory.store(word_ir_base, word_ir_total)

    # Lots of zero-init here to keep angr from complaining
    # about uninitialized registers (even though it might not matter)
    state.regs.rip = get_phobos_start
    state.regs.rsi = word_ir_base
    state.regs.rdi = 0x0
    state.regs.rcx = 0x55
    state.regs.r13 = 0
    state.regs.r9 = 0
    state.regs.r11 = 0
    state.regs.r12 = 0
    state.regs.rbp = 0
    state.regs.rbx = 0

    sim.move('found', 'active')

    sim.explore(find=stop2)
    phobos = claripy.Concat(*phobos_out_contracted)

    return (sim.one_found, phobos)

def main():
    wire, wirc = gen_word_ir_vars()
    state, phobos = solve(wire, wirc)
    phobi = state.solver.eval_upto(phobos, 500)
    word_ir_total = claripy.Concat(*wire)

    print(f"Found {len(phobi)} phobi. Outputting phobos.out...")
    with open('phobos.out', 'w') as f:
        for phobos_ in phobi:
            hex_rep = hex(phobos_)[2:].rjust(14, '0')
            example_word = state.solver.eval(word_ir_total, extra_constraints=(phobos == phobos_,))
            example_word = word_for_expanded_ir(example_word)
            f.write(f"{hex_rep} {example_word}\n")

if __name__ == '__main__':
    main()
