from emulator import *
import unicornafl
import argparse

ENTRY_POINT = 0xffff0000
STACK_ADDRESS = 0xfffe3800
    
# FASTBOOT_CMD_HANDLER = 0xffff9758
TEST_OFFSET = 0xfffa0000 + 0x8000
TEST_REQ_BUFFER = TEST_OFFSET + 0x800
TEST_CONTEXT_BUFFER = TEST_OFFSET + 0x9000

debug_functions = [
    # (start, end)
    # (0xffff9bc4, 0xffff9d6c), # Fastboot
    # (0xffff66d8, 0xffff6754),
]
  
def test_fb_cmd(cmd=b'getvar:version', device="S905X3"):
    emulator = Amlogic_Emulator(device=device, debug=True)
    emulator.debug = True
    emulator.place_fastboot_command(cmd)
    res = emulator.run_fastboot_cmd()
    pass

def afl_fuzzer():
    emulator = Amlogic_Emulator()
    # emulator.debug = True

    def _place_fb_command(uc, input, persistent_round, data):
        # hexdump(bytes(input), "input")
        if len(input) > 0x200:
            return False
        # Filter some unsupported commands:
        if input[:4] == b"boot":
            return False
        emulator.place_fastboot_command(input)
        return True

    def _run(uc, data):
        emulator.run_fastboot_cmd()
        return 0
        
    unicornafl.uc_afl_fuzz_custom(emulator.uc, "input/getvar", _place_fb_command, _run, persistent_iters=1)
    
if __name__ == "__main__":
    args = argparse.ArgumentParser("Amlogic BootROM Fuzzer")
    args.add_argument("--device", "-d", help="Device to test", default="S905X3")
    args.add_argument("--test", "-t", help="Test fastboot command", default=False, action="store_true")
    args.add_argument("--input", "-i", help="Input file for crash", default=None)
    
    args = args.parse_args()
    
    if args.test:
        test_fb_cmd(device=args.device)
        exit(0)
        
    if args.input:
        # Run a single comand
        test_fb_cmd(open(args.input, 'rb').read())
    else:
        # Run AFL
        afl_fuzzer()