#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"

const wchar_t *DEVNAME_HUB = L"$hub";
const wchar_t *DEVNAME_DONGLE = L"$dongle";
const wchar_t HUB_DRIVER_KEY[] = L"{36fc9e60-c465-11cf-8056-444553540000}\\0001";
const DEVINST HUB_DEVINST = 573;
static struct es3sec_config config;
static HANDLE dongle_fd;
static HANDLE hub_fd;
static IID hubs_iid;
DEVINST root_dev_inst;

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_handle_dongle_irp(struct irp *irp);
static HRESULT es3sec_handle_dongle_open(struct irp *irp);
static HRESULT es3sec_handle_dongle_close(struct irp *irp);
static HRESULT es3sec_handle_dongle_ioctl(struct irp *irp);

static HRESULT es3sec_hub_handle_driverkey(struct irp *irp);

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, char *vid, char *pid)
{
    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);
    hr = iohook_push_handler(es3sec_handle_dongle_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);

    dprintf("ES3 Dongle: init with VID %s PID %s (root node DEVINST %lx)\n", vid, pid, root_dev_inst);

    memcpy(&config, cfg, sizeof(*cfg));
    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)) {
        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);
        default: dprintf("ES3 Dongle: Unknown hub IOCTL %X\n", irp->ioctl); return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
    }

    return iohook_invoke_next(irp);
}

static HRESULT es3sec_hub_handle_driverkey(struct irp *irp)
{
    int actual_length = sizeof(HUB_DRIVER_KEY);
    PUSB_HCD_DRIVERKEY_NAME usb_hcd_driver_key_name = (PUSB_HCD_DRIVERKEY_NAME) malloc(sizeof(USB_HCD_DRIVERKEY_NAME) + 2 * actual_length);

    usb_hcd_driver_key_name->ActualLength = actual_length;
    errno_t err = wcscpy_s(
        usb_hcd_driver_key_name->DriverKeyName,
        _countof(HUB_DRIVER_KEY),
        HUB_DRIVER_KEY
    );

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

    iobuf_write(&irp->read, usb_hcd_driver_key_name, sizeof(usb_hcd_driver_key_name));

    dprintf("ES3 Dongle: Get Hub Driver Key\n");
    return S_OK;
}

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

    switch (irp->op) {
    case IRP_OP_OPEN:   return es3sec_handle_dongle_open(irp);
    case IRP_OP_CLOSE:  return es3sec_handle_dongle_close(irp);
    case IRP_OP_IOCTL:  return es3sec_handle_dongle_ioctl(irp);
    default:            return HRESULT_FROM_WIN32(ERROR_INVALID_FUNCTION);
    }
}

static HRESULT es3sec_handle_dongle_open(struct irp *irp)
{
    if (!wstr_ieq(irp->open_filename, DEVNAME_DONGLE)) {
        return iohook_invoke_next(irp);
    }

    dprintf("ES3 Dongle: Open dongle\n");
    irp->fd = dongle_fd;
    return S_OK;
}

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

static HRESULT es3sec_handle_dongle_ioctl(struct irp *irp)
{
    dprintf("ES3 Dongle: IOCTL %X\n", irp->ioctl);
    return iohook_invoke_next(irp);
}

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: wcscpy_s(Buffer, _countof(L"Disk drive"), L"Disk drive"); break;

        case CM_DRP_DRIVER: wcscpy_s(Buffer, _countof(HUB_DRIVER_KEY), HUB_DRIVER_KEY); break;

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

    dprintf("ES3 Dongle: my_CM_Get_DevNode_Registry_PropertyW %lX\n", ulProperty);
    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);
    }

    dprintf("ES3 Dongle: Adding fake hub to root node\n");
    *pdnDevInst = HUB_DEVINST;
    return CR_SUCCESS;
}