void ProcessDiskInfo(wstring apiKey, wstring envKey, SOCKET socketHandle, sockaddr_in remoteServAddr, int interval) { StringBuilder sb; int sleepTime = interval * 1000; wchar_t driveLetter[_MAX_PATH + 1]; HRESULT hr; ULARGE_INTEGER freeBytesAvailable; ULARGE_INTEGER totalNumberOfBytes; ULARGE_INTEGER totalNumberOfFreeBytes; while (true) { if (g_servicePaused == FALSE) { sb.clear(); sb.push_back(L"{"); sb.push_back(L"\"" + StorageInfo::Members::disk + L"\":"); sb.push_back(L"["); DWORD dwNeedLength = ::GetLogicalDriveStrings(0, NULL); wchar_t *pWchBuf = nullptr; wchar_t *pHead = nullptr; do { if (dwNeedLength != 0) { pWchBuf = new wchar_t[dwNeedLength + 1]; if (pWchBuf == nullptr) { break; } pHead = pWchBuf; DWORD dwResult = ::GetLogicalDriveStrings(dwNeedLength, pWchBuf); while (dwResult > 0) { hr = StringCchPrintf(driveLetter, _MAX_PATH, L"%s", pWchBuf); if (FAILED(hr) == TRUE) { break; } int length = wcslen(driveLetter) + 1; if (length < 2) { break; } dwResult -= length; pWchBuf += length; if (::GetDiskFreeSpaceEx(driveLetter, &freeBytesAvailable, &totalNumberOfBytes, &totalNumberOfFreeBytes) == FALSE) { break; } driveLetter[1] = '\0'; sb.push_back(L"{ \"" + DiskInfo::Members::name + L"\": \""); sb.push_back(driveLetter); sb.push_back(L"\", "); sb.push_back(L"\"" + DiskInfo::Members::size + L"\": "); sb.push_back((__int64)totalNumberOfBytes.QuadPart); sb.push_back(L", "); sb.push_back(L"\"" + DiskInfo::Members::current + L"\": "); sb.push_back((__int64)(totalNumberOfBytes.QuadPart - totalNumberOfFreeBytes.QuadPart)); sb.push_back(L"}"); if (dwResult > 0) { sb.push_back(L","); } } } } while (false); if (pHead != nullptr) { delete[] pHead; } sb.push_back(L"],"); sb.push_back(L"\"" + PacketBase::Members::groupKey + L"\":"); sb.push_back(L"\""); sb.push_back(apiKey); sb.push_back(L"\","); sb.push_back(L"\"" + PacketBase::Members::machineId + L"\":"); sb.push_back(L"\""); sb.push_back(envKey); sb.push_back(L"\""); sb.push_back(L"}"); } #if defined(_DEBUG) // ::OutputDebugString(sb.ToString().c_str()); #endif SendToServer(socketHandle, remoteServAddr, sb); if (g_isConsoleApp == TRUE) { printf(":"); } DWORD dwEventResult = ::WaitForSingleObject(g_killServiceEvent, sleepTime); if (dwEventResult == WAIT_TIMEOUT) { continue; } #if defined(_DEBUG) ::OutputDebugString(L"ProcessDiskInfo-thread exited."); #endif break; } }
void DoRegistration(wstring apiKey, wstring envKey, string remoteServAddr, int port, vector<int> intervalTimes) { OutputConsole(L"Installing Service...\n"); DoUnregistration(); SC_HANDLE scm = OpenSCManager(0, 0, SC_MANAGER_CREATE_SERVICE); if (!scm) { OutputError(L"OpenSCManager fails! (%d)\n", GetLastError()); return; } SC_HANDLE myService = NULL; HKEY hkey = NULL; wchar_t *pBuf = nullptr; do { wstring filePath = g_moduleFilePath; wchar_t safeFilePath[MAX_PATH]; StringCchPrintf(safeFilePath, MAX_PATH, L"\"%s\"", filePath.c_str()); OutputConsole(L"Opened Service Control Manager...\n"); myService = CreateService( scm, SERVICE_NAME, // the internal service name used by the SCM L"StrawAgent Service", // the external label seen in the Service Control applet SERVICE_ALL_ACCESS, // We want full access to control the service SERVICE_WIN32_OWN_PROCESS, // The service is a single app and not a driver SERVICE_AUTO_START, // The service will be started by us manually SERVICE_ERROR_NORMAL, // If error during service start, don't misbehave. safeFilePath, 0, 0, 0, 0, 0); if (!myService) { OutputError(L"CreateService fails! (%d)\n", GetLastError()); break; } wchar_t szBuffer[MAX_PATH] = { 0 }; StringBuilder sb; sb.push_back(L"asNTservice=1"); sb.push_back(L"apiKey=" + apiKey); sb.push_back(L"envKey=" + envKey); sb.push_back("server=" + remoteServAddr); StringCchPrintf(szBuffer, MAX_PATH, L"port=%d", port); sb.push_back(szBuffer); if (g_debugMode == TRUE) { sb.push_back(L"debug=1"); } if (RegOpenKey(HKEY_LOCAL_MACHINE, REG_SERVICE, &hkey) != ERROR_SUCCESS) { OutputError(L"RegOpenKey fails! (%d)\n", GetLastError()); break; } int written = 0; pBuf = sb.ToStringMultiLine(&written); if (RegSetValueEx(hkey, L"Environment", 0, REG_MULTI_SZ, (const BYTE *)pBuf, written) != ERROR_SUCCESS) { OutputError(L"RegSetValueEx fails! (%d)\n", GetLastError()); break; } OutputConsole(L"Service successfully installed.\n"); } while (false); if (pBuf != nullptr) { delete[] pBuf; } if (hkey != NULL) { RegCloseKey(hkey); } if (myService != NULL) { CloseServiceHandle(myService); } if (scm != NULL) { CloseServiceHandle(scm); } }
void ProcessCpuMemInfo(wstring apiKey, wstring envKey, SOCKET socketHandle, sockaddr_in remoteServAddr, int interval) { StringBuilder sb; int sleepTime = interval * 1000; while (true) { if (g_servicePaused == FALSE) { sb.clear(); sb.push_back(L"{"); { sb.push_back(L"\"" + SystemInfo::Members::cpuUsage + L"\":"); sb.push_back(L"{"); { float totalUsage = 0.0f; sb.push_back(L"\"" + CpuInfo::Members::unit + L"\":["); if (RetrieveCpuInfo(sb, &totalUsage) == false) { Sleep(1000); continue; } sb.push_back(L"]"); wchar_t buf[40]; StringCchPrintf(buf, 40, L", \"%s\": %.2f", CpuInfo::Members::total.c_str(), totalUsage / 100); sb.push_back(buf); } sb.push_back(L"},"); __int64 maxMemory; __int64 currentUsage; GetMemoryInfo(&maxMemory, ¤tUsage); sb.push_back(L"\"" + SystemInfo::Members::memoryUsage + L"\":"); { sb.push_back(L"{\"" + MemoryInfo::Members::max + L"\":"); sb.push_back(maxMemory); sb.push_back(L", \"" + MemoryInfo::Members::current + L"\":"); sb.push_back(currentUsage); sb.push_back(L"},"); } sb.push_back(L"\"" + PacketBase::Members::groupKey + L"\":"); sb.push_back(L"\""); sb.push_back(apiKey); sb.push_back(L"\","); sb.push_back(L"\"" + PacketBase::Members::machineId + L"\":"); sb.push_back(L"\""); sb.push_back(envKey); sb.push_back(L"\""); } sb.push_back(L"}"); SendToServer(socketHandle, remoteServAddr, sb); if (g_isConsoleApp == TRUE) { printf("."); } } if (::WaitForSingleObject(g_killServiceEvent, sleepTime) == WAIT_TIMEOUT) { continue; } #if defined(_DEBUG) ::OutputDebugString(L"ProcessCpuMemInfo-thread exited."); #endif break; } }