#include "precomp.h"


NTSTATUS UDF_CallUSBDI(IN PDEVICE_OBJECT pDevObj, IN PVOID UrbEtc)
{
    IO_STATUS_BLOCK IoStatus;
    KEVENT event;
    NTSTATUS status;
    PIRP Irp = NULL;
    PIO_STACK_LOCATION NextIrpStack = NULL;

    //g_Temp = pDevObj;
    //g_Temp2 = UrbEtc;

    // Initialise IRP completion event
    KeInitializeEvent(&event, NotificationEvent, FALSE);

    // Build Internal IOCTL IRP
    Irp = IoBuildDeviceIoControlRequest(
        IOCTL_INTERNAL_USB_SUBMIT_URB, pDevObj,
        NULL, 0, // Input buffer
        NULL, 0, // Output buffer
        TRUE, &event, &IoStatus);

    // Get IRP stack location for next driver down (already set up)
    NextIrpStack = IoGetNextIrpStackLocation(Irp);
    // Store pointer to the URB etc
    NextIrpStack->Parameters.Others.Argument1 = UrbEtc;
    NextIrpStack->Parameters.Others.Argument2 = (PVOID)0;

    // Call the driver and wait for completion if necessary
    status = IoCallDriver(pDevObj, Irp);
    if (status == STATUS_PENDING)
    {
        KLogEx(DEBUG_TRACE_ERROR, "waiting for URB completion\n");
        status = KeWaitForSingleObject(&event, Suspended, KernelMode, FALSE, NULL);
        status = IoStatus.Status;
    }

    // return IRP completion status
    KLogEx(DEBUG_TRACE_ERROR, "returned %x\n", status);
    return status;
}

PURB AllocUrb(USHORT Size)
{
    PURB urb = ExAllocatePoolWithTag(NonPagedPool, Size, 'PURB');
    if (urb) RtlZeroMemory(urb, Size);
    return urb;
}

VOID FreeUsbDeviceInfo(IN PUSB_PARSED_INFO pInfo)
{
    if (!pInfo) return;

    // 문자열 메모리 해제
    if (pInfo->ManufacturerStr) ExFreePool(pInfo->ManufacturerStr);
    if (pInfo->ProductStr)      ExFreePool(pInfo->ProductStr);
    if (pInfo->SerialNumberStr) ExFreePool(pInfo->SerialNumberStr);

    // Configuration Descriptor 해제
    if (pInfo->ConfigDesc)      ExFreePool(pInfo->ConfigDesc);

    // 구조체 자체 해제
    ExFreePool(pInfo);
}

// ---------------------------------------------------------------------------
// 1. Device Descriptor 가져오기
// ---------------------------------------------------------------------------
NTSTATUS GetUsbDeviceDescriptor(
    IN PDEVICE_OBJECT DeviceObject,
    OUT PUSB_DEVICE_DESCRIPTOR* ppDescriptor
)
{
    PURB urb = NULL;
    PUSB_DEVICE_DESCRIPTOR descriptor = NULL;
    NTSTATUS status;

    urb = AllocUrb(sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));
    if (!urb) return STATUS_INSUFFICIENT_RESOURCES;

    descriptor = ExAllocatePoolWithTag(NonPagedPool, sizeof(USB_DEVICE_DESCRIPTOR), 'DST');
    if (!descriptor) {
        ExFreePool(urb);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    RtlZeroMemory(descriptor, sizeof(USB_DEVICE_DESCRIPTOR));

    UsbBuildGetDescriptorRequest(
        urb,
        sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST),
        USB_DEVICE_DESCRIPTOR_TYPE,
        0, 0,
        descriptor, NULL,
        sizeof(USB_DEVICE_DESCRIPTOR), NULL
    );

    status = UDF_CallUSBDI(DeviceObject, urb);

    if (!NT_SUCCESS(status)) {
        ExFreePool(descriptor);
        *ppDescriptor = NULL;
    }
    else {
        *ppDescriptor = descriptor;
    }

    ExFreePool(urb);
    return status;
}

// ---------------------------------------------------------------------------
// 2. String Descriptor 가져오기 (특정 Index, Language ID)
// ---------------------------------------------------------------------------
NTSTATUS GetUsbStringDescriptor(
    IN PDEVICE_OBJECT DeviceObject,
    IN UCHAR Index,
    IN USHORT LangId,
    OUT PWCHAR* ppStringData
)
{
    PURB urb = NULL;
    PUSB_STRING_DESCRIPTOR pStrDesc = NULL;
    NTSTATUS status;
    ULONG size = 256; // 넉넉하게 잡음

    if (Index == 0) return STATUS_INVALID_PARAMETER;

    urb = AllocUrb(sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));
    if (!urb) return STATUS_INSUFFICIENT_RESOURCES;

    pStrDesc = ExAllocatePoolWithTag(NonPagedPool, size, 'STR');
    if (!pStrDesc) {
        ExFreePool(urb);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    RtlZeroMemory(pStrDesc, size);

    UsbBuildGetDescriptorRequest(
        urb,
        sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST),
        USB_STRING_DESCRIPTOR_TYPE,
        Index,
        LangId,
        pStrDesc, NULL,
        size, NULL
    );

    status = UDF_CallUSBDI(DeviceObject, urb);

    if (NT_SUCCESS(status) && pStrDesc->bLength > 2) {
        // bString은 WCHAR 배열입니다. NULL Termination을 위해 별도 버퍼에 복사
        ULONG strLenBytes = pStrDesc->bLength - 2;
        PWCHAR outBuf = ExAllocatePoolWithTag(NonPagedPool, strLenBytes + sizeof(WCHAR), 'STR2');

        if (outBuf) {
            RtlZeroMemory(outBuf, strLenBytes + sizeof(WCHAR));
            RtlCopyMemory(outBuf, pStrDesc->bString, strLenBytes);
            *ppStringData = outBuf;
        }
        else {
            status = STATUS_INSUFFICIENT_RESOURCES;
        }
    }
    else {
        *ppStringData = NULL;
        status = STATUS_UNSUCCESSFUL;
    }

    ExFreePool(pStrDesc);
    ExFreePool(urb);
    return status;
}

// ---------------------------------------------------------------------------
// 3. Configuration Descriptor 전체 가져오기 (Interface, Endpoint 포함)
// ---------------------------------------------------------------------------
NTSTATUS GetUsbFullConfigDescriptor(
    IN PDEVICE_OBJECT DeviceObject,
    OUT PUSB_CONFIGURATION_DESCRIPTOR* ppConfigDesc
)
{
    PURB urb = NULL;
    PUSB_CONFIGURATION_DESCRIPTOR pHeader = NULL;
    PUSB_CONFIGURATION_DESCRIPTOR pFullDesc = NULL;
    NTSTATUS status;
    ULONG fullSize = 0;

    // 1단계: 헤더(9바이트)만 먼저 읽어서 전체 크기(wTotalLength) 파악
    urb = AllocUrb(sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));
    if (!urb) return STATUS_INSUFFICIENT_RESOURCES;

    pHeader = ExAllocatePoolWithTag(NonPagedPool, sizeof(USB_CONFIGURATION_DESCRIPTOR), 'CFG');
    if (!pHeader) {
        ExFreePool(urb);
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    UsbBuildGetDescriptorRequest(
        urb, sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST),
        USB_CONFIGURATION_DESCRIPTOR_TYPE, 0, 0,
        pHeader, NULL, sizeof(USB_CONFIGURATION_DESCRIPTOR), NULL
    );

    status = UDF_CallUSBDI(DeviceObject, urb);
    if (!NT_SUCCESS(status)) {
        ExFreePool(pHeader);
        ExFreePool(urb);
        return status;
    }

    fullSize = pHeader->wTotalLength;
    ExFreePool(pHeader); // 헤더는 이제 필요 없음 (크기 알았으므로)

    // 2단계: 전체 크기만큼 할당 후 다시 요청
    pFullDesc = ExAllocatePoolWithTag(NonPagedPool, fullSize, 'CFG');
    if (!pFullDesc) {
        ExFreePool(urb);
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // URB 재사용을 위해 ZeroMemory (버퍼 포인터 등이 바뀌므로)
    RtlZeroMemory(urb, sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));

    UsbBuildGetDescriptorRequest(
        urb, sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST),
        USB_CONFIGURATION_DESCRIPTOR_TYPE, 0, 0,
        pFullDesc, NULL, fullSize, NULL
    );

    status = UDF_CallUSBDI(DeviceObject, urb);
    if (NT_SUCCESS(status)) {
        *ppConfigDesc = pFullDesc;
    }
    else {
        ExFreePool(pFullDesc);
        *ppConfigDesc = NULL;
    }

    ExFreePool(urb);
    return status;
}

// ---------------------------------------------------------------------------
// Main Logic: 모든 정보를 수집하고 출력(Logging)하는 함수
// ---------------------------------------------------------------------------
VOID GetAllUsbDeviceInfo(PDEVICE_OBJECT DeviceObject)
{
    NTSTATUS status;
    PUSB_DEVICE_DESCRIPTOR pDevDesc = NULL;
    PUSB_CONFIGURATION_DESCRIPTOR pConfigDesc = NULL;
    USHORT LangId = 0x0409; // Default English

    KLogEx(DEBUG_TRACE_INFO, "===== START USB INFO COLLECTION =====\n");

    // 1. Device Descriptor 수집
    status = GetUsbDeviceDescriptor(DeviceObject, &pDevDesc);
    if (NT_SUCCESS(status) && pDevDesc)
    {
        KLogEx(DEBUG_TRACE_INFO, "[Device Descriptor]\n");
        KLogEx(DEBUG_TRACE_INFO, "  bLength: 0x%X\n", pDevDesc->bLength);
        KLogEx(DEBUG_TRACE_INFO, "  bDescriptorType: 0x%X\n", pDevDesc->bDescriptorType);
        KLogEx(DEBUG_TRACE_INFO, "  bcdUSB: 0x%X\n", pDevDesc->bcdUSB);
        KLogEx(DEBUG_TRACE_INFO, "  idVendor: 0x%04X\n", pDevDesc->idVendor);
        KLogEx(DEBUG_TRACE_INFO, "  idProduct: 0x%04X\n", pDevDesc->idProduct);
        KLogEx(DEBUG_TRACE_INFO, "  bcdDevice: 0x%X\n", pDevDesc->bcdDevice);
        KLogEx(DEBUG_TRACE_INFO, "  bNumConfigurations: 0x%X\n", pDevDesc->bNumConfigurations);

        // String Descriptor 수집 (Manufacturer, Product, Serial)
        // 실제로는 String Index 0을 먼저 호출하여 지원 언어를 확인해야 하지만, 
        // 여기서는 편의상 0x0409(English)를 시도합니다.

        if (pDevDesc->iManufacturer) 
        {
            PWCHAR pStr = NULL;
            if (NT_SUCCESS(GetUsbStringDescriptor(DeviceObject, pDevDesc->iManufacturer, LangId, &pStr)) && pStr) {
                KLogEx(DEBUG_TRACE_INFO, "  Manufacturer: %ws\n", pStr);
                ExFreePool(pStr);
            }
        }

        if (pDevDesc->iProduct) 
        {
            PWCHAR pStr = NULL;
            if (NT_SUCCESS(GetUsbStringDescriptor(DeviceObject, pDevDesc->iProduct, LangId, &pStr)) && pStr) {
                KLogEx(DEBUG_TRACE_INFO, "  Product: %ws\n", pStr);
                ExFreePool(pStr);
            }
        }

        if (pDevDesc->iSerialNumber)
        {
            PWCHAR pStr = NULL;
            if (NT_SUCCESS(GetUsbStringDescriptor(DeviceObject, pDevDesc->iSerialNumber, LangId, &pStr)) && pStr) {
                KLogEx(DEBUG_TRACE_INFO, "  SerialNumber: %ws\n", pStr);
                ExFreePool(pStr);
            }
        }
    }

    // 2. Configuration Descriptor (Full) 수집 및 파싱
    status = GetUsbFullConfigDescriptor(DeviceObject, &pConfigDesc);
    if (NT_SUCCESS(status) && pConfigDesc) 
    {
        PUCHAR pCurr = (PUCHAR)pConfigDesc;
        PUCHAR pEnd = pCurr + pConfigDesc->wTotalLength;

        KLogEx(DEBUG_TRACE_INFO, "[Configuration Descriptor]\n");
        KLogEx(DEBUG_TRACE_INFO, "  wTotalLength: 0x%X\n", pConfigDesc->wTotalLength);
        KLogEx(DEBUG_TRACE_INFO, "  bNumInterfaces: 0x%X\n", pConfigDesc->bNumInterfaces);
        KLogEx(DEBUG_TRACE_INFO, "  bConfigurationValue: 0x%X\n", pConfigDesc->bConfigurationValue);
        KLogEx(DEBUG_TRACE_INFO, "  MaxPower: 0x%X", pConfigDesc->MaxPower);

        // Loop를 돌면서 Interface 및 Endpoint 파싱
        // 첫 번째 Descriptor는 Configuration이므로 건너뜀
        pCurr += pConfigDesc->bLength;

        while (pCurr < pEnd)
        {
            PUSB_COMMON_DESCRIPTOR pCommon = (PUSB_COMMON_DESCRIPTOR)pCurr;

            if (pCommon->bLength == 0) break; // 에러 방지

            if (pCommon->bDescriptorType == USB_INTERFACE_DESCRIPTOR_TYPE)
            {
                PUSB_INTERFACE_DESCRIPTOR pIntf = (PUSB_INTERFACE_DESCRIPTOR)pCommon;
                KLogEx(DEBUG_TRACE_INFO, "[Interface Descriptor]\n");
                KLogEx(DEBUG_TRACE_INFO, "  bInterfaceNumber: 0x%X\n", pIntf->bInterfaceNumber);
                KLogEx(DEBUG_TRACE_INFO, "  bAlternateSetting: 0x%X\n", pIntf->bAlternateSetting);
                KLogEx(DEBUG_TRACE_INFO, "  bNumEndpoints: 0x%X\n", pIntf->bNumEndpoints);
                KLogEx(DEBUG_TRACE_INFO, "  bInterfaceClass: 0x%X\n", pIntf->bInterfaceClass);
                KLogEx(DEBUG_TRACE_INFO, "  bInterfaceSubClass: 0x%X\n", pIntf->bInterfaceSubClass);
                KLogEx(DEBUG_TRACE_INFO, "  bInterfaceProtocol: 0x%X\n", pIntf->bInterfaceProtocol);
            }
            else if (pCommon->bDescriptorType == USB_ENDPOINT_DESCRIPTOR_TYPE) 
            {
                PUSB_ENDPOINT_DESCRIPTOR pEp = (PUSB_ENDPOINT_DESCRIPTOR)pCommon;
                KLogEx(DEBUG_TRACE_INFO, "[Endpoint Descriptor]\n");
                KLogEx(DEBUG_TRACE_INFO, "  bEndpointAddress: 0x%X\n", pEp->bEndpointAddress);
                KLogEx(DEBUG_TRACE_INFO, "  bmAttributes: 0x%X\n", pEp->bmAttributes);
                KLogEx(DEBUG_TRACE_INFO, "  wMaxPacketSize: 0x%X\n", pEp->wMaxPacketSize);
                KLogEx(DEBUG_TRACE_INFO, "  bInterval: 0x%X\n", pEp->bInterval);
            }

            pCurr += pCommon->bLength;
        }
    }

    // 메모리 정리
    if (pDevDesc) ExFreePool(pDevDesc);
    if (pConfigDesc) ExFreePool(pConfigDesc);

    KLogEx(DEBUG_TRACE_INFO, "===== END USB INFO COLLECTION =====\n");
}


NTSTATUS CollectUsbDeviceInfo(
    IN PDEVICE_OBJECT DeviceObject,
    OUT PUSB_PARSED_INFO* ppInfo
)
{
    NTSTATUS status;
    PUSB_PARSED_INFO pInfo = NULL;
    PUSB_DEVICE_DESCRIPTOR pDevDesc = NULL;
    USHORT LangId = 0x0409; // English (US)

    *ppInfo = NULL;

    if (KeGetCurrentIrql() > APC_LEVEL)
    {
        KLogEx(DEBUG_TRACE_ERROR, "not APC_LEVEL\n");
        return STATUS_INVALID_DEVICE_STATE;
    }

    // 1. 구조체 할당
    pInfo = ExAllocatePoolWithTag(NonPagedPool, sizeof(USB_PARSED_INFO), 'INFO');
    if (!pInfo) return STATUS_INSUFFICIENT_RESOURCES;
    RtlZeroMemory(pInfo, sizeof(USB_PARSED_INFO));

    // 2. Device Descriptor 수집
    status = GetUsbDeviceDescriptor(DeviceObject, &pDevDesc);
    if (!NT_SUCCESS(status) || !pDevDesc) {
        FreeUsbDeviceInfo(pInfo);
        return status;
    }

    // 구조체에 복사 후 임시 버퍼 해제
    RtlCopyMemory(&pInfo->DeviceDesc, pDevDesc, sizeof(USB_DEVICE_DESCRIPTOR));
    ExFreePool(pDevDesc);

    // 3. String Descriptor 수집 (Manufacturer, Product, Serial)
    if (pInfo->DeviceDesc.iManufacturer) {
        GetUsbStringDescriptor(DeviceObject, pInfo->DeviceDesc.iManufacturer, LangId, &pInfo->ManufacturerStr);
    }
    if (pInfo->DeviceDesc.iProduct) {
        GetUsbStringDescriptor(DeviceObject, pInfo->DeviceDesc.iProduct, LangId, &pInfo->ProductStr);
    }
    if (pInfo->DeviceDesc.iSerialNumber) {
        GetUsbStringDescriptor(DeviceObject, pInfo->DeviceDesc.iSerialNumber, LangId, &pInfo->SerialNumberStr);
    }

    // 4. Configuration Descriptor (Full) 수집
    // ConfigDesc는 내부적으로 할당된 메모리를 그대로 포인터로 연결
    status = GetUsbFullConfigDescriptor(DeviceObject, &pInfo->ConfigDesc);
    if (NT_SUCCESS(status) && pInfo->ConfigDesc) 
    {
        pInfo->ConfigDescSize = pInfo->ConfigDesc->wTotalLength;

        PUCHAR pCurr = (PUCHAR)pInfo->ConfigDesc;
        PUCHAR pEnd = pCurr + pInfo->ConfigDescSize;
        LONG currentInterfaceIndex = -1; // 현재 파싱 중인 인터페이스 번호

        if (pCurr + sizeof(USB_CONFIGURATION_DESCRIPTOR) < pEnd)
            pCurr += ((PUSB_COMMON_DESCRIPTOR)pCurr)->bLength;

        while (pCurr < pEnd)
        {
            PUSB_COMMON_DESCRIPTOR pCommon = (PUSB_COMMON_DESCRIPTOR)pCurr;

            // 유효성 체크
            if (pCommon->bLength == 0 || pCurr + pCommon->bLength > pEnd) 
                break;

            // Interface Descriptor 확인
            if (pCommon->bDescriptorType == USB_INTERFACE_DESCRIPTOR_TYPE) {
                PUSB_INTERFACE_DESCRIPTOR pIntf = (PUSB_INTERFACE_DESCRIPTOR)pCommon;
                currentInterfaceIndex = pIntf->bInterfaceNumber;
            }
            // [B] HID Descriptor 발견 (0x21)
            else if (pCommon->bDescriptorType == USB_DESCRIPTOR_TYPE_HID) 
            {
                PUSB_HID_DESCRIPTOR pHid = (PUSB_HID_DESCRIPTOR)pCommon;

                // 저장 공간이 있고, 유효한 인터페이스 내부에 있는 경우
                if (pInfo->HidInterfaceCount < MAX_HID_INTERFACES && currentInterfaceIndex >= 0) 
                {
                    ULONG idx = pInfo->HidInterfaceCount;

                    pInfo->HidEntries[idx].InterfaceNumber = (UCHAR)currentInterfaceIndex;
                    pInfo->HidEntries[idx].HidDesc = pHid;

                    pInfo->HidInterfaceCount++;

                    KLogEx(DEBUG_TRACE_INFO, "  Found HID Desc at Interface %d (Count: %d)\n",
                        currentInterfaceIndex, pInfo->HidInterfaceCount);
                }
            }

            pCurr += pCommon->bLength;
        }
    }
    else 
    {
        pInfo->ConfigDesc = NULL;
        pInfo->ConfigDescSize = 0;
    }

    // 성공 시 포인터 반환
    *ppInfo = pInfo;
    return STATUS_SUCCESS;
}