UCHEBA/5 СЕМ/БСИТ/2 ЛАБА/WMI.cpp

#define _SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING

#include <iostream>
#include <comdef.h>
#include <Wbemidl.h>
#include <Windows.h>
#include <filesystem>
#include <fstream>
#include <map>
#include <string>
#include <codecvt>

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

using namespace std;
namespace fs = filesystem;

wstring replaceBackslashes(const wstring& input);
void folderTraversal(const wstring& path, IWbemServices* pSvc);
void checkFiles(const wstring& path, IWbemServices* pSvc, map<wstring, wstring>& fileMap);
void checkDirs(const wstring& path, IWbemServices* pSvc, map<wstring, wstring>& fileMap);
void checkObjects(const wstring& path, IWbemServices* pSvc,
    const wstring& query, const wstring& objectType, map<wstring, wstring>& fileMap);

#define CHECK_SECONDS_INTERVAL 2

int main(int argc, char** argv) {
    setlocale(LC_ALL, "");
    HRESULT hres;
    IWbemServices* pSvc = nullptr;
    IWbemLocator* pLoc = nullptr;

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

    // Initialize COM Security
    hres = CoInitializeSecurity(
        NULL, -1, NULL, NULL,
        RPC_C_AUTHN_LEVEL_DEFAULT,
        RPC_C_IMP_LEVEL_IMPERSONATE,
        NULL, EOAC_NONE, NULL);

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

    // Create the IWbemLocator object
    hres = CoCreateInstance(CLSID_WbemLocator, 0, 
        CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*)&pLoc);
    if (FAILED(hres)) {
        cout << "Failed to create IWbemLocator object. " << "Error code = 0x" << hex << hres << endl;
        return 1;
    }

    // Connect to WMI
    hres = pLoc->ConnectServer(
        _bstr_t(L"ROOT\\CIMV2"),
        NULL, NULL, 0, NULL,
        0, 0, &pSvc);

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

    wstring path = L"C:\\Users\\user\\Desktop\\WMI\\config.txt"; // Нужный путь для мониторинга
    folderTraversal(path, pSvc);
    pLoc->Release();
    pSvc->Release();
    CoUninitialize();
    return 0;
}

wstring replaceBackslashes(const wstring& input) {
    wstring output;
    size_t pos = 0;
    size_t found;

    while ((found = input.find(L"\\", pos)) != wstring::npos) {
        output.append(input, pos, found - pos); // Добавляем часть строки до найденного символа
        output += L"\\\\"; // Добавляем заменяемую строку
        pos = found + 1; // Продолжаем поиск с символа после найденного
    }
    output.append(input, pos, input.length() - pos); // Добавляем оставшуюся часть строки

    return output;
}

// ПРИМЕР ЗАПРОСА
// SELECT * FROM CIM_Directory WHERE Drive = 'C:' AND Path = '\\Users\\user\\Desktop\\тестовая директория\\.vs\\'

// Говно, надо через MultiByteToWideChar, но лень
wstring utf8_to_wstring(const string& str) {
    wstring_convert<codecvt_utf8<wchar_t>> converter;
    return converter.from_bytes(str);
}

// сюда должен идти путь до конфига
void folderTraversal(const wstring& path, IWbemServices* pSvc) {

    ifstream configFile(path); // Открываем файл конфигурации
    if (!configFile.is_open()) {
        wcerr << L"Не удалось открыть файл конфигурации: " << path << endl;
        return;
    }

    // если хочу отслеживать файл, то надо сразу в карту его закинуть и не кидать в вектор,
    // потому что этот вектор это вектор путей к директориям
    map<wstring, wstring> prevFileMap; // Карта предыдущих файлов
    vector<wstring> paths; // Вектор для хранения путей к папкам
    string line;

    // Чтение путей из файла
    while (getline(configFile, line)) {
        wstring wline = utf8_to_wstring(line);
        // wcout << wline << endl;
        paths.push_back(wline);
    }
    configFile.close(); // Закрываем файл

    for (ULONGLONG i = 0; ; ++i) {
        try {
            map<wstring, wstring> currFileMap; // Карта текущих файлов

            // Обработка файлов и директорий текущей директории

            for (const auto& path : paths) {
                checkFiles(path, pSvc, currFileMap);
                checkDirs(path, pSvc, currFileMap);
            }

            if (i) { // чтобы не сработало для нулевого i (выведутся все файлы и папки при запуске)
                // Сравнение карт
                for (const auto& file : currFileMap) {
                    if (prevFileMap.find(file.first) == prevFileMap.end()) {
                        wcout << L"Новый файл: " << file.first << endl;
                    }
                }

                for (const auto& file : prevFileMap) {
                    if (currFileMap.find(file.first) == currFileMap.end()) {
                        wcout << L"Удаленный файл: " << file.first << endl;
                    }
                }
            }
            prevFileMap = currFileMap; // Обновляем предыдущую карту
        }
        catch (const exception& e) {
            wcerr << L"Общая ошибка: " << e.what() << endl;
            break;
        }

        Sleep(CHECK_SECONDS_INTERVAL * 1000);
    }
}

void checkFiles(const wstring& path, IWbemServices* pSvc, map<wstring, wstring>& fileMap) {
    const wstring query = L"SELECT * FROM CIM_DataFile WHERE Drive = 'C:' AND Path = '" + path + L"'";
    checkObjects(path, pSvc, query, L"F", fileMap);
}

void checkDirs(const wstring& path, IWbemServices* pSvc, map<wstring, wstring>& fileMap) {
    const wstring query = L"SELECT * FROM CIM_Directory WHERE Drive = 'C:' AND Path = '" + path + L"'";
    checkObjects(path, pSvc, query, L"D", fileMap);
}

void checkObjects(const wstring& path, IWbemServices* pSvc,
    const wstring& query, const wstring& objectType, map<wstring, wstring>& fileMap) {

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

    if (FAILED(hres)) {
        cout << "ExecQuery failed. Error code = 0x" << hex << hres << endl;
        return;
    }

    IWbemClassObject* pclsObj = nullptr;
    ULONG uReturn = 0, count = 0;

    while (pEnumerator) {
        HRESULT hr = pEnumerator->Next(WBEM_INFINITE, 1, &pclsObj, &uReturn);
        if (0 == uReturn) {
            break; // Нет больше объектов
        }

        // Получение имени объекта
        VARIANT vtObjectName;
        VariantInit(&vtObjectName); // Инициализируем VARIANT

        if (SUCCEEDED(pclsObj->Get(L"Name", 0, &vtObjectName, 0, 0))) {
            wstring fileName = vtObjectName.bstrVal;
            fileMap[fileName] = L""; // Добавляем файл в карту
            // wcout << "FILENAME: " << vtObjectName.bstrVal << " SIZE: " << GetFileSizeOrFolderSize(vtObjectName.bstrVal) << endl;
            // Если это директория, рекурсивно проверяем её содержимое
            if (objectType == L"D") {
                wstring subDirPath = (wstring)vtObjectName.bstrVal + L"\\";

                checkDirs(replaceBackslashes(subDirPath.substr(2)), pSvc, fileMap);
                checkFiles(replaceBackslashes(subDirPath.substr(2)), pSvc, fileMap);
            }
        }
        else {
            wcerr << L"Failed to get the Name property." << endl;
        }
        VariantClear(&vtObjectName); // Очищаем VARIANT после использования

        count++; // Увеличиваем счетчик объектов
        // Освобождение объекта
        pclsObj->Release();
    }
    pEnumerator->Release();
}