#include <windows.h>
#include <cfgmgr32.h>
#include <usbioctl.h>

#include "hook/table.h"
#include "hook/iohook.h"
#include "hooklib/setupapi.h"

#include "util/dprintf.h"
#include "util/str.h"
#include "es3sec.h"

#define DONGLE_STR_IDX_MANUFACTURER 1
#define DONGLE_STR_IDX_PRODUCT 2
#define DONGLE_STR_IDX_SERIAL_NUMBER 3

static const wchar_t DEVNAME_HUB[] = L"$hub";
static const wchar_t DEVNAME_HUB_[] = L"\\\\.\\$hub";
static const wchar_t DEVNAME_DONGLE[] = L"$dongle";
static const wchar_t HUB_HCD_DRIVER_KEY[] = L"{36fc9e60-c465-11cf-8056-444553540000}\\0001";
// This is what a real dongle is set to
static const wchar_t HUB_NODE_CONNECTION_DRIVER_KEY[] = L"{c8b76578-d062-4834-0001-f8b6f2162a22}\\0003";
static const wchar_t root_hub_name[] = L"Fake Root Hub";
static const wchar_t FILLER_LANGIDS[] = L"\uffff\uffff";
static const DEVINST HUB_DEVINST = 573;
static const DEVINST DONGLE_DEVINST = 5730;
static struct es3sec_config config;
static HANDLE dongle_fd;
static HANDLE hub_fd;
static IID hubs_iid;
static DEVINST root_dev_inst;
static USHORT dongle_vid;
static USHORT dongle_pid;
static wchar_t dongle_mftr[MAX_PATH];
static wchar_t dongle_product[MAX_PATH];

static HRESULT es3sec_handle_hub_irp(struct irp *irp);
static HRESULT es3sec_handle_hub_open(struct irp *irp);
static HRESULT es3sec_handle_hub_close(struct irp *irp);
static HRESULT es3sec_handle_hub_ioctl(struct irp *irp);

static HRESULT es3sec_hub_handle_driverkey(struct irp *irp);
static HRESULT es3sec_hub_handle_roothub(struct irp *irp);
static HRESULT es3sec_hub_handle_nodeinfo(struct irp *irp);
static HRESULT es3sec_hub_connection_info_ex(struct irp *irp);
static HRESULT es3sec_hub_descriptor_from_node(struct irp *irp);
static HRESULT es3sec_hub_connection_driver_key_name(struct irp *irp);
static HRESULT es3sec_hub_get_config_descriptor(struct irp *irp, PUSB_DESCRIPTOR_REQUEST req, UCHAR dataRequested);
static HRESULT es3sec_hub_get_string_descriptor(struct irp *irp, PUSB_DESCRIPTOR_REQUEST req, UCHAR dataRequested);

static CONFIGRET my_CM_Get_Child(PDEVINST pdnDevInst, DEVINST dnDevInst, ULONG ulFlags);
static CONFIGRET (*next_CM_Get_Child)(PDEVINST pdnDevInst, DEVINST dnDevInst, ULONG ulFlags);

static CONFIGRET my_CM_Get_DevNode_Registry_PropertyW(
    DEVINST dnDevInst,
    ULONG ulProperty,
    PULONG pulRegDataType,
    PVOID Buffer,
    PULONG pulLength,
    ULONG ulFlags
);
static CONFIGRET (*next_CM_Get_DevNode_Registry_PropertyW)(
    DEVINST dnDevInst,
    ULONG ulProperty,
    PULONG pulRegDataType,
    PVOID Buffer,
    PULONG pulLength,
    ULONG ulFlags
);

static const struct hook_symbol cm_syms[] = {
    {
        .name = "CM_Get_DevNode_Registry_PropertyW",
        .patch = my_CM_Get_DevNode_Registry_PropertyW,
        .link = (void **) &next_CM_Get_DevNode_Registry_PropertyW
    },
    {
        .name = "CM_Get_Child",
        .patch = my_CM_Get_Child,
        .link = (void **) &next_CM_Get_Child
    },
};

HRESULT es3sec_hook_init(
    const struct es3sec_config *cfg, 
    USHORT vid, 
    USHORT pid,
    const wchar_t *manufacturer, 
    const wchar_t *product)
{
    HRESULT hr;
    assert(cfg != NULL);

    if (!cfg->enable) {
        return S_FALSE;
    }

    IIDFromString(L"{3ABF6F2D-71C4-462a-8A92-1E6861E6AF27}", &hubs_iid);
    setupapi_add_phantom_dev(&hubs_iid, DEVNAME_HUB);

    hr = iohook_open_nul_fd(&dongle_fd);
    hr = iohook_open_nul_fd(&hub_fd);

    if (FAILED(hr)) {
        return hr;
    }

    hr = iohook_push_handler(es3sec_handle_hub_irp);

    if (FAILED(hr))
    {
        return hr;
    }

    hook_table_apply(NULL, "setupapi.dll", cm_syms, _countof(cm_syms));

    CM_Locate_DevNodeW(&root_dev_inst, NULL, CM_LOCATE_DEVNODE_NORMAL);

    dongle_vid = vid;
    dongle_pid = pid;
    wcscpy_s(dongle_mftr, _countof(dongle_mftr), manufacturer);
    wcscpy_s(dongle_product, _countof(dongle_product), product); // Invalid param?
    memcpy(&config, cfg, sizeof(*cfg));

    dprintf("ES3 Dongle: init\tVID: %04X | PID: %04X | Manufacturer: %ls | Product: %ls | S/N: %ls\n", 
        vid, 
        pid, 
        dongle_mftr, 
        dongle_product, 
        config.serial);

    return S_OK;
}

static HRESULT es3sec_handle_hub_irp(struct irp *irp)
{
    assert(irp != NULL);
    if (irp->op != IRP_OP_OPEN && irp->fd != hub_fd)
    {
        return iohook_invoke_next(irp);
    }

    switch (irp->op)
    {
    case IRP_OP_OPEN:
        return es3sec_handle_hub_open(irp);
    case IRP_OP_CLOSE:
        return es3sec_handle_hub_close(irp);
    case IRP_OP_IOCTL:
        return es3sec_handle_hub_ioctl(irp);
    default:
        return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
    }
}

static HRESULT es3sec_handle_hub_open(struct irp *irp)
{
    if (!wstr_ieq(irp->open_filename, DEVNAME_HUB) && !wstr_ieq(irp->open_filename, DEVNAME_HUB_))
    {
        return iohook_invoke_next(irp);
    }

    dprintf("ES3 Dongle: Open USB Hub\n");
    irp->fd = hub_fd;
    return S_OK;
}

static HRESULT es3sec_handle_hub_close(struct irp *irp)
{
    dprintf("ES3 Dongle: Close Hub\n");
    return S_OK;
}

static HRESULT es3sec_handle_hub_ioctl(struct irp *irp)
{
    switch (irp->ioctl) {
        case 0x220424: return es3sec_hub_handle_driverkey(irp);
        case 0x220448: return es3sec_hub_connection_info_ex(irp);
        case 0x220410: return es3sec_hub_descriptor_from_node(irp);
        case 0x220420: return es3sec_hub_connection_driver_key_name(irp);
        case 0x220408:
            if (irp->read.nbytes == sizeof(USB_NODE_INFORMATION))
                return es3sec_hub_handle_nodeinfo(irp);

            else if (irp->read.nbytes >= sizeof(USB_ROOT_HUB_NAME))
                return es3sec_hub_handle_roothub(irp);

            else
            {
                dprintf("ES3 Dongle: Bad size for IOCTL %X\n", irp->ioctl);
                return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
            }

        default:
            dprintf("ES3 Dongle: Unknown hub IOCTL %X\n", irp->ioctl);
            return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
    }
}

static HRESULT es3sec_hub_handle_driverkey(struct irp *irp)
{
    size_t size_of_driver_key = sizeof(HUB_HCD_DRIVER_KEY);
    ULONG actual_length = size_of_driver_key + sizeof(USB_ROOT_HUB_NAME);
    HRESULT hr;

    if (irp->write.nbytes == sizeof(USB_HCD_DRIVERKEY_NAME)) // Requesting key size
    {
        USB_HCD_DRIVERKEY_NAME usb_hcd_driver_key_name;
        usb_hcd_driver_key_name.ActualLength = actual_length;
        hr = iobuf_write(&irp->read, &usb_hcd_driver_key_name, sizeof(usb_hcd_driver_key_name));

        if (FAILED(hr))
        {
            dprintf("ES3 Dongle: es3sec_hub_handle_driverkey iobuf_write failed! %lx\n", hr);
        }

        return hr;
    }

    // Requesting the key itself
    PUSB_HCD_DRIVERKEY_NAME usb_hcd_driver_key_name = (PUSB_HCD_DRIVERKEY_NAME)malloc(sizeof(USB_HCD_DRIVERKEY_NAME) + actual_length);
    usb_hcd_driver_key_name->ActualLength = actual_length;

    errno_t err = wcscpy_s(
        usb_hcd_driver_key_name->DriverKeyName,
        _countof(HUB_HCD_DRIVER_KEY),
        HUB_HCD_DRIVER_KEY);

    if (err)
    {
        dprintf("ES3 Dongle: es3sec_hub_handle_driverkey wcscpy_s failed with %X", err);
        return E_FAIL;
    }

    hr = iobuf_write(&irp->read, usb_hcd_driver_key_name, actual_length);

    if (FAILED(hr))
    {
        dprintf("ES3 Dongle: es3sec_hub_handle_driverkey iobuf_write failed! %lx\n", hr);
    }

    return hr;
}

static HRESULT es3sec_hub_handle_roothub(struct irp *irp)
{
    size_t size_of_hub_name = sizeof(DEVNAME_HUB);
    ULONG actual_length = size_of_hub_name + sizeof(USB_ROOT_HUB_NAME);
    HRESULT hr;
   
    if (irp->read.nbytes == sizeof(USB_ROOT_HUB_NAME)) { // Root hub name size
        USB_ROOT_HUB_NAME rhub;
        rhub.ActualLength = actual_length;
        hr = iobuf_write(&irp->read, &rhub, sizeof(rhub));
        
        if (FAILED(hr)) {
            dprintf("ES3 Dongle: es3sec_hub_handle_roothub iobuf_write failed! %lx\n", hr);
        }
        
        return hr;
    }

    //Root hub name
    PUSB_ROOT_HUB_NAME rhub = (PUSB_ROOT_HUB_NAME)malloc(actual_length);
    rhub->ActualLength = actual_length;

    errno_t err = wcscpy_s(
        rhub->RootHubName,
        _countof(DEVNAME_HUB),
        DEVNAME_HUB
    );

    if (err) {
        dprintf("ES3 Dongle: es3sec_hub_handle_roothub wcscpy_s failed with %X", err);
        return E_FAIL;
    }

    hr = iobuf_write(&irp->read, rhub, sizeof(USB_ROOT_HUB_NAME) + size_of_hub_name);

    if (FAILED(hr)) {
        dprintf("ES3 Dongle: iobuf_write failed! %lx\n", hr);
    }

    return hr;
}

static HRESULT es3sec_hub_handle_nodeinfo(struct irp *irp)
{
    USB_NODE_INFORMATION node_info;
    node_info.NodeType = UsbHub;
    node_info.u.HubInformation.HubDescriptor.bNumberOfPorts = 1;
    HRESULT hr = iobuf_write(&irp->read, &node_info, sizeof(node_info));

    if (FAILED(hr))
    {
        dprintf("ES3 Dongle: es3sec_hub_handle_nodeinfo iobuf_write failed! 0x%lX\n", hr);
    }

    return hr;
}

static HRESULT es3sec_hub_connection_info_ex(struct irp *irp)
{
    HRESULT hr;
    PUSB_NODE_CONNECTION_INFORMATION_EX conn_info = (PUSB_NODE_CONNECTION_INFORMATION_EX) malloc(irp->write.nbytes);
    hr = iobuf_read(&irp->write, conn_info, irp->write.nbytes);

    if (FAILED(hr)) {
        dprintf("ES3 Dongle: es3sec_hub_connection_info_ex Failed to read IRP %lx\n", hr);
        return hr;
    }

    conn_info->ConnectionStatus = DeviceConnected;
    conn_info->DeviceIsHub = false;

    conn_info->DeviceDescriptor.idVendor = dongle_vid;
    conn_info->DeviceDescriptor.idProduct = dongle_pid;
    conn_info->DeviceDescriptor.bLength = sizeof(conn_info->DeviceDescriptor);
    conn_info->DeviceDescriptor.iManufacturer = DONGLE_STR_IDX_MANUFACTURER;
    conn_info->DeviceDescriptor.iProduct = DONGLE_STR_IDX_PRODUCT;
    conn_info->DeviceDescriptor.iSerialNumber = DONGLE_STR_IDX_SERIAL_NUMBER;

    hr = iobuf_write(&irp->read, conn_info, irp->read.nbytes);

    if (FAILED(hr)) {
        dprintf("ES3 Dongle: es3sec_hub_connection_info_ex Failed to write IRP %lx\n", hr);
    }

    return hr;
}

static HRESULT es3sec_hub_descriptor_from_node(struct irp *irp)
{
    HRESULT hr;
    UCHAR req_type;
    UCHAR req_data_requested;
    PUSB_DESCRIPTOR_REQUEST req = (PUSB_DESCRIPTOR_REQUEST) malloc(irp->write.nbytes);
    
    hr = iobuf_read(&irp->write, req, irp->write.nbytes);
    if (FAILED(hr)) {
        dprintf("ES3 Dongle: es3sec_hub_descriptor_from_node Failed to read IRP %lx\n", hr);
        return hr;
    }

    req_type = req->SetupPacket.wValue >> 8;
    req_data_requested = req->SetupPacket.wValue & 0xFF;

    switch (req_type) {
        case USB_CONFIGURATION_DESCRIPTOR_TYPE:
            return es3sec_hub_get_config_descriptor(irp, req, req_data_requested);
        
        case USB_STRING_DESCRIPTOR_TYPE:
            return es3sec_hub_get_string_descriptor(irp, req, req_data_requested);
        
        default:
            dprintf("ES3 Dongle: es3sec_hub_descriptor_from_node Unknown request type %x\n", req_type);
            return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
    }
    hr = iobuf_write(&irp->read, req, irp->read.nbytes);

    if (FAILED(hr)) {
        dprintf("ES3 Dongle: es3sec_hub_descriptor_from_node Failed to write IRP %lx\n", hr);
    }
    return hr;
}

static HRESULT es3sec_hub_get_config_descriptor(struct irp *irp, PUSB_DESCRIPTOR_REQUEST req, UCHAR dataRequested)
{
    dprintf("ES3 Dongle: Get Hub Config Descriptor from Node Connection\n");

    PUSB_CONFIGURATION_DESCRIPTOR config = (PUSB_CONFIGURATION_DESCRIPTOR)malloc(sizeof(USB_CONFIGURATION_DESCRIPTOR));
    config->bLength = sizeof(*config);
    config->bDescriptorType = USB_CONFIGURATION_DESCRIPTOR_TYPE;
    config->wTotalLength = sizeof(*req);

    int cpy = memcpy_s(req->Data, 9, config, sizeof(*config));
    if (cpy) {
        dprintf("ES3 Dongle: es3sec_hub_get_config_descriptor memcpy_s failed %d\n", cpy);
    }
    return iobuf_write(&irp->read, req, irp->read.nbytes);
}

static HRESULT es3sec_hub_get_string_descriptor(struct irp *irp, PUSB_DESCRIPTOR_REQUEST req, UCHAR dataRequested)
{
    PUSB_STRING_DESCRIPTOR str_desc;
    int additional_length = sizeof(FILLER_LANGIDS);
    str_desc = (PUSB_STRING_DESCRIPTOR)malloc(sizeof(USB_STRING_DESCRIPTOR) + additional_length);
    str_desc->bDescriptorType = USB_STRING_DESCRIPTOR_TYPE;
    wcscpy_s(str_desc->bString, _countof(FILLER_LANGIDS), FILLER_LANGIDS);

    switch (dataRequested) {
        case DONGLE_STR_IDX_MANUFACTURER:
            dprintf("ES3 Dongle: Get Hub String Descriptor For Dongle Manufacturer\n");
            additional_length = sizeof(dongle_mftr);
            str_desc = (PUSB_STRING_DESCRIPTOR)realloc(str_desc, sizeof(USB_STRING_DESCRIPTOR) + additional_length);
            str_desc->bDescriptorType = USB_STRING_DESCRIPTOR_TYPE;
            wcscpy_s(str_desc->bString, _countof(dongle_mftr), dongle_mftr);
            break;
        
        case DONGLE_STR_IDX_PRODUCT:
            dprintf("ES3 Dongle: Get Hub String Descriptor For Dongle Product Name\n");
            additional_length = sizeof(dongle_product);
            str_desc = (PUSB_STRING_DESCRIPTOR)realloc(str_desc, sizeof(USB_STRING_DESCRIPTOR) + additional_length);
            str_desc->bDescriptorType = USB_STRING_DESCRIPTOR_TYPE;
            wcscpy_s(str_desc->bString, _countof(dongle_product), dongle_product);
            break;
        
        case DONGLE_STR_IDX_SERIAL_NUMBER:
            dprintf("ES3 Dongle: Get Hub String Descriptor For Dongle Serial Number\n");
            additional_length = sizeof(L"000000000000");
            str_desc = (PUSB_STRING_DESCRIPTOR)realloc(str_desc, sizeof(USB_STRING_DESCRIPTOR) + additional_length);
            str_desc->bDescriptorType = USB_STRING_DESCRIPTOR_TYPE;
            wcscpy_s(str_desc->bString, _countof(config.serial), config.serial);
            break;
    }
    
    str_desc->bLength = sizeof(USB_STRING_DESCRIPTOR) + additional_length;
    PUSB_DESCRIPTOR_REQUEST resp = (PUSB_DESCRIPTOR_REQUEST)malloc(sizeof(USB_DESCRIPTOR_REQUEST) + sizeof(USB_STRING_DESCRIPTOR) + additional_length);

    int cpy = memcpy_s(resp->Data, sizeof(USB_STRING_DESCRIPTOR) + additional_length, str_desc, sizeof(USB_STRING_DESCRIPTOR) + additional_length);
    if (cpy) {
        dprintf("ES3 Dongle: es3sec_hub_get_config_descriptor memcpy_s failed %d\n", cpy);
    }
    return iobuf_write(&irp->read, resp, sizeof(USB_DESCRIPTOR_REQUEST) + sizeof(USB_STRING_DESCRIPTOR) + additional_length);
}

static HRESULT es3sec_hub_connection_driver_key_name(struct irp *irp)
{
    size_t size_of_driver_key = sizeof(HUB_NODE_CONNECTION_DRIVER_KEY);
    ULONG actual_length = size_of_driver_key + sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME);
    HRESULT hr;

    if (irp->write.nbytes == sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME))
    {
        dprintf("ES3 Dongle: Get Hub Connection Driver Key size\n");
        USB_NODE_CONNECTION_DRIVERKEY_NAME usb_node_conn_driver_key_name;
        usb_node_conn_driver_key_name.ActualLength = actual_length;
        hr = iobuf_write(&irp->read, &usb_node_conn_driver_key_name, sizeof(usb_node_conn_driver_key_name));

        if (FAILED(hr))
        {
            dprintf("ES3 Dongle: iobuf_write failed! %lx\n", hr);
        }

        return hr;
    }

    dprintf("ES3 Dongle: Get Hub Connection Driver Key\n");

    PUSB_NODE_CONNECTION_DRIVERKEY_NAME usb_node_conn_driver_key_name = (PUSB_NODE_CONNECTION_DRIVERKEY_NAME)malloc(sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME) + actual_length);
    usb_node_conn_driver_key_name->ConnectionIndex = 1;
    usb_node_conn_driver_key_name->ActualLength = actual_length;

    errno_t err = wcscpy_s(
        usb_node_conn_driver_key_name->DriverKeyName,
        _countof(HUB_NODE_CONNECTION_DRIVER_KEY),
        HUB_NODE_CONNECTION_DRIVER_KEY);

    if (err)
    {
        dprintf("ES3 Dongle: es3sec_hub_connection_driver_key_name wcscpy_s failed with %X", err);
        return E_FAIL;
    }

    hr = iobuf_write(&irp->read, usb_node_conn_driver_key_name, actual_length);

    if (FAILED(hr))
    {
        dprintf("ES3 Dongle: iobuf_write failed! %lx\n", hr);
    }

    return hr;
}

static CONFIGRET my_CM_Get_DevNode_Registry_PropertyW(
    DEVINST dnDevInst,
    ULONG ulProperty,
    PULONG pulRegDataType,
    PVOID Buffer,
    PULONG pulLength,
    ULONG ulFlags)
{
    CONFIGRET cr = next_CM_Get_DevNode_Registry_PropertyW(dnDevInst, ulProperty, pulRegDataType, Buffer, pulLength, ulFlags);
    if (dnDevInst != HUB_DEVINST)
    {
        return cr;
    }

    switch (ulProperty)
    {
    case CM_DRP_DEVICEDESC:
        dprintf("ES3 Dongle: Get Hub Device Description\n");
        wcscpy_s(Buffer, _countof(L"Fake USB Hub"), L"Fake USB Hub");
        break;

    case CM_DRP_DRIVER:
        dprintf("ES3 Dongle: Get Hub Driver Key\n");
        wcscpy_s(Buffer, _countof(HUB_HCD_DRIVER_KEY), HUB_HCD_DRIVER_KEY);
        break;

    default:
        dprintf("ES3 Dongle: my_CM_Get_DevNode_Registry_PropertyW Unhandled property 0x%lX\n", ulProperty);
        return CR_FAILURE;
    }

    return CR_SUCCESS;
}

static CONFIGRET my_CM_Get_Child(PDEVINST pdnDevInst, DEVINST dnDevInst, ULONG ulFlags)
{
    if (dnDevInst != root_dev_inst)
    {
        return next_CM_Get_Child(pdnDevInst, dnDevInst, ulFlags);
    }

    *pdnDevInst = HUB_DEVINST;
    return CR_SUCCESS;
}