from emulator import *
import unicornafl
import argparse, os

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_fastboot(device="S905X3"):
    emulator = Amlogic_Emulator(device=device)
    # 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, "fuzz_fastboot_input/getvar", _place_fb_command, _run, persistent_iters=1)
    
def afl_fuzzer_verifybl(device="S905X3"):
    emulator = Amlogic_Emulator(device=device)
    # emulator.debug = True

    def place_bootloader_header(uc, input, persistent_round, data):
        # hexdump(bytes(input), "input")
        if len(input) != 0x40:
            return False
        # Filter some unsupported commands:
        if input[:4] == b"@AML":
            return False
        emulator = Amlogic_Emulator(device=device)
        emulator.place_bootloader(input)
        return True

    def _run(uc, data):
        emulator.run_check_bootloader_header()
        return 0
        
    unicornafl.uc_afl_fuzz_custom(emulator.uc, "fuzz_verifybl_input/fuzz_bl.bin", place_bootloader_header, _run, persistent_iters=1)
    
    
if __name__ == "__main__":
    # print(sys.argv)
    # 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)
    #     print("", flush=True)
    #     exit(0)
        
    # if args.input:
    #     # Run a single comand
    #     test_fb_cmd(open(args.input, 'rb').read())
    # else:
    #     pass
    # Run AFL
    
    if os.getenv("AFL_FUZZER") == "verifybl":
        afl_fuzzer_verifybl()
    elif os.getenv("AFL_FUZZER") == "fastboot":
        afl_fuzzer_fastboot()
    else:
        print("No AFL fuzzer selected")
        exit(1)