from unicorn import *
from unicorn.arm64_const import *
from ghidra_assistant.utils.utils import *
from ghidra_assistant.utils.archs.arm64.arm64_emulator import ARM64UC_Emulator
from ghidra_assistant.utils.archs.arm64.asm_arm64 import ShellcodeCrafterARM64
import sys, pathlib
from add_memory_maps import *
from tools import *
from amlogic_devices import *
from amlogic_rom_db import *

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

def get_local_file(name):
    return pathlib.Path(pathlib.Path(__file__).parent.resolve() / name)

class Amlogic_Emulator(ARM64UC_Emulator):
    def __init__(self, file = "bin/BootROM_s905x3.bin", debug=False, device="S905X3") -> None:
        '''
            Amlogic bootrom emulator.
        '''
        super().__init__(True)
        
        self.logger = setup_logger("GSCEmulator")
        self.logger.setLevel(logging.DEBUG)
        self.debug = debug

        self.bootrom = open(file, 'rb').read()
        self.uc = Uc(UC_ARCH_ARM64, UC_MODE_LITTLE_ENDIAN)
        self.sc = ShellcodeCrafterARM64(None, None)
        self.device_offsets = amlogic_rom_db[device]
        self.setup_memory()
        self.setup_registers()
        self.setup_hooks()
        self.setup_devices()
        self.apply_patches()
        self.debug = False
        self.afl_input = []
        self.em = None # By default GA not used
        self.wait_for_usb = False # USB not ready

        self.I2C_inject_buffer = b""

        self.pc_trace = []
        self.enable_trace = False

    def write(self, address, data):
        self.uc.mem_write(address, data)

    def read(self, address, size):
        return bytes(self.uc.mem_read(address, size))
    
    def apply_patches(self):
        self.write(self.device_offsets['fastboot_response'], self.sc.ret_ins)
        self.write(self.device_offsets['usb_ep_queue'], self.sc.ret_ins)

    def setup_memory(self):
        #Bootrom
        self.uc.mem_map(ENTRY_POINT, len(self.bootrom), UC_PROT_EXEC | UC_PROT_READ)
        self.write(ENTRY_POINT, self.bootrom)

        # Timer  0xff800000|0xff800fff|rti
        self.uc.mem_map(0xff800000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # SRAM 0xfffa0000|0xfffe7fff |
        self.uc.mem_map(0xfffa0000, 0x48000, UC_PROT_READ | UC_PROT_WRITE)

        # ISA 0xffd0f000|0xffd0ffff|isa
        self.uc.mem_map(0xffd0f000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # eFuse 0xff630000|0xff631fff|efuse
        self.uc.mem_map(0xff630000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)
        self.uc.mem_write(0xff630000, open(get_local_file("bin/efuse_ff630000_ff632000.bin"), 'rb').read())

        # UART 0xff803000|0xff803fff|uart and add hook
        self.uc.mem_map(0xff803000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # ?? 0xffd08000|0xffd08fff|assist
        self.uc.mem_map(0xffd08000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # 0xff634400|0xff6347ff|periphs_reg
        self.uc.mem_map(0xff634400, 0x400, UC_PROT_READ | UC_PROT_WRITE)


        # Reset? 0xffd01000|0xffd01fff|reset
        self.uc.mem_map(0xffd01000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # RAM for USB? ff63c000
        self.uc.mem_map(0xff63c000, 0x1000, UC_PROT_READ | UC_PROT_WRITE)

        # UNKNOWN
        self.uc.mem_map(0xffe09000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)

        # DMA 0xff63e000|0xff63ffff|dma
        self.uc.mem_map(0xff63e000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)

        # 0xff63a000|0xff63bfff|usbphy21
        self.uc.mem_map(0xff63a000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)

        # USB1 0xff400000|0xff4fffff|usb1
        self.uc.mem_map(0xff400000, 0x100000, UC_PROT_READ | UC_PROT_WRITE) #1MB

        # UNKNOWN 0xffd14000|0xffd14fff|spifc
        self.uc.mem_map(0xffd14000, 0x1000, UC_PROT_READ | UC_PROT_WRITE) #1MB

        # Flash 0xf6000000|0xf9ffffff|flash #TODO add flash device for fuzzing?
        self.uc.mem_map(0xf6000000, 0xf9ffffff + 1 - 0xf6000000, UC_PROT_READ | UC_PROT_WRITE)

        # Emmc 0xffe07000|0xffe08fff|emmcC TODO add eMMC device for fuzzing?
        self.uc.mem_map(0xffe07000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)
        # 0xffe05000|0xffe06fff|emmcB
        self.uc.mem_map(0xffe05000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)
        # FFE03000 FFE04FFF emmcA
        self.uc.mem_map(0xFFE03000, 0x2000, UC_PROT_READ | UC_PROT_WRITE)

    def write_ptr(self, at, ptr):
        return self.uc.mem_write(at, p32(ptr))

    def read_ptr(self, at):
        return u32(self.uc.mem_read(at, 4))

    def setup_registers(self):
        self.pc = ENTRY_POINT
        self.sp = STACK_ADDRESS

    def setup_hooks(self):
        if self.debug:
            self.uc.hook_add(UC_HOOK_CODE, self.hook_code)
        self.uc.hook_add(UC_HOOK_INSN_INVALID, self.hook_invalid_ins)
        self.uc.hook_add(UC_HOOK_INTR, self.hook_exception)
        self.uc.hook_add(UC_HOOK_MEM_READ|UC_HOOK_MEM_WRITE, self.hook_devices, None, 0xff3f0000, 0xff30cffe + 1)

        #Add hook usb device
        self.uc.hook_add(UC_HOOK_MEM_READ|UC_HOOK_MEM_WRITE, self.hook_devices, None, 0xff63a000, 0xff63bfff + 1)
        self.uc.hook_add(UC_HOOK_MEM_READ|UC_HOOK_MEM_WRITE, self.hook_devices, None, 0xff400000, 0xff4fffff + 1)

    def setup_devices(self):
        self.devices = [
            # USBDevice(self),
            UARTDevice(self),
            TimerDevice(self),
            GPIODevice(self),
            EfuseDevice(self),
            USBPHYDevice(self),
            USB1Device(self),
        ]

    def hook_devices(self, handle, access, address, size, value, user_data):
        if address == 0xff630044 and access == UC_MEM_WRITE:
            pass
        for device in self.devices:
            if address >= device.BASE and address < device.BASE + device.SIZE:
                if access == UC_MEM_READ:
                    device.read(address, size)
                else:
                    device.write(address, size, value)
                return True
        return True

    def get_device(self, name):
        for dev in self.devices:
            if dev.NAME == name:
                return dev
        return None
    
    def run(self):
        while True:
            self.uc.emu_start(self.pc, self.pc + 4)

    def run_until_usb(self):
        info("Run until USB")
        try:
            while True:
                self.uc.emu_start(self.pc, self.pc + 4)
                if self.wait_for_usb:
                    return
        except Exception as e:
            print(bytes(self.get_device("UART").get_rx()).decode(errors="ignore"))
            self.print_ctx()
            print(str(e))
            if e.errno == UC_ERR_READ_UNMAPPED or e.errno == UC_ERR_WRITE_UNMAPPED:
                pass 
                # print(find_memory())
            exit(0)

    @property
    def pc(self):
        return self.uc.reg_read(UC_ARM64_REG_PC)

    @pc.setter
    def pc(self, value):
        self.uc.reg_write(UC_ARM64_REG_PC, value)

    @property
    def sp(self):
        return self.uc.reg_read(UC_ARM64_REG_SP)

    @sp.setter
    def sp(self, value):
        self.uc.reg_write(UC_ARM64_REG_SP, value)
        
    @property
    def X0(self):
        return self.uc.reg_read(UC_ARM64_REG_X0)

    @X0.setter
    def X0(self, value):
        self.uc.reg_write(UC_ARM64_REG_X0, value)

    @property
    def X1(self):
        return self.uc.reg_read(UC_ARM64_REG_X1)

    @X1.setter
    def X1(self, value):
        self.uc.reg_write(UC_ARM64_REG_X1, value)

    def pc_in_debug(self):
        # for start, end in debug_functions:
        #     if self.pc >= start and self.pc <= end:
        #         return True
        return False

    def dump_stack(self, n_offset = 0, size = 0x200):
        hexdump(bytes(self.uc.mem_read(self.sp - n_offset, size)))
        
    def add_breakpoint(self, at):
        def _breakpoint(uc : Uc, address, size, em : Amlogic_Emulator):
            em.print_ctx(print)
            pass
        self.uc.hook_add(UC_HOOK_CODE, _breakpoint, self, at, at + 1)

    def hook_code(self, uc : Uc, address, size, user_data):
        '''
        Hook startup code to continue boot process
        '''
        if self.enable_trace:
            self.pc_trace.append(hex(self.pc))
        # Reset expects value in DRAM
        if self.pc == 0xffff00b0:
            self.write_ptr(0xff800228, uc.reg_read(UC_ARM64_REG_X1))
        if self.pc == 0xffff62dc:
            uc.reg_write(UC_ARM64_REG_W21, uc.reg_read(UC_ARM64_REG_W2))

        # Hooks for I2c ==========
        if self.pc == 0xffff569c:
            # Address for HDMI i2c is 0xa4
            # self.logger.debug(f"i2c reading from address: {hex(self.uc.reg_read(UC_ARM64_REG_W0))}")
            self.get_device("GPIO").I2C_Address = self.uc.reg_read(UC_ARM64_REG_W0)
        if self.pc == 0xffff56e8:
            # self.logger.info(f"Data from i2c address={hex(self.get_device('GPIO').I2C_Address)} value={hex(self.uc.reg_read(UC_ARM64_REG_W0))}")
            pass
        #Dirty hook to inject I2C data, instead of implementing a full I2CDevice
        if self.pc == 0xffff5760:
            self.I2C_inject_buffer = b"boot@USB\x00"
            self.i2c_offset = 0
        if len(self.I2C_inject_buffer) > 0 and self.pc == 0xffff56d8:
            self.uc.reg_write(UC_ARM64_REG_W20, self.I2C_inject_buffer[self.i2c_offset])
            self.i2c_offset += 1
        
        # Add hook for injecting i2c boot string:
        if self.pc == 0xffff5790:
            self.logger.info(f"I2C boot string: {self.uc.mem_read(self.sp + 0x10, 8).decode(errors='ignore')}")
        # ====================
        if self.get_device("UART").has_rx():
            print(bytes(self.get_device("UART").get_rx()).decode(errors="ignore"), end="")
        
        if self.pc_in_debug():
            self.print_ctx()
            # breakpoint()

        # USB is here
        if self.pc == 0xffff9d18:
            self.wait_for_usb = True
            uc.emu_stop()
            pass

    def hook_invalid_ins(self, handle, intno, user_data):
        pass

    def read_string(self, at):
        if at == 0:
            return b''
        s = b''
        while 1:
            b = self.uc.mem_read(at, 1)
            at += 12
            if b == b'\0':
                return s
            s += b
        return s

    def hook_exception(self, handle : Uc, intno, user_data):
        if intno == 1:
            handle.reg_write(UC_ARM64_REG_PC, handle.reg_read(UC_ARM64_REG_PC) + 4)
            return True
        else:
            # Invalid instruction
            pass
        
    def place_fastboot_command(self, input):
        assert len(input) <= 0x200, "Too much data!"
        
        # Setup fake request in X1
        self.X1 = TEST_OFFSET
        # Set request length to 0x200
        self.write_ptr(TEST_OFFSET + 8, 0x200)
        
        # Write input buffer
        self.write_ptr(TEST_OFFSET, TEST_REQ_BUFFER)
        self.write(TEST_REQ_BUFFER, input)
        
        # Setup context
        self.write_ptr(self.X1 + 0x28 + 0x4, 64)
        
    def run_fastboot_cmd(self):
        def _hook_usb_ep_que(uc : Uc, address, size, em : Amlogic_Emulator):
            # em.uc.emu_stop()
            return True
        # self.uc.hook_add(UC_HOOK_CODE, _hook_usb_ep_que, self, 0xffff8998, 0xffff8998 + 1)
            
        def _hook_fastboot_tx_write(uc :Uc, address, size, em : Amlogic_Emulator):
            if em.debug:
                hexdump(em.read(em.X0, em.X1))
            # em.uc.emu_stop()
            return True
        self.uc.hook_add(UC_HOOK_CODE, _hook_fastboot_tx_write, self, self.device_offsets['fastboot_response'], self.device_offsets['fastboot_response'] + 1)
        
        self.pc = FASTBOOT_CMD_HANDLER
        self.sp = STACK_ADDRESS
        
        # Run 
        if self.debug:
            self.enable_trace = True
        self.uc.emu_start(self.pc, 0)
        return 0

if __name__ == "__main__":
    device = "S905X3"
    
    emulator = Amlogic_Emulator()