#include "precomp.h"

// ---------------------------------------------------------------------------
// [헬퍼 함수] URB 전송 및 메모리 해제 로직 (기존 코드 유지)
// ---------------------------------------------------------------------------
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;
    
    KeInitializeEvent(&event, NotificationEvent, FALSE);

    Irp = IoBuildDeviceIoControlRequest(
        IOCTL_INTERNAL_USB_SUBMIT_URB, pDevObj,
        NULL, 0, NULL, 0, TRUE, &event, &IoStatus);

    NextIrpStack = IoGetNextIrpStackLocation(Irp);
    NextIrpStack->Parameters.Others.Argument1 = UrbEtc;
    NextIrpStack->Parameters.Others.Argument2 = (PVOID)0;

    status = IoCallDriver(pDevObj, Irp);
    if (status == STATUS_PENDING)
    {
        status = KeWaitForSingleObject(&event, Suspended, KernelMode, FALSE, NULL);
        status = IoStatus.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);
    if (pInfo->ConfigDesc)
        ExFreePool(pInfo->ConfigDesc);

    ExFreePool(pInfo);
}

NTSTATUS GetUsbDeviceDescriptor(IN PDEVICE_OBJECT DeviceObject, OUT PUSB_DEVICE_DESCRIPTOR* ppDescriptor)
{
    PURB urb = AllocUrb(sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));
    if (!urb) return STATUS_INSUFFICIENT_RESOURCES;

    PUSB_DEVICE_DESCRIPTOR 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);

    NTSTATUS status = UDF_CallUSBDI(DeviceObject, urb);
    if (!NT_SUCCESS(status)) { ExFreePool(descriptor); *ppDescriptor = NULL; }
    else { *ppDescriptor = descriptor; }

    ExFreePool(urb);
    return status;
}

NTSTATUS GetUsbStringDescriptor(IN PDEVICE_OBJECT DeviceObject, IN UCHAR Index, IN USHORT LangId, OUT PWCHAR* ppStringData)
{
    if (Index == 0) return STATUS_INVALID_PARAMETER;

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

    ULONG size = 256;
    PUSB_STRING_DESCRIPTOR 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);

    NTSTATUS status = UDF_CallUSBDI(DeviceObject, urb);
    if (NT_SUCCESS(status) && pStrDesc->bLength > 2) {
        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;
}

NTSTATUS GetUsbFullConfigDescriptor(IN PDEVICE_OBJECT DeviceObject, OUT PUSB_CONFIGURATION_DESCRIPTOR* ppConfigDesc)
{
    PURB urb = AllocUrb(sizeof(struct _URB_CONTROL_DESCRIPTOR_REQUEST));
    if (!urb) return STATUS_INSUFFICIENT_RESOURCES;

    PUSB_CONFIGURATION_DESCRIPTOR 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);

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

    ULONG fullSize = pHeader->wTotalLength;
    ExFreePool(pHeader);

    PUSB_CONFIGURATION_DESCRIPTOR pFullDesc = ExAllocatePoolWithTag(NonPagedPool, fullSize, 'CFG');
    if (!pFullDesc) { ExFreePool(urb); return STATUS_INSUFFICIENT_RESOURCES; }

    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;
}

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) 
        return STATUS_INVALID_DEVICE_STATE;

    pInfo = ExAllocatePoolWithTag(NonPagedPool, sizeof(USB_PARSED_INFO), 'INFO');
    if (!pInfo) return STATUS_INSUFFICIENT_RESOURCES;
    RtlZeroMemory(pInfo, sizeof(USB_PARSED_INFO));

    status = GetUsbDeviceDescriptor(DeviceObject, &pDevDesc);
    if (!NT_SUCCESS(status) || !pDevDesc) 
    { 
        FreeUsbDeviceInfo(pInfo); 
        return status; 
    }

    RtlCopyMemory(&pInfo->DeviceDesc, pDevDesc, sizeof(USB_DEVICE_DESCRIPTOR));
    ExFreePool(pDevDesc);

    // String Descriptor 수집
    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);

    
    // Configuration Descriptor (Full) 수집 및 원-패스 파싱
    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;
        
        KLogEx(DEBUG_TRACE_INFO, "[Configuration Descriptor] bNumInterfaces: 0x%X", pInfo->ConfigDesc->bNumInterfaces);

        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;              
                UCHAR cls = pIntf->bInterfaceClass;
                UCHAR subCls = pIntf->bInterfaceSubClass;
                UCHAR proto = pIntf->bInterfaceProtocol;

                KLogEx(DEBUG_TRACE_INFO, "[Interface Descriptor] bInterfaceNumber: 0x%X, bAlternateSetting: 0x%X, bNumEndpoints: 0x%X, bInterfaceClass: 0x%X, bInterfaceSubClass: 0x%X, bInterfaceProtocol: 0x%X",
                    pIntf->bInterfaceNumber,
                    pIntf->bAlternateSetting,
                    pIntf->bNumEndpoints,
                    pIntf->bInterfaceClass,
                    pIntf->bInterfaceSubClass,
                    pIntf->bInterfaceProtocol
                    );

                currentInterfaceIndex = pIntf->bInterfaceNumber;
                pInfo->FoundClassCount++;
                if (pInfo->FoundClassCount > 1) 
                    pInfo->IsCompositeDevice = TRUE;

                switch (cls)
                {
                case 0x01: pInfo->IsAudio = TRUE; break;
                case 0x02: pInfo->IsCDC = TRUE; break;
                case 0x06:
                case 0x0E: pInfo->IsImageOrVideo = TRUE; break;
                case 0x07: pInfo->IsPrinter = TRUE; break;
                case 0x08: pInfo->IsMassStorage = TRUE; break;
                case 0x09: pInfo->IsHub = TRUE; break;
                case 0x03: // HID (키보드, 마우스, 기타)
                    if (subCls == 0x01)
                    {
                        if (proto == 0x01)      pInfo->IsBootKeyboard = TRUE;
                        else if (proto == 0x02) pInfo->IsBootMouse = TRUE;
                        else                    pInfo->IsAbnormalHID = TRUE;
                    }
                    else if (subCls == 0x00) 
                    { 
                        pInfo->IsCustomHID = TRUE;
                    }
                    else 
                    { 
                        pInfo->IsAbnormalHID = TRUE; 
                    }

                    if (!pInfo->IsBootKeyboard && !pInfo->IsBootMouse) 
                    {
                        pInfo->OtherHidSubClass = subCls;
                        pInfo->OtherHidProtocol = proto;
                    }
                    break;
                }
            }
            // =================================================================
            // Endpoint Descriptor 파싱 (BadUSB 위협 분석용)
            // =================================================================
            else if (pCommon->bDescriptorType == USB_ENDPOINT_DESCRIPTOR_TYPE)
            {
                PUSB_ENDPOINT_DESCRIPTOR pEp = (PUSB_ENDPOINT_DESCRIPTOR)pCommon;               
                UCHAR epType = pEp->bmAttributes & 0x03; // USB_ENDPOINT_TYPE_MASK

			    KLogEx(DEBUG_TRACE_INFO, "[Endpoint Descriptor]  bEndpointAddress: 0x%X, bmAttributes: 0x%X, wMaxPacketSize: 0x%X, bInterval: 0x%X\n", 
				    pEp->bEndpointAddress,
				    pEp->bmAttributes,
				    pEp->wMaxPacketSize,
				    pEp->bInterval
			    );

                pInfo->TotalEndpoints++;
                if (epType == USB_ENDPOINT_TYPE_INTERRUPT)
                {
                    pInfo->InterruptEndpoints++;
                }
                else if (epType == USB_ENDPOINT_TYPE_BULK) 
                {
                    pInfo->BulkEndpoints++;
                }
            }
            // =================================================================
            // [C] HID Descriptor 별도 보관 (필요 시 Report 분석용)
            // =================================================================
            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++;
                }
            }

            pCurr += pCommon->bLength;
        }

        // =================================================================
        // [결론] 최종 BadUSB 휴리스틱 판별 (루프 종료 후)
        // =================================================================
        // 1. Bulk와 Interrupt가 혼용되었는지 (아이폰, 프린터 등 복합장치 특성)
        if (pInfo->InterruptEndpoints > 0 && pInfo->BulkEndpoints > 0) 
        {
            pInfo->HasMixedEndpoints = TRUE;
        }

        BOOLEAN hasHID = pInfo->IsBootKeyboard || pInfo->IsBootMouse;
        BOOLEAN hasBothHID = pInfo->IsBootKeyboard && pInfo->IsBootMouse;

        // 2. 최종 판별 로직 (안전한 신형 로직 적용)
        // 조건: 자기가 키보드나 마우스라고 주장하는데, 아래 중 하나라도 해당하면 BadUSB로 차단!
        // - 키보드와 마우스가 동시에 존재 (해커들이 자주 씀)
        // - 인터럽트 엔드포인트가 3개 이상으로 비정상적으로 많음
        // - 엔드포인트가 전혀 없음 (가짜 장치)
        // - 키보드라면서 대용량 전송용 BULK 엔드포인트가 섞여 있음
        if (hasBothHID || (hasHID && (pInfo->InterruptEndpoints >= 3 || pInfo->TotalEndpoints == 0 || pInfo->HasMixedEndpoints)))
        {
            pInfo->IsBadUSB = TRUE;
            KLogEx(DEBUG_TRACE_ERROR, "[!!!] MALICIOUS BAD-USB DETECTED! (IntEps: %d, BulkEps: %d)\n",
                pInfo->InterruptEndpoints, pInfo->BulkEndpoints);
        }
    }
    else
    {
        pInfo->ConfigDesc = NULL;
        pInfo->ConfigDescSize = 0;
    }

    *ppInfo = pInfo;
    return STATUS_SUCCESS;
}


VOID PrintUsbParsedInfo(IN PUSB_PARSED_INFO pInfo)
{
    if (!pInfo)
    {
        KLogEx(DEBUG_TRACE_ERROR, "[-] PrintUsbParsedInfo: pInfo is NULL\n");
        return;
    }

    // HID 정보 출력
    if (pInfo->IsBootKeyboard)  KLogEx(DEBUG_TRACE_INFO, "  - IsBootKeyboard     : TRUE\n");
    if (pInfo->IsBootMouse)     KLogEx(DEBUG_TRACE_INFO, "  - IsBootMouse        : TRUE\n");
    if (pInfo->IsCustomHID)     KLogEx(DEBUG_TRACE_INFO, "  - IsCustomHID        : TRUE\n");
    if (pInfo->IsAbnormalHID)   KLogEx(DEBUG_TRACE_INFO, "  - IsAbnormalHID      : TRUE\n");

    // 기타 HID일 경우에만 서브클래스와 프로토콜 값을 출력
    if (pInfo->IsCustomHID || pInfo->IsAbnormalHID) {
        KLogEx(DEBUG_TRACE_INFO, "    > OtherHidSubClass : 0x%02X\n", pInfo->OtherHidSubClass);
        KLogEx(DEBUG_TRACE_INFO, "    > OtherHidProtocol : 0x%02X\n", pInfo->OtherHidProtocol);
    }

    // 디바이스 클래스 정보 출력
    if (pInfo->IsMassStorage)   KLogEx(DEBUG_TRACE_INFO, "  - IsMassStorage      : TRUE\n");
    if (pInfo->IsAudio)         KLogEx(DEBUG_TRACE_INFO, "  - IsAudio            : TRUE\n");
    if (pInfo->IsPrinter)       KLogEx(DEBUG_TRACE_INFO, "  - IsPrinter          : TRUE\n");
    if (pInfo->IsImageOrVideo)  KLogEx(DEBUG_TRACE_INFO, "  - IsImageOrVideo     : TRUE\n");
    if (pInfo->IsHub)           KLogEx(DEBUG_TRACE_INFO, "  - IsHub              : TRUE\n");
    if (pInfo->IsCDC)           KLogEx(DEBUG_TRACE_INFO, "  - IsCDC              : TRUE\n");
    if (pInfo->IsCompositeDevice) KLogEx(DEBUG_TRACE_INFO, " - IsCompositeDevice  : TRUE\n");

    KLogEx(DEBUG_TRACE_INFO, "  - FoundClassCount    : %lu\n", pInfo->FoundClassCount);

    // 3. BadUSB 및 엔드포인트 정보 출력
    KLogEx(DEBUG_TRACE_INFO, "\n[BadUSB & Endpoint Information]\n");
    KLogEx(DEBUG_TRACE_INFO, "  - TotalEndpoints     : %lu\n", pInfo->TotalEndpoints);
    KLogEx(DEBUG_TRACE_INFO, "  - InterruptEndpoints : %lu\n", pInfo->InterruptEndpoints);
    KLogEx(DEBUG_TRACE_INFO, "  - BulkEndpoints      : %lu\n", pInfo->BulkEndpoints);

    if (pInfo->HasMixedEndpoints)
        KLogEx(DEBUG_TRACE_INFO, "  - HasMixedEndpoints  : TRUE\n");

    // BadUSB 탐지 여부는 눈에 잘 띄게 출력
    if (pInfo->IsBadUSB) {
        KLogEx(DEBUG_TRACE_INFO, "  - IsBadUSB           : *** TRUE (MALICIOUS DEVICE DETECTED!) ***\n");
    }
}