usb_power.cpp

#include "usb_power.h"
#include "logger.h"
#include <SetupAPI.h>
#include <usbioctl.h>
#include <usb.h>
#include <devguid.h>
#include <cfgmgr32.h>
#include <sstream>
#include <iomanip>
#include <algorithm>

#pragma comment(lib, "setupapi.lib")

// GUID for USB Hub interface
DEFINE_GUID(GUID_DEVINTERFACE_USB_HUB,
    0xf18a0e88, 0xc30c, 0x11d0, 0x88, 0x15, 0x00, 0xa0, 0xc9, 0x06, 0xbe, 0xd8);

std::vector<std::wstring> USBPowerMonitor::getUSBHubs() {
    std::vector<std::wstring> hubs;

    HDEVINFO deviceInfoSet = SetupDiGetClassDevsW(
        &GUID_DEVINTERFACE_USB_HUB,
        nullptr,
        nullptr,
        DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);

    if (deviceInfoSet == INVALID_HANDLE_VALUE) {
        return hubs;
    }

    SP_DEVICE_INTERFACE_DATA deviceInterfaceData;
    deviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

    for (DWORD i = 0; SetupDiEnumDeviceInterfaces(deviceInfoSet, nullptr,
        &GUID_DEVINTERFACE_USB_HUB, i, &deviceInterfaceData); i++) {

        DWORD requiredSize = 0;
        SetupDiGetDeviceInterfaceDetailW(deviceInfoSet, &deviceInterfaceData,
            nullptr, 0, &requiredSize, nullptr);

        if (requiredSize == 0) continue;

        std::vector<BYTE> buffer(requiredSize);
        PSP_DEVICE_INTERFACE_DETAIL_DATA_W detailData =
            reinterpret_cast<PSP_DEVICE_INTERFACE_DETAIL_DATA_W>(buffer.data());
        detailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);

        if (SetupDiGetDeviceInterfaceDetailW(deviceInfoSet, &deviceInterfaceData,
            detailData, requiredSize, nullptr, nullptr)) {
            hubs.push_back(detailData->DevicePath);
        }
    }

    SetupDiDestroyDeviceInfoList(deviceInfoSet);
    return hubs;
}

USBPortInfo USBPowerMonitor::getPortInfo(HANDLE hubHandle, ULONG portIndex) {
    USBPortInfo info;
    info.portNumber = portIndex;

    // Allocate buffer for connection info
    ULONG nBytes = sizeof(USB_NODE_CONNECTION_INFORMATION_EX);
    std::vector<BYTE> buffer(nBytes);

    PUSB_NODE_CONNECTION_INFORMATION_EX connectionInfo =
        reinterpret_cast<PUSB_NODE_CONNECTION_INFORMATION_EX>(buffer.data());
    connectionInfo->ConnectionIndex = portIndex;

    DWORD bytesReturned = 0;
    if (!DeviceIoControl(hubHandle,
        IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX,
        connectionInfo, nBytes,
        connectionInfo, nBytes,
        &bytesReturned, nullptr)) {
        return info;
    }

    info.isConnected = (connectionInfo->ConnectionStatus == DeviceConnected);

    if (info.isConnected) {
        info.vendorId = connectionInfo->DeviceDescriptor.idVendor;
        info.productId = connectionInfo->DeviceDescriptor.idProduct;

        // Power is in 2mA units in USB descriptor
        info.maxPowerMa = connectionInfo->DeviceDescriptor.bMaxPacketSize0;

        // Check USB speed for USB3
        info.isUSB3 = (connectionInfo->Speed >= UsbHighSpeed);

        // Get actual power from configuration descriptor
        // Need to query the configuration descriptor for accurate power
        ULONG configDescSize = sizeof(USB_DESCRIPTOR_REQUEST) + sizeof(USB_CONFIGURATION_DESCRIPTOR);
        std::vector<BYTE> configBuffer(configDescSize + 256);

        PUSB_DESCRIPTOR_REQUEST descRequest =
            reinterpret_cast<PUSB_DESCRIPTOR_REQUEST>(configBuffer.data());
        descRequest->ConnectionIndex = portIndex;
        descRequest->SetupPacket.wValue = (USB_CONFIGURATION_DESCRIPTOR_TYPE << 8) | 0;
        descRequest->SetupPacket.wIndex = 0;
        descRequest->SetupPacket.wLength = 256;

        if (DeviceIoControl(hubHandle,
            IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION,
            descRequest, configDescSize + 256,
            descRequest, configDescSize + 256,
            &bytesReturned, nullptr)) {

            PUSB_CONFIGURATION_DESCRIPTOR configDesc =
                reinterpret_cast<PUSB_CONFIGURATION_DESCRIPTOR>(descRequest->Data);

            // bMaxPower is in 2mA units for USB 2.0, 8mA units for USB 3.0
            if (info.isUSB3) {
                info.maxPowerMa = configDesc->MaxPower * 8;
            } else {
                info.maxPowerMa = configDesc->MaxPower * 2;
            }
        }

        // Try to get device name from string descriptor
        if (connectionInfo->DeviceDescriptor.iProduct != 0) {
            info.deviceName = getStringDescriptor(hubHandle, portIndex,
                connectionInfo->DeviceDescriptor.iProduct);
        }

        if (info.deviceName.empty()) {
            std::wostringstream oss;
            oss << L"USB Device (VID:" << std::uppercase << std::hex
                << std::setfill(L'0') << std::setw(4) << info.vendorId
                << L" PID:" << std::setw(4) << info.productId << L")";
            info.deviceName = oss.str();
        }
    }

    return info;
}

std::wstring USBPowerMonitor::getStringDescriptor(HANDLE hubHandle, ULONG portIndex, UCHAR stringIndex) {
    if (stringIndex == 0) return L"";

    ULONG bufferSize = sizeof(USB_DESCRIPTOR_REQUEST) + sizeof(USB_STRING_DESCRIPTOR) + 256;
    std::vector<BYTE> buffer(bufferSize);

    PUSB_DESCRIPTOR_REQUEST request = reinterpret_cast<PUSB_DESCRIPTOR_REQUEST>(buffer.data());
    request->ConnectionIndex = portIndex;
    request->SetupPacket.wValue = (USB_STRING_DESCRIPTOR_TYPE << 8) | stringIndex;
    request->SetupPacket.wIndex = 0x0409; // English
    request->SetupPacket.wLength = 256;

    DWORD bytesReturned = 0;
    if (DeviceIoControl(hubHandle,
        IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION,
        request, bufferSize,
        request, bufferSize,
        &bytesReturned, nullptr)) {

        PUSB_STRING_DESCRIPTOR stringDesc =
            reinterpret_cast<PUSB_STRING_DESCRIPTOR>(request->Data);

        if (stringDesc->bLength > 2) {
            return std::wstring(stringDesc->bString,
                (stringDesc->bLength - 2) / sizeof(WCHAR));
        }
    }

    return L"";
}

USBHubInfo USBPowerMonitor::getHubInfo(const std::wstring& hubPath) {
    USBHubInfo hub;
    hub.hubPath = hubPath;

    // Extract hub name from path
    size_t lastSlash = hubPath.rfind(L'\\');
    if (lastSlash != std::wstring::npos) {
        hub.hubName = hubPath.substr(lastSlash + 1);
    } else {
        hub.hubName = hubPath;
    }

    // Check if root hub
    std::wstring upperPath = hubPath;
    std::transform(upperPath.begin(), upperPath.end(), upperPath.begin(), ::towupper);
    hub.isRootHub = (upperPath.find(L"ROOT") != std::wstring::npos);

    // Open hub
    HANDLE hubHandle = CreateFileW(
        hubPath.c_str(),
        GENERIC_WRITE,
        FILE_SHARE_WRITE,
        nullptr,
        OPEN_EXISTING,
        0,
        nullptr);

    if (hubHandle == INVALID_HANDLE_VALUE) {
        return hub;
    }

    // Get hub node information
    USB_NODE_INFORMATION nodeInfo;
    DWORD bytesReturned = 0;

    if (DeviceIoControl(hubHandle,
        IOCTL_USB_GET_NODE_INFORMATION,
        &nodeInfo, sizeof(nodeInfo),
        &nodeInfo, sizeof(nodeInfo),
        &bytesReturned, nullptr)) {

        if (nodeInfo.NodeType == UsbHub) {
            hub.portCount = nodeInfo.u.HubInformation.HubDescriptor.bNumberOfPorts;
            hub.isUSB3 = nodeInfo.u.HubInformation.HubIsBusPowered ? false : true;
            hub.maxPowerMa = hub.isUSB3 ? 900 : 500;
        }
    }

    // Enumerate all ports
    for (ULONG i = 1; i <= hub.portCount; i++) {
        USBPortInfo portInfo = getPortInfo(hubHandle, i);
        if (portInfo.isConnected) {
            hub.ports.push_back(portInfo);
            hub.totalPowerMa += portInfo.maxPowerMa;
        }
    }

    CloseHandle(hubHandle);
    return hub;
}

std::vector<USBHubInfo> USBPowerMonitor::enumerateUSBTopology() {
    std::vector<USBHubInfo> topology;
    auto hubs = getUSBHubs();

    for (const auto& hubPath : hubs) {
        USBHubInfo hubInfo = getHubInfo(hubPath);
        if (hubInfo.portCount > 0) {
            topology.push_back(hubInfo);
        }
    }

    return topology;
}

void USBPowerMonitor::printTopology() {
    auto topology = enumerateUSBTopology();

    LOG_SECTION(L"USB POWER TOPOLOGY");

    if (topology.empty()) {
        LOG_INFO(L"No USB hubs found");
        return;
    }

    for (const auto& hub : topology) {
        if (hub.ports.empty()) continue;

        std::wostringstream hubHeader;
        hubHeader << (hub.isRootHub ? L"[ROOT] " : L"[HUB]  ");
        hubHeader << L"USB" << (hub.isUSB3 ? L"3" : L"2") << L" Hub";
        hubHeader << L" - " << hub.ports.size() << L"/" << hub.portCount << L" ports used";
        hubHeader << L" - Power: " << hub.totalPowerMa << L"mA";

        // Warning if over 80% capacity
        uint32_t maxTotal = hub.portCount * hub.maxPowerMa;
        if (hub.totalPowerMa > maxTotal * 0.8) {
            hubHeader << L" [WARNING: HIGH LOAD]";
        }

        LOG_INFO(hubHeader.str());

        // List devices on this hub
        for (const auto& port : hub.ports) {
            std::wostringstream portLine;
            portLine << L"  Port " << std::setw(2) << port.portNumber << L": ";
            portLine << port.deviceName;
            portLine << L" - " << port.maxPowerMa << L"mA";

            if (port.maxPowerMa >= 500) {
                portLine << L" [HIGH]";
            }

            Logger::instance().log(LogLevel::USB_DEVICE, portLine.str());
        }
    }
}

void USBPowerMonitor::printPowerSummary() {
    auto topology = enumerateUSBTopology();

    LOG_SECTION(L"USB POWER SUMMARY");

    uint32_t totalPower = 0;
    uint32_t deviceCount = 0;

    // Collect all devices with power
    std::vector<std::pair<std::wstring, uint32_t>> devices;

    for (const auto& hub : topology) {
        for (const auto& port : hub.ports) {
            totalPower += port.maxPowerMa;
            deviceCount++;
            devices.push_back({port.deviceName, port.maxPowerMa});
        }
    }

    // Sort by power consumption (highest first)
    std::sort(devices.begin(), devices.end(),
        [](const auto& a, const auto& b) { return a.second > b.second; });

    // Print top consumers
    LOG_INFO(L"Top power consumers:");
    int count = 0;
    for (const auto& device : devices) {
        if (count++ >= 10) break;

        std::wostringstream line;
        line << L"  " << std::setw(4) << device.second << L"mA - " << device.first;
        LOG_USB(line.str());
    }

    std::wostringstream summary;
    summary << L"Total: " << deviceCount << L" devices, " << totalPower << L"mA";
    LOG_INFO(summary.str());
}