import struct, sys, usb1, libusb1, ctypes, usb, argparse
from keystone import *
from capstone import *
from ghidra_assistant.utils.utils import *
from ghidra_assistant.concrete_device import *
from ghidra_assistant.utils.debugger.debugger_archs.ga_arm64 import GA_arm64_debugger
from qiling.const import QL_ARCH
import os, tqdm, datetime

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)

class ExynosDevice():
    """
    Class to exploit a Exynos device (8890/8895) using the USB stack.
    """

    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):
        """Setup proper connection, and ensure the connection is alive"""
        self.context = usb1.USBContext()

        while True:
            self.handle = self.context.openByVendorIDAndProductID(
                vendor_id=self.idVendor,
                product_id=self.idProduct,
                skip_on_error=False
            )
            if self.handle == None:
                continue
            break

        try:
            self.handle.getDevice().getSerialNumber()
        except Exception as e:
            if e.value == usb1.libusb1.LIBUSB_ERROR_TIMEOUT or e.value == usb1.libusb1.LIBUSB_ERROR_IO:
                print("Device disconnected / not connected. Reconnect USB?")
                sys.exit(1)
            else:
                raise e

        # claim usb interface
        self.handle.claimInterface(0)
        print(f"Connected device! {hex(self.idVendor)} {hex(self.idProduct)}")


    def disconnect(self):
        """Disconnect the device"""
        self.handle.releaseInterface(0)
        self.handle.close()
        self.context.exit()


    def write(self, data):
        """Write data to the device"""
        transferred = ctypes.c_int()
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, data, len(data), ctypes.byref(transferred), 0)
        assert(res == 0), "Could not perform bulk transfer"
        return res
    
    
    def usb_write(self, data):
        assert len(data) <= 0x200, "Data too big"
        transferred = ctypes.c_int()
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_OUT, data, len(data), ctypes.byref(transferred), 300)

        assert res == 0, f"Error sending data {res}"
        assert transferred.value == len(data), f"Invalid transfered size {transferred.value} != {len(data)}"
        return transferred.value
    

    def usb_read(self, size):
        transferred = ctypes.c_int()
        buf = ctypes.c_buffer(b"", size)
        res = libusb1.libusb_bulk_transfer(self.handle._USBDeviceHandle__handle, ENDPOINT_BULK_IN, buf, len(buf), ctypes.byref(transferred), 300)
        assert res == 0, f"Error receiving data {res}"
        return buf.raw[:transferred.value]