namespace Encryption {

const int RC4BOXSIZE = 256;

class RC4 {
public:
    RC4(const std::string &keyStr) : m_SBox(nullptr) {
        const char *keydata = keyStr.c_str();
        const int keySize = keyStr.size();

        m_SBox = new unsigned char[RC4BOXSIZE];
        for (int i = 0; i < RC4BOXSIZE; ++i) {
            m_SBox[i] = i;
        }

        for (int i = 0, j = 0; i < RC4BOXSIZE; ++i) {
            j = (j + m_SBox[i] + keydata[i % keySize] + 20) % RC4BOXSIZE;
            std::swap(m_SBox[i], m_SBox[j]); // 搅乱box数据
        }
    }
    ~RC4() {
        delete[] m_SBox;
        m_SBox = nullptr;
    }

    std::string encry(const std::string& inputStr) {
        return encry(inputStr.c_str(), inputStr.size());
    }

    char *encry(const char *inputDatas, int len) {
        if (!inputDatas || len == 0 || !m_SBox)
            return "";

        unsigned char *cpBox = new unsigned char[RC4BOXSIZE];
        memcpy(cpBox, m_SBox, RC4BOXSIZE);

        std::string outputStr;
        char *outDatas = new char[len];
        int i = 0, j = 0;
        for (int k = 0; k < len; ++k) {
            i = (i + 1) % RC4BOXSIZE;
            j = (j + cpBox[i]) % RC4BOXSIZE;
            std::swap(cpBox[i], cpBox[j]);

            //(m_SBox[i] + cpBox[j]) % 256  的值是一致的，异或反转
            outDatas[k] = inputDatas[k] ^ cpBox[(m_SBox[i] + cpBox[j]) % 256]; 
        }
        return outDatas;
    }
private:
    unsigned char *m_SBox;
};

}

void testRc4() {
    string keyStr = "this is a test";
    Encryption::RC4* rc4 = new Encryption::RC4(keyStr);

    const char* datas = "start encryption";
    // 这里应当是个固定的大小，加密后长度不变
    size_t len = datas, strlen(datas);
    char *newData = rc4->encry(datas, len);  
    char *newData1 = rc4->encry(newData, len);


    assert(datas != newData1, "error");

    delete[] newData;
    delete[] newData1;
    delete rc4;
}

    inline unsigned long long qNextPowerOfTwo(unsigned int v)
    {
        v |= v >> 1;
        v |= v >> 2;
        v |= v >> 4;
        v |= v >> 8;
        v |= v >> 16;
        ++v;
        return v;
    }

    inline unsigned long long qNextPowerOfTwo(unsigned long long v)
    {
        v |= v >> 1;    // 步骤1
        v |= v >> 2;    // 步骤2
        v |= v >> 4;    // 步骤3
        v |= v >> 8;
        v |= v >> 16;
        v |= v >> 32; 
        ++v;            //  最终
        return v;
    } 

    const unsigned align = 8;
    unsigned int RoundUp(unsigned int v) 
    {
        return (v + align - 1) & ~(align - 1);
    }

    size_t countOneV1(int v) {
        size_t c = 0;
        while (v) 
        {
            if (v & 1) ++c;
            v >>= 1;
        }
        return c;
    }

    size_t countOneV2(int v) {
        size_t c = 0;
        while (v) 
        {
            ++c;
            v = v & (v - 1);
        }
        return c;
    }

    size_t countZero(int v) {
        size_t c = 0;
        while(v) 
        {
            ++c;
            v = v | (v + 1);
        }
        return c;
    }

HRESULT hres;

// Step 1: --------------------------------------------------
// Initialize COM. ------------------------------------------

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

// Step 2: --------------------------------------------------
// Set general COM security levels --------------------------
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
);

// note: 此处步骤1，2 如果已经初始化过，第二次会false，所以视情况调用

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

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

IWbemLocator* pLoc = NULL;

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;
    CoUninitialize();
    return 1;                 // Program has failed.
}

// Step 4: -----------------------------------------------------
// Connect to WMI through the IWbemLocator::ConnectServer method

IWbemServices* pSvc = NULL;

// Connect to the root\cimv2 namespace with
// the current user and obtain pointer pSvc
// to make IWbemServices calls.
hres = pLoc->ConnectServer(
    _bstr_t(L"ROOT\\CIMV2"), // 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 (for example, Kerberos)
    0,                       // Context object 
    &pSvc                    // pointer to IWbemServices proxy
);

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

cout << "Connected to ROOT\\CIMV2 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
    nullptr,                        // Server principal name 
    RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx 
    RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
    nullptr,                        // 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();
    CoUninitialize();
    return 1;               // Program has failed.
}

// Step 6: --------------------------------------------------
// Use the IWbemServices pointer to make requests of WMI ----

// For example, get the name of the operating system
IEnumWbemClassObject* pEnumerator = nullptr;

hres = pSvc->ExecQuery(
    bstr_t("WQL"),
    bstr_t("SELECT * FROM Win32_DiskDrive"),
    WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
    NULL,
    &pEnumerator);

// note: 
// 1. 这里通过查询得到DEVICEID = \\\\.\\PHYSICALDRIVE0, 然而还需要将\ 替换为 \\

// 2. WQL语句中“=”左右不能有空格，否则ExecQuery之后会查询失败
//    如：ASSOCIATORS OF {Win32_DiskDrive.DeviceID=\" xxx \"}就会失败

// 3. gcc环境中 bstr_t("xxx") 可能会编译报错, 使用 "BSTR xxx = SysAllocString(L"sql .. ") " 和 SysFreeString(xxx) 组合代替
//    如：bstr_t("WQL") ——> BSTR strQueryLanguage = SysAllocString(L"WQL"); SysFreeString(strQueryLanguage) 

//hres = pSvc->ExecQuery(
//    bstr_t("WQL"),
//    bstr_t("ASSOCIATORS OF {Win32_DiskDrive.DeviceID=\"\\\\\\\\.\\\\PHYSICALDRIVE0\"} WHERE AssocClass=Win32_DiskDriveToDiskPartition"),
//    //bstr_t("ASSOCIATORS OF {Win32_LogicalDisk.DeviceID=\"C:\"}"),
//    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();
    CoUninitialize();
    return 1;               // Program has failed.
}

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

IWbemClassObject* pclsObj = nullptr;
ULONG uReturn = 0;

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

    if (0 == uReturn)
    {
        break;
    }

    VARIANT vtProp;

    // Get the value of the Name property
    hr = pclsObj->Get(L"DeviceID", 0, &vtProp, 0, 0);
    //wcout << "DeviceID : " << vtProp.bstrVal << endl;
    wstring deviceID = vtProp.bstrVal;
    wcout << deviceID << endl;

    hr = pclsObj->Get(L"SerialNumber", 0, &vtProp, 0, 0);
    wstring serialNumber = vtProp.bstrVal;
    wcout << "SerialNumber : " << serialNumber << endl;

    VariantClear(&vtProp);
    pclsObj->Release();

    replace(deviceID, L"\\", L"\\\\");

    wstring query;
    query.append(L"ASSOCIATORS OF {Win32_DiskDrive.DeviceID=\"");
    query.append(deviceID);
    query.append(L"\"} WHERE AssocClass=Win32_DiskDriveToDiskPartition");

    IEnumWbemClassObject* partitionEnumerator = nullptr;
    hres = pSvc->ExecQuery(
        bstr_t("WQL"),
        bstr_t(query.c_str()),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        nullptr,
        &partitionEnumerator);

    IWbemClassObject* diskPartition = nullptr;
    while (partitionEnumerator)
    {
        HRESULT hr = partitionEnumerator->Next(WBEM_INFINITE, 1,
            &diskPartition, &uReturn);

        if (0 == uReturn)
        {
            break;
        }
        VARIANT vtProp;

        // Get the value of the Name property
        hr = diskPartition->Get(L"DeviceID", 0, &vtProp, 0, 0);
        wstring diskPartitionDeviceID = vtProp.bstrVal;
        wcout << "diskPartitionDeviceID : " << diskPartitionDeviceID << endl;

        VariantClear(&vtProp);

        diskPartition->Release();
        diskPartition = nullptr;

        wstring query;
        query.append(L"ASSOCIATORS OF {Win32_DiskPartition.DeviceID=\"");
        query.append(diskPartitionDeviceID);
        query.append(L"\"} WHERE AssocClass=Win32_LogicalDiskToPartition");

        IEnumWbemClassObject* logicalDiskEnumerator = nullptr;
        hres = pSvc->ExecQuery(
            bstr_t("WQL"),
            bstr_t(query.c_str()),
            WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
            nullptr,
            &logicalDiskEnumerator);

        IWbemClassObject* logicalPartition = nullptr;
        while (logicalDiskEnumerator) {

            HRESULT hr = logicalDiskEnumerator->Next(WBEM_INFINITE, 1,
                &logicalPartition, &uReturn);

            if (0 == uReturn)
            {
                break;
            }

            VARIANT vtProp;

            hr = logicalPartition->Get(L"DeviceID", 0, &vtProp, 0, 0);
            wstring logicalDiskDeviceID = vtProp.bstrVal;
            wcout << "logicalDiskDeviceID : " << logicalDiskDeviceID << endl; // 通过上方得到的SerialNumber和此处的logicalDiskDeviceID建立map就完成了步骤2

            VariantClear(&vtProp);

            logicalPartition->Release();
            logicalPartition = nullptr;

        }
    }
}

// Cleanup
// ========

pSvc->Release();
pLoc->Release();
pEnumerator->Release();
CoUninitialize();

return 0;   // Program successfully completed.