import os
import time
import logging

import angr
import angrop  # pylint: disable=unused-import
from angrop.errors import RopException

l = logging.getLogger(__name__)

bin_path = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..', 'binaries'))
tests_dir = os.path.join(bin_path, 'tests')
data_dir = os.path.join(bin_path, 'tests_data', 'angrop_gadgets_cache')


"""
Suggestions on how to debug angr changes that break angrop.

If the gadget is completely missing after your changes. Pick the address that didn't work and run the following.
The logging should say why the gadget was discarded.

rop = p.analyses.ROP()
angrop.gadget_analyzer.l.setLevel("DEBUG")
rop._gadget_analyzer.analyze_gadget(addr)

If a gadget is missing memory reads / memory writes / memory changes, the actions are probably missing.
Memory changes require a read action followed by a write action to the same address.
"""

def test_badbyte():
    cache_path = os.path.join(data_dir, "bronze_ropchain")
    proj = angr.Project(os.path.join(tests_dir, "i386", "bronze_ropchain"), auto_load_libs=False)
    rop = proj.analyses.ROP()
    rop.set_badbytes([0x0, 0x0a])

    if os.path.exists(cache_path):
        rop.load_gadgets(cache_path)
    else:
        rop.find_gadgets()
        rop.save_gadgets(cache_path)

    # make sure it can set 0 first
    chain = rop.set_regs(eax=0)
    state = chain.exec()
    assert not state.regs.eax.symbolic
    assert state.solver.eval(state.regs.eax == 0)
    assert all(x not in chain.payload_str() for x in [0, 0xa])

    # make sure it can set 0x16, which requires gadgets like `mov eax, 0x16`
    chain = rop.set_regs(eax=0x16)
    state = chain.exec()
    assert not state.regs.eax.symbolic
    assert state.solver.eval(state.regs.eax == 0x16)
    assert all(x not in chain.payload_str() for x in [0, 0xa])

    # make sure it can set 0xb
    # this binary does not have gadget that sets 0xb directly, it needs to do calculation
    # something like `mov eax, 8` and then `add eax, 3`
    chain = rop.set_regs(eax=0xb)
    state = chain.exec()
    assert not state.regs.eax.symbolic
    assert state.solver.eval(state.regs.eax == 0xb)
    assert all(x not in chain.payload_str() for x in [0, 0xa])

    # make sure it can write '/bin/sh\x00\n' into memory, notice that '\x00' and '\n' are bad bytes
    ptr = rop.chain_builder._mem_writer._get_ptr_to_writable(9+4)
    chain = rop.write_to_mem(ptr, b'/bin/sh\x00\n')
    state = chain.exec()
    assert state.solver.eval(state.memory.load(ptr, 9), cast_to=bytes) == b'/bin/sh\x00\n'
    assert all(x not in chain.payload_str() for x in [0, 0xa])

    # finally, make sure setting multiple registers can work
    nullptr = rop.chain_builder._mem_writer._get_ptr_to_null()
    chain = rop.set_regs(eax=0xb, ebx=ptr, ecx=nullptr, edx=nullptr)
    state = chain.exec()
    assert not state.regs.eax.symbolic
    assert state.solver.eval(state.regs.eax == 0xb)
    assert not state.regs.ebx.symbolic
    assert state.solver.eval(state.regs.ebx == ptr)
    assert not state.regs.ecx.symbolic
    assert state.solver.eval(state.regs.ecx == nullptr)
    assert not state.regs.edx.symbolic
    assert state.solver.eval(state.regs.edx == nullptr)
    assert all(x not in chain.payload_str() for x in [0, 0xa])

def test_badbyte_transform():
    cache_path = os.path.join(data_dir, "bronze_ropchain")
    proj = angr.Project(os.path.join(tests_dir, "i386", "bronze_ropchain"), auto_load_libs=False)
    rop = proj.analyses.ROP()

    if os.path.exists(cache_path):
        rop.load_gadgets(cache_path)
    else:
        rop.find_gadgets()
        rop.save_gadgets(cache_path)

    rop.set_badbytes([0x00, 0x0A])
    chain = rop.write_to_mem(0xdeadbeef, b"\x00", fill_byte=b"A")
    state = chain.exec()
    assert state.memory.load(0xdeadbeef, 1).concrete_value == 0x00

    payload = chain.payload_str()
    assert b"\x00" not in payload
    assert b"\x0A" not in payload

def test_badbyte_multibyte():
    cache_path = os.path.join(data_dir, "bronze_ropchain")
    proj = angr.Project(os.path.join(tests_dir, "i386", "bronze_ropchain"), auto_load_libs=False)
    rop = proj.analyses.ROP()

    if os.path.exists(cache_path):
        rop.load_gadgets(cache_path)
    else:
        rop.find_gadgets()
        rop.save_gadgets(cache_path)

    rop.set_badbytes([0x00, 0x0A])
    target = b"\x00\x00\x00\x42"
    chain = rop.write_to_mem(0xdeadbf00, target, fill_byte=b"\xff")
    state = chain.exec()
    data = state.solver.eval(state.memory.load(0xdeadbf00, len(target)), cast_to=bytes)
    assert data == target

    payload = chain.payload_str()
    assert b"\x00" not in payload
    assert b"\x0A" not in payload

def test_hard_regs_loop():
    cache_path = os.path.join(data_dir, "bronze_ropchain")
    proj = angr.Project(os.path.join(tests_dir, "i386", "bronze_ropchain"), auto_load_libs=False)
    rop = proj.analyses.ROP()

    if os.path.exists(cache_path):
        rop.load_gadgets(cache_path)
    else:
        rop.find_gadgets()
        rop.save_gadgets(cache_path)

    rop.set_badbytes([0x00, 0x0A])
    start = time.time()
    try:
        rop.set_regs(eax=0x4200009a, edx=0xdeadbefc)
    except RopException:
        pass

    # if should take less than 0.1s, but let's be lenient here
    assert time.time() - start < 2

def run_all():
    functions = globals()
    all_functions = {x:y for x, y in functions.items() if x.startswith('test_')}
    for f in sorted(all_functions.keys()):
        if hasattr(all_functions[f], '__call__'):
            all_functions[f]()


if __name__ == "__main__":
    logging.getLogger("angrop.rop").setLevel(logging.DEBUG)

    import sys
    if len(sys.argv) > 1:
        globals()['test_' + sys.argv[1]]()
    else:
        run_all()
