// ==========================================================================
// Author:  Yee Hsu
// Date:    6/7/2008
// File:    WMI.cpp
//
// Desc:    WMI Wrapper that wraps around the Windows Management Instrumental
//          API. This wrapper puts the procedures into an Object Oriented
//          Model for easy usage.
// ==========================================================================

#include "stdafx.h"
#include "Wmi.h"
#include "IntCommon.h"

// ==========================================================================
// Identifier:  WMI()
//
// Description: WMI Constructor
// ==========================================================================

WMI::WMI()
{
}

// ==========================================================================
// Identifier:  WMI()
//
// Description: WMI Destructor
// ==========================================================================

WMI::~WMI()
{
}

// ==========================================================================
// Identifier:  Connect()
//
// Description: Connect to WMI DB by root path
// ==========================================================================

bool WMI::Connect(const OLECHAR* pRootPath)
{
    // Step 1: --------------------------------------------------
    // Initialize COM. ------------------------------------------

    //hres =  CoInitializeEx(NULL, COINIT_MULTITHREADED); 
    hres = CoInitialize(NULL);

    if(hres == RPC_E_CHANGED_MODE)
    {
        //cout <<"Cannot Change com apartment model!";
        return false;
    }

    if (FAILED(hres))
    {
        //cout << "Failed to initialize COM library. Error code = 0x" << hex << hres << endl;
        return false;                  // Program has failed.
    }

    // Step 2: --------------------------------------------------
    // Set general COM security levels --------------------------
    // Note: If you are using Windows 2000, you need to specify -
    // the default authentication credentials for a user by using
    // a SOLE_AUTHENTICATION_LIST structure in the pAuthList ----
    // parameter of CoInitializeSecurity ------------------------

    hres =  CoInitializeSecurity(
        NULL, 
        -1,                          // COM authentication
        NULL,                        // Authentication services
        NULL,                        // Reserved
        RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication 
        RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation  
        NULL,                        // Authentication info
        EOAC_NONE,                   // Additional capabilities 
        NULL                         // Reserved
        );


    if( (hres != RPC_E_TOO_LATE) && FAILED(hres) )
    {
        cout << "Failed to initialize security. Error code = 0x" << hex << hres << endl;
        this->Release();
        return false;                    // Program has failed.
    }

    // Step 3: ---------------------------------------------------
    // Obtain the initial locator to WMI -------------------------

    hres = CoCreateInstance(
        CLSID_WbemLocator,             
        0, 
        CLSCTX_INPROC_SERVER, 
        IID_IWbemLocator, (LPVOID *) &pLoc);

    if (FAILED(hres))
    {
        //cout << "Failed to create IWbemLocator object." << " Err code = 0x" << hex << hres << endl;
        this->Release();
        return false;                 // Program has failed.
    }

    // Step 4: -----------------------------------------------------
    // Connect to WMI through the IWbemLocator::ConnectServer method
    BSTR WMI_DB = SysAllocString(pRootPath);

    // Connect to the root\cimv2 namespace with
    // the current user and obtain pointer pSvc
    // to make IWbemServices calls.
    hres = pLoc->ConnectServer(
        WMI_DB,                  // Object path of WMI namespace
        NULL,                    // User name. NULL = current user
        NULL,                    // User password. NULL = current
        0,                       // Locale. NULL indicates current
        NULL,                    // Security flags.
        0,                       // Authority (e.g. Kerberos)
        0,                       // Context object 
        &pSvc                    // pointer to IWbemServices proxy
        );

    if (FAILED(hres))
    {
        //cout << "Could not connect. Error code = 0x" << hex << hres << endl;
        //pLoc->Release();     
        this->Release();
        return false;                // Program has failed.
    }

    //cout << "Connected to " << (_bstr_t) WMI_DB << " WMI namespace" << endl;

    // Step 5: --------------------------------------------------
    // Set security levels on the proxy -------------------------

    hres = CoSetProxyBlanket(
        pSvc,                        // Indicates the proxy to set
        RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
        RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
        NULL,                        // Server principal name 
        RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx 
        RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
        NULL,                        // client identity
        EOAC_NONE                    // proxy capabilities 
        );

    if (FAILED(hres))
    {
        //cout << "Could not set proxy blanket. Error code = 0x" << hex << hres << endl;
        //pSvc->Release();
        //pLoc->Release();     
        this->Release();
        return false;               // Program has failed.
    }
    return true;
}

// ==========================================================================
// Identifier:  QueryEx()
//
// Description: Construct the WMI SQL query
// ==========================================================================

bool WMI::QueryEx(const OLECHAR* pSQLQuery)
{
    // Step 5.5: --------------------------------------------------
    // Assign SQL statement

    this->pSQLQuery = (OLECHAR*) pSQLQuery;
    return true;
}

// ==========================================================================
// Identifier:  GetStringValue()
//
// Description: Gets WMI string value
// ==========================================================================

bool WMI::GetStringValue(LPCWSTR wszName, std::string& sString)
{
    bool bRet   = this->GetWmiValue(wszName, WMI::QUERY_STRING);
    sString     = this->sRetString;

    return bRet;
}

// ==========================================================================
// Identifier:  GetBoolValue()
//
// Description: Gets WMI bool value
// ==========================================================================

bool WMI::GetBoolValue(LPCWSTR wszName, bool& bBool)
{
    bool bRet   = this->GetWmiValue(wszName, WMI::QUERY_BOOL);
    bBool       = this->bRetBool;

    return bRet;
}

// ==========================================================================
// Identifier:  GetIntValue()
//
// Description: Gets WMI int value
// ==========================================================================

bool WMI::GetIntValue(LPCWSTR wszName, int& nInt)
{
    bool bRet   = this->GetWmiValue(wszName, WMI::QUERY_SINT32);
    nInt        = this->nRetInt;

    return bRet;
}

// ==========================================================================
// Identifier:  GetWmiValue()
//
// Description: Gets WMI value
// ==========================================================================

bool WMI::GetWmiValue(LPCWSTR wszName, eWmiInfo eInfo)
{
    // Step 6: --------------------------------------------------
    // Use the IWbemServices pointer to make requests of WMI ----

    // For example, get the name of the operating system
    BSTR Language   = SysAllocString(L"WQL");
    BSTR Query      = SysAllocString(pSQLQuery);

    hres = pSvc->ExecQuery(
        Language, 
        Query,
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, 
        NULL,
        &pEnumerator);

    if (FAILED(hres))
    {
        //cout << "Query for operating system name failed." << " Error code = 0x"  << hex << hres << endl;
        //pSvc->Release();
        //pLoc->Release();
        this->Release();
        return false;               // Program has failed.
    }

    // Step 7: -------------------------------------------------
    // Get the data from the query in step 6 -------------------

    ULONG   uReturn = 0;
    bool    bInfoExist = false;

    while (pEnumerator)
    {
        HRESULT hr = pEnumerator->Next(WBEM_INFINITE, 1L, &pclsObj, &uReturn);

        if (0 == uReturn)
        {
            break;
        }

        VARIANT vtProp;
        VariantInit(&vtProp);

        switch (eInfo)
        {
            case WMI::QUERY_STRING:
                vtProp.bstrVal = SysAllocString(L"");
                hr = pclsObj->Get(wszName, 0, &vtProp, 0, 0);
                sRetString.clear();
                sRetString = (_bstr_t) vtProp.bstrVal;
                break;

            case WMI::QUERY_BOOL:
                vtProp.boolVal = VARIANT_FALSE;
                hr = pclsObj->Get(wszName, 0, &vtProp, 0, 0);
                bRetBool = ((vtProp.boolVal == VARIANT_TRUE) ? true : false);
                break;

            case WMI::QUERY_SINT32:
                hr = pclsObj->Get(wszName, 0, &vtProp, 0, 0);
                nRetInt = vtProp.intVal;
                break;
        }
        VariantClear(&vtProp);
        bInfoExist = true;
    }
    return bInfoExist;
}

// ==========================================================================
// Identifier:  Release()
//
// Description: Releases all WMI objects
// ==========================================================================

bool WMI::Release()
{
    /*
    //Cleanup
    //========


    if (pSvc)
    pSvc->Release();

    if (pLoc)
    pLoc->Release();

    if (pEnumerator)
    pEnumerator->Release();

    if (pclsObj)
    pclsObj->Release();
    */

    CoUninitialize();

    return true;
}

// ==========================================================================
// Class: ExWMI
//
// Description: WMI extension. An extended version of WMI for even easier
//              access. 
// ==========================================================================

// static counter
int ExWMI::nCounter = 0;

// ==========================================================================
// Identifier:  ExWMI()
//
// Description: ExWMI Constructor
// ==========================================================================

ExWMI::ExWMI()
{
    this->pRootPath = NULL;
    this->pSQLQuery = NULL;
    this->mrswmi.clear();
}

// ==========================================================================
// Identifier:  ExWMI()
//
// Description: ExWMI Destructor
// ==========================================================================

ExWMI::~ExWMI()
{
    this->pRootPath = NULL;
    this->pSQLQuery = NULL;
    this->mrswmi.clear();
}

// ==========================================================================
// Identifier:  GetType()
//
// Description: Get type of WMI variable
// ==========================================================================

string ExWMI::GetType(const int nType)
{
    string sRet = "";

    switch (nType)
    {
        case ExWMI::QUERY_STRING:
            sRet = "string";
            break;

        case ExWMI::QUERY_BOOL:
            sRet = "boolean";
            break;

        case ExWMI::QUERY_SINT32:
            sRet = "sint32";
            break;

        default:
            sRet = "unknown";
            break;
    }
    return sRet;
}

// ==========================================================================
// Identifier:  GetEntryFor()
//
// Description: Gets the WMI entry
// ==========================================================================

void ExWMI::GetEntryFor(LPCWSTR sName, const int nType)
{
    rsWMI rswmi;

    rswmi.sName = sName;
    rswmi.nType = nType;
    rswmi.sValue.clear();

    this->mrswmi[nCounter++] = rswmi;
}

// ==========================================================================
// Identifier:  ExecuteGetWmi()
//
// Description: Executes the WMI query and sets the value for return
// ==========================================================================

void ExWMI::ExecuteGetWmi()
{
    if (this->Connect(this->pRootPath))
    {
        if (this->QueryEx(this->pSQLQuery))
        {
            for (MapRSWMI::iterator iwmi = this->mrswmi.begin(); iwmi != this->mrswmi.end(); iwmi++)
            {
                switch (iwmi->second.nType)
                {
                    case WMI::QUERY_STRING:
                        this->GetStringValue(iwmi->second.sName, iwmi->second.sValue);
                        break;

                    case WMI::QUERY_BOOL:
                        bool bBool;
                        this->GetBoolValue(iwmi->second.sName, bBool);
                        iwmi->second.sValue = (bBool == true ? "true" : "false");
                        break;

                    case WMI::QUERY_SINT32:
                        int nInt;
                        char szInt[32];

                        this->GetIntValue(iwmi->second.sName, nInt);
                        sprintf(szInt, "%d", nInt);
                        iwmi->second.sValue = szInt;
                        break;

                    default:
                        iwmi->second.sValue.clear();
                        break;
                }
            }
        }
    }
}