from __future__ import print_function

import os
import subprocess as sp

import angr

from claripy import BVS, BVV

src = """
#include <stdio.h>
#include <stdlib.h>

int verify(char* user, char* pass) {
    char* u = user;
    char* p = pass;
    if (p[0] == u[1] ^ u[2]) {
        return 1;
    } else {
        puts("nope");
        exit(1);
    }
}

int main(int argc, char* argv[]) {
    char* username = argv[1];
    char* password = argv[2];
    int b = verify(username, password);
    return b;
}
"""

if not os.path.exists("./test"):
    print("compiling binary. Test this on x86_64.")
    with open("./test.c", "w") as f:
        f.write(src)
    # I compiled with gcc version 6.3.1 20161221 (Red Hat 6.3.1-1) (GCC)
    # on fedora 25 x86_64
    sp.check_call(['make', 'test'])

proj = angr.Project('./test',
                    load_options={"auto_load_libs": False})
verify_addr = proj.loader.main_bin.get_symbol('verify').addr
block = proj.factory.block(verify_addr)
print("got the function")
block.pp()
"""
0x400546:	push	rbp
0x400547:	mov	rbp, rsp
0x40054a:	sub	rsp, 0x20
0x40054e:	mov	qword ptr [rbp - 0x18], rdi
0x400552:	mov	qword ptr [rbp - 0x20], rsi
0x400556:	mov	rax, qword ptr [rbp - 0x18]
...
"""
start_addr = block.capstone.insns[5].insn.address
alloca = 0x20
offset = 0x18

bvs = BVS("bvs", 12 * 8)
bvs_addr = BVV(0x12345678, 64)
# create a initial state
initial_state = proj.factory.blank_state(addr=start_addr)

print("setting up state")
# manually do function prologue
initial_state.regs.rbp = initial_state.regs.rsp
initial_state.regs.rsp -= alloca
initial_state.memory.store(bvs_addr, bvs)
# this is what I would intuitively do
initial_state.memory.store((initial_state.regs.rbp - offset),
                           bvs_addr)

print("rsp =", initial_state.regs.rsp)
print("rbp =", initial_state.regs.rbp)
print("bvs @ ", bvs_addr)
print("ptr @ ", initial_state.regs.rbp - offset)
block = proj.factory.block(initial_state.ip.args[0])
print("executing the following instruction:")
print(str(block.capstone.insns[0]))
path = proj.factory.path(initial_state)
p = path.step(num_inst=1)[0]
print("expecting rax == [rbp - 0x{:x}]".format(offset))
print("rax =", p.state.regs.rax)
print("[rbp - 0x{:x}] =".format(offset),
      p.state.memory.load(p.state.regs.rbp - offset))
# the endianess is reversed here
assert p.state.regs.rax.reversed.args[0] == bvs_addr.args[0]
# this is what I would expect
assert p.state.regs.rax.args[0] == bvs_addr.args[0]
"""
setting up state
rsp = <BV64 0x7fffffffffeffe0>
rbp = <BV64 0x7ffffffffff0000>
bvs @  <BV64 0x12345678>
ptr @  <BV64 0x7fffffffffeffe8>
executing the following instruction:
0x400556:	mov	rax, qword ptr [rbp - 0x18]
expecting rax == [rbp - 0x18]
rax = <BV64 0x7856341200000000>
[rbp - 0x18] = <BV64 0x12345678>
"""