From cbedf45796864cf0b1ead7c53b0af1ac1a4d11b7 Mon Sep 17 00:00:00 2001 From: seraphima Date: Tue, 9 Jul 2024 20:18:42 +0200 Subject: [PATCH] get command line args in the snapshot --- .../ProjectsSnapshotTool.vcxproj | 4 +- .../ProjectsSnapshotTool/SnapshotUtils.cpp | 142 +++++++++++++++++- .../Projects/ProjectsSnapshotTool/main.cpp | 9 ++ 3 files changed, 150 insertions(+), 5 deletions(-) diff --git a/src/modules/Projects/ProjectsSnapshotTool/ProjectsSnapshotTool.vcxproj b/src/modules/Projects/ProjectsSnapshotTool/ProjectsSnapshotTool.vcxproj index a434873b57..d7302e6c79 100644 --- a/src/modules/Projects/ProjectsSnapshotTool/ProjectsSnapshotTool.vcxproj +++ b/src/modules/Projects/ProjectsSnapshotTool/ProjectsSnapshotTool.vcxproj @@ -104,7 +104,7 @@ Windows true - shcore.lib;Shell32.lib;propsys.lib;DbgHelp.lib;%(AdditionalDependencies) + shcore.lib;Shell32.lib;propsys.lib;DbgHelp.lib;wbemuuid.lib @@ -122,7 +122,7 @@ true true true - shcore.lib;Shell32.lib;propsys.lib;DbgHelp.lib;%(AdditionalDependencies) + shcore.lib;Shell32.lib;propsys.lib;DbgHelp.lib;wbemuuid.lib diff --git a/src/modules/Projects/ProjectsSnapshotTool/SnapshotUtils.cpp b/src/modules/Projects/ProjectsSnapshotTool/SnapshotUtils.cpp index 3775920823..351031963a 100644 --- a/src/modules/Projects/ProjectsSnapshotTool/SnapshotUtils.cpp +++ b/src/modules/Projects/ProjectsSnapshotTool/SnapshotUtils.cpp @@ -1,23 +1,156 @@ #include "pch.h" #include "SnapshotUtils.h" -#include +#include +#include #include #include #include #include -#include namespace SnapshotUtils { + class WbemHelper + { + public: + WbemHelper() = default; + ~WbemHelper() + { + if (m_services) + { + m_services->Release(); + } + + if (m_locator) + { + m_locator->Release(); + } + } + + bool Initialize() + { + // Obtain the initial locator to WMI. + HRESULT hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, reinterpret_cast(&m_locator)); + if (FAILED(hres)) + { + Logger::error(L"Failed to create IWbemLocator object. Error: {}", get_last_error_or_default(hres)); + return false; + } + + // Connect to WMI through the IWbemLocator::ConnectServer method. + hres = m_locator->ConnectServer(_bstr_t(L"ROOT\\CIMV2"), NULL, NULL, 0, NULL, 0, 0, &m_services); + if (FAILED(hres)) + { + Logger::error(L"Could not connect to WMI. Error: {}", get_last_error_or_default(hres)); + return false; + } + + // Set security levels on the proxy. + hres = CoSetProxyBlanket(m_services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE); + if (FAILED(hres)) + { + Logger::error(L"Could not set proxy blanket. Error: {}", get_last_error_or_default(hres)); + return false; + } + + return true; + } + + std::wstring GetCommandLineArgs(DWORD processID) const + { + static std::wstring property = L"CommandLine"; + std::wstring query = L"SELECT " + property + L" FROM Win32_Process WHERE ProcessId = " + std::to_wstring(processID); + return Query(query, property); + } + + std::wstring GetExecutablePath(DWORD processID) const + { + static std::wstring property = L"ExecutablePath"; + std::wstring query = L"SELECT " + property + L" FROM Win32_Process WHERE ProcessId = " + std::to_wstring(processID); + return Query(query, property); + } + + private: + std::wstring Query(const std::wstring& query, const std::wstring& propertyName) const + { + if (!m_locator || !m_services) + { + return L""; + } + + IEnumWbemClassObject* pEnumerator = NULL; + + HRESULT hres = m_services->ExecQuery(bstr_t("WQL"), bstr_t(query.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumerator); + if (FAILED(hres)) + { + Logger::error(L"Query for process failed. Error: {}", get_last_error_or_default(hres)); + return L""; + } + + IWbemClassObject* pClassObject = NULL; + ULONG uReturn = 0; + std::wstring result = L""; + while (pEnumerator) + { + HRESULT hr = pEnumerator->Next(WBEM_INFINITE, 1, &pClassObject, &uReturn); + if (uReturn == 0) + { + break; + } + + VARIANT vtProp; + hr = pClassObject->Get(propertyName.c_str(), 0, &vtProp, 0, 0); + if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) + { + result = vtProp.bstrVal; + } + VariantClear(&vtProp); + + pClassObject->Release(); + } + + pEnumerator->Release(); + + return result; + } + + IWbemLocator* m_locator = NULL; + IWbemServices* m_services = NULL; + }; + + std::wstring GetCommandLineArgs(DWORD processID, const WbemHelper& wbemHelper) + { + std::wstring executablePath = wbemHelper.GetExecutablePath(processID); + std::wstring commandLineArgs = wbemHelper.GetCommandLineArgs(processID); + + if (!commandLineArgs.empty()) + { + auto pos = commandLineArgs.find(executablePath); + if (pos != std::wstring::npos) + { + commandLineArgs = commandLineArgs.substr(pos + executablePath.size()); + auto spacePos = commandLineArgs.find_first_of(' '); + if (spacePos != std::wstring::npos) + { + commandLineArgs = commandLineArgs.substr(spacePos + 1); + } + } + } + + return commandLineArgs; + } + std::vector GetApps(const std::function getMonitorNumberFromWindowHandle) { std::vector apps{}; auto installedApps = Utils::Apps::GetAppsList(); auto windows = WindowEnumerator::Enumerate(WindowFilter::Filter); + + WbemHelper wbemHelper; + wbemHelper.Initialize(); for (const auto window : windows) { @@ -42,6 +175,9 @@ namespace SnapshotUtils continue; } + DWORD pid{}; + GetWindowThreadProcessId(window, &pid); + auto data = Utils::Apps::GetApp(processPath, installedApps); if (!data.has_value() || data->name.empty()) { @@ -53,7 +189,7 @@ namespace SnapshotUtils .title = title, .path = processPath, .packageFullName = data.value().packageFullName, - .commandLineArgs = L"", + .commandLineArgs = GetCommandLineArgs(pid, wbemHelper), .isMinimized = WindowUtils::IsMinimized(window), .isMaximized = WindowUtils::IsMaximized(window), .position = Project::Application::Position{ diff --git a/src/modules/Projects/ProjectsSnapshotTool/main.cpp b/src/modules/Projects/ProjectsSnapshotTool/main.cpp index b2bb309c57..e0b497d884 100644 --- a/src/modules/Projects/ProjectsSnapshotTool/main.cpp +++ b/src/modules/Projects/ProjectsSnapshotTool/main.cpp @@ -36,6 +36,15 @@ int APIENTRY WinMain(HINSTANCE hInst, HINSTANCE hInstPrev, LPSTR cmdLine, int cm return -1; } + // Set general COM security levels. + comInitHres = CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, NULL); + if (FAILED(comInitHres)) + { + Logger::error(L"Failed to initialize security. Error code: {}", get_last_error_or_default(comInitHres)); + CoUninitialize(); + return -1; + } + std::wstring fileName = JsonUtils::ProjectsFile(); std::string cmdLineStr(cmdLine); if (!cmdLineStr.empty())