
#include "StdAfx.h"
#include "DevRule.h"

#ifdef _MANAGED
#pragma managed(push, off)
#endif

BOOL APIENTRY DllMain(HMODULE hModule, DWORD dwReasonForCall, LPVOID lpReserved)
{
	switch (dwReasonForCall)
	{
	    case DLL_PROCESS_ATTACH:
	    case DLL_THREAD_ATTACH:
	    case DLL_THREAD_DETACH:
	    case DLL_PROCESS_DETACH:
		    break;
	}

    return TRUE;
}

#ifdef _MANAGED
#pragma managed(pop)
#endif

static BOOL ToggleDevice(HDEVINFO hDevInfo, DWORD StateChange, LPCTSTR DevInstanceId)
{
    BOOL bNeedReboot = FALSE;

    LOG_DEBUG(((StateChange == DICS_ENABLE) ? "Enable" : "Disable") <<
        " Device: " << DevInstanceId);

    SP_DEVINFO_DATA DevInfo;
    ZeroMemory(&DevInfo, sizeof(DevInfo));
    DevInfo.cbSize = sizeof(DevInfo);

    if (::SetupDiOpenDeviceInfo(hDevInfo, DevInstanceId, NULL, 0, &DevInfo))
    {
        SP_PROPCHANGE_PARAMS pcp;

        if (StateChange == DICS_ENABLE)
        {
            //
            // enable both on global and config-specific profile
            // do global first and see if that succeeded in enabling the device
            // (global enable doesn't mark reboot required if device is still
            // disabled on current config whereas vice-versa isn't true)
            //
            pcp.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
            pcp.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
            pcp.StateChange = DICS_ENABLE;
            pcp.Scope = DICS_FLAG_GLOBAL;
            pcp.HwProfile = 0;
            //
            // don't worry if this fails, we'll get an error when we try config-
            // specific.
            //
            if (::SetupDiSetClassInstallParams(hDevInfo, &DevInfo, &pcp.ClassInstallHeader, sizeof(pcp)) &&
                ::SetupDiCallClassInstaller(DIF_PROPERTYCHANGE, hDevInfo, &DevInfo))
            {

            }
            else
            {
                LOG_DEBUG("Setting up global profile failed but don't worry. Error = " << ::GetLastError());
            }
        }

        //
        // operate on config-specific profile
        //
        pcp.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);
        pcp.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE;
        pcp.StateChange = StateChange;
        pcp.Scope = DICS_FLAG_CONFIGSPECIFIC;
        pcp.HwProfile = 0;

        if (::SetupDiSetClassInstallParams(hDevInfo, &DevInfo, &pcp.ClassInstallHeader, sizeof(pcp)) &&
            ::SetupDiCallClassInstaller(DIF_PROPERTYCHANGE, hDevInfo, &DevInfo))
        {
            //
            // see if device needs reboot
            //
            SP_DEVINSTALL_PARAMS DevParams;
            ZeroMemory(&DevParams, sizeof(DevParams));
            DevParams.cbSize = sizeof(DevParams);

            if ((::SetupDiGetDeviceInstallParams(hDevInfo, &DevInfo, &DevParams)) &&
                (DevParams.Flags & (DI_NEEDRESTART | DI_NEEDREBOOT)))
            {
                bNeedReboot = TRUE;
            }
        }
        else
        {
            LOG_DEBUG("Setting up config-specific profile failed. Error = " << ::GetLastError());
        }
    }
    else
    {
        LOG_DEBUG("SetupDiOpenDeviceInfo failed. Error = " << ::GetLastError());
    }

    return bNeedReboot;
}

static BOOL ToggleDevices(DWORD StateChange, LPCTSTR DevInstanceId)
{
    BOOL bNeedReboot = FALSE;

    HDEVINFO hDevInfo = ::SetupDiGetClassDevs(NULL, 0, 0, DIGCF_PRESENT | DIGCF_ALLCLASSES);
    if (hDevInfo != INVALID_HANDLE_VALUE)
    {
        DWORD i = 0;

        SP_DEVINFO_DATA DeviceInfoData;
        ZeroMemory(&DeviceInfoData, sizeof(DeviceInfoData));
        DeviceInfoData.cbSize = sizeof(DeviceInfoData);

        while (::SetupDiEnumDeviceInfo(hDevInfo, i++, &DeviceInfoData))
        {
            TCHAR DevId[MAX_DEVICE_ID_LEN];
            if (::CM_Get_Device_ID_Ex(DeviceInfoData.DevInst, DevId, MAX_DEVICE_ID_LEN, 0, NULL) != CR_SUCCESS)
            {
                lstrcpy(DevId, TEXT("?"));
            }

            LOG_DEBUG("Enum Device: " << DevId);

            if (!_tcsnicmp(DevInstanceId, DevId, _tcslen(DevInstanceId)))
            {
                if (ToggleDevice(hDevInfo, StateChange, DevId) == TRUE)
                {
                    bNeedReboot = TRUE;
                }
            }
        }

        ::SetupDiDestroyDeviceInfoList(hDevInfo);
    }
    else
    {
        LOG_DEBUG("SetupDiGetClassDevs failed. Error = " << ::GetLastError());
    }

    LOG_DEBUG("Reboot required: " << (bNeedReboot ? "True" : "False"));
    return bNeedReboot;
}

DEVRULE_API BOOL DisableDevice(LPCTSTR DevInstanceId)
{
    if (DevInstanceId == NULL) return FALSE;
    LOG_DEBUG("Device Instance Id - " << DevInstanceId);
    return ToggleDevices(DICS_DISABLE, DevInstanceId);
}

DEVRULE_API BOOL EnableDevice(LPCTSTR DevInstanceId)
{
    if (DevInstanceId == NULL) return FALSE;
    LOG_DEBUG("Device Instance Id - " << DevInstanceId);
    return ToggleDevices(DICS_ENABLE, DevInstanceId);
}

DEVRULE_API BOOL IsDeviceExist(LPCTSTR DevInstanceId)
{
    BOOL bExist = FALSE;

    if (DevInstanceId == NULL) return FALSE;
    LOG_DEBUG("Device Instance Id - " << DevInstanceId);

    HDEVINFO hDevInfo = ::SetupDiGetClassDevs(NULL, 0, 0, DIGCF_PRESENT | DIGCF_ALLCLASSES);
    if (hDevInfo != INVALID_HANDLE_VALUE)
    {
        DWORD i = 0;

        SP_DEVINFO_DATA DeviceInfoData;
        ZeroMemory(&DeviceInfoData, sizeof(DeviceInfoData));
        DeviceInfoData.cbSize = sizeof(DeviceInfoData);

        while (::SetupDiEnumDeviceInfo(hDevInfo, i++, &DeviceInfoData))
        {
            TCHAR DevId[MAX_DEVICE_ID_LEN];
            if (::CM_Get_Device_ID_Ex(DeviceInfoData.DevInst, DevId, MAX_DEVICE_ID_LEN, 0, NULL) != CR_SUCCESS)
            {
                lstrcpy(DevId, TEXT("?"));
            }
            
            LOG_DEBUG("Enum Device: " << DevId);

            if (!_tcsnicmp(DevInstanceId, DevId, _tcslen(DevInstanceId)))
            {
                bExist = TRUE;
                break;
            }
        }

        ::SetupDiDestroyDeviceInfoList(hDevInfo);
    }
    else
    {
        LOG_DEBUG("SetupDiGetClassDevs failed. Error = " << ::GetLastError());
    }

    LOG_DEBUG("Result: " << (bExist ? "True" : "False"));
    return bExist;
}

DEVRULE_API BOOL RescanDevices()
{
    BOOL bRescan = FALSE;
    
    DEVINST hDevRoot;
    if (::CM_Locate_DevNode_Ex(&hDevRoot, NULL, CM_LOCATE_DEVNODE_NORMAL, NULL) == CR_SUCCESS)
    {
        if (::CM_Reenumerate_DevNode_Ex(hDevRoot, 0, NULL) == CR_SUCCESS)
        {
            bRescan = TRUE;
        }
    }

    LOG_DEBUG("Result: " << (bRescan ? "True" : "False"));
    return bRescan;
}
