From 21e072332c326f294d3b72a045fd8ecc2f8147a8 Mon Sep 17 00:00:00 2001 From: Jean Flach Date: Tue, 23 Jun 2015 13:15:55 +0200 Subject: [PATCH] Add path exclusion to check_disk.exe Also do some refactoring and make -m do something fixes #9201 --- plugins/check_disk.cpp | 90 ++++++++++++++++++++++++++++-------------- plugins/check_disk.h | 13 +++--- 2 files changed, 68 insertions(+), 35 deletions(-) diff --git a/plugins/check_disk.cpp b/plugins/check_disk.cpp index 3a5d1d70f..c7a565d01 100644 --- a/plugins/check_disk.cpp +++ b/plugins/check_disk.cpp @@ -20,6 +20,8 @@ #include #include #include +#include +#include #include "check_disk.h" @@ -45,7 +47,7 @@ INT wmain(INT argc, WCHAR **argv) printInfo.crit.legal = !printInfo.crit.legal; if (printInfo.drives.empty()) - ret = check_drives(vDrives); + ret = check_drives(vDrives, printInfo.exclude_drives); else ret = check_drives(vDrives, printInfo); @@ -62,7 +64,7 @@ INT wmain(INT argc, WCHAR **argv) return printOutput(printInfo, vDrives); } -INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& printInfo) +static INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& printInfo) { WCHAR namePath[MAX_PATH]; GetModuleFileName(NULL, namePath, MAX_PATH); @@ -77,23 +79,18 @@ INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& p ("warning,w", po::wvalue(), "Warning threshold") ("critical,c", po::wvalue(), "Critical threshold") ("path,p", po::wvalue>()->multitoken(), "Declare explicitly which drives to check (default checks all)") - ("unit,u", po::wvalue(), "Assign unit possible are: B, kB, MB, GB, TB") + ("exclude_device,x", po::wvalue>()->multitoken(), "Exclude these drives from check") + ("exclude-type,X", po::wvalue>()->multitoken(), "Exclude partition types (ignored)") + ("unit,u", po::wvalue(), "Assign unit possible are: B, kB, MB, GB, TB")\ + ("megabytes,m", "use megabytes, overridden by -unit") ; po::basic_command_line_parser parser(ac, av); try { - po::options_description allDesc; - - allDesc.add(desc); - allDesc.add_options() - ("exclude-type,X", po::wvalue>()->multitoken(), "exclude partition types (ignored)") - ("megabytes,m", "use megabytes") - ; - po::store( parser - .options(allDesc) + .options(desc) .style( po::command_line_style::unix_style | po::command_line_style::allow_long_disguise) @@ -124,7 +121,7 @@ INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& p L" 0\tOK,\n\tNo Thresholds were broken or the programs check part was not executed\n" L" 1\tWARNING,\n\tThe warning, but not the critical threshold was broken\n" L" 2\tCRITICAL,\n\tThe critical threshold was broken\n" - L" 3\tUNKNOWN, \n\tThe program experienced an INTernal or input error\n\n" + L" 3\tUNKNOWN, \n\tThe program experienced an internal or input error\n\n" L"Threshold syntax:\n\n" L"-w THRESHOLD\n" L"warn if threshold is broken, which means VALUE < THRESHOLD\n\n" @@ -168,6 +165,9 @@ INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& p if (vm.count("path")) printInfo.drives = vm["path"].as>(); + if (vm.count("exclude_device")) + printInfo.exclude_drives = vm["exclude_device"].as>(); + if (vm.count("unit")) { try { printInfo.unit = parseBUnit(vm["unit"].as()); @@ -175,8 +175,12 @@ INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& p std::wcout << "Unknown unit Type " << vm["unit"].as() << '\n'; return 3; } - } else - printInfo.unit = BunitB; + } else { + if (vm.count("megabytes")) + printInfo.unit = BunitMB; + else + printInfo.unit = BunitB; + } if (vm.count("debug")) debug = TRUE; @@ -184,7 +188,7 @@ INT parseArguments(INT ac, WCHAR **av, po::variables_map& vm, printInfoStruct& p return -1; } -INT printOutput(printInfoStruct& printInfo, std::vector& vDrives) +static INT printOutput(printInfoStruct& printInfo, std::vector& vDrives) { if (debug) std::wcout << "Constructing output string\n"; @@ -231,7 +235,7 @@ INT printOutput(printInfoStruct& printInfo, std::vector& vDrives) return state; } -INT check_drives(std::vector& vDrives) +static INT check_drives(std::vector& vDrives, std::vector& vExclude_Drives) { DWORD dwResult, dwSize = 0, dwVolumePathNamesLen = MAX_PATH + 1; WCHAR szLogicalDrives[1024], szVolumeName[MAX_PATH], *szVolumePathNames = NULL; @@ -248,7 +252,7 @@ INT check_drives(std::vector& vDrives) if (dwResult > MAX_PATH) goto die; if (debug) - std::wcout << "Splitting string INTo single drive names\n"; + std::wcout << "Splitting string intoo single drive names\n"; LPTSTR szSingleDrive = szLogicalDrives; while (*szSingleDrive) { @@ -259,9 +263,9 @@ INT check_drives(std::vector& vDrives) std::wcout << "Got: " << drname << '\n'; } - if (debug) - std::wcout << "Getting volume mountpoINTs (includes NTFS folders)\n" - << "Getting first volume\n"; + if (debug) + std::wcout << "Getting volume mountpoints (includes NTFS folders)\n" + << "Getting first volume\n"; hVolume = FindFirstVolume(szVolumeName, MAX_PATH); if (hVolume == INVALID_HANDLE_VALUE) @@ -276,7 +280,8 @@ INT check_drives(std::vector& vDrives) volumeNameEnd = wcslen(szVolumeName) - 1; szVolumePathNames = reinterpret_cast(new WCHAR[dwVolumePathNamesLen]); - while (!GetVolumePathNamesForVolumeName(szVolumeName, szVolumePathNames, dwVolumePathNamesLen, &dwVolumePathNamesLen)) { + while (!GetVolumePathNamesForVolumeName(szVolumeName, szVolumePathNames, dwVolumePathNamesLen, + &dwVolumePathNamesLen)) { if (GetLastError() != ERROR_MORE_DATA) break; delete[] reinterpret_cast(szVolumePathNames); @@ -303,8 +308,19 @@ INT check_drives(std::vector& vDrives) FindVolumeClose(hVolume); if (szVolumePathNames) delete[] reinterpret_cast(szVolumePathNames); + + if (!vExclude_Drives.empty()) { + if (debug) + std::wcout << "Removing excluded drives\n"; + + BOOST_FOREACH(const std::wstring wsDriveName, vExclude_Drives) + { + vDrives.erase(std::remove_if(vDrives.begin(), vDrives.end(), + boost::bind(checkName, _1, wsDriveName + L'\\')), vDrives.end()); + } + } return -1; - + die: if (hVolume) FindVolumeClose(hVolume); @@ -312,17 +328,27 @@ die: return 3; } -INT check_drives(std::vector& vDrives, printInfoStruct& printInfo) + + +static INT check_drives(std::vector& vDrives, printInfoStruct& printInfo) { - WCHAR *slash = L"\\"; + if (!printInfo.exclude_drives.empty()) { + if (debug) + std::wcout << "Removing excluded drives from user input drives\n"; + BOOST_FOREACH(std::wstring wsDrive, printInfo.exclude_drives) + { + printInfo.drives.erase(std::remove(printInfo.drives.begin(), printInfo.drives.end(), wsDrive), + printInfo.drives.end()); + } + } if (debug) std::wcout << "Parsing user input drive names\n"; for (std::vector::iterator it = printInfo.drives.begin(); it != printInfo.drives.end(); ++it) { - if (it->at(it->length() - 1) != *slash) - it->append(slash); + if (it->at(it->length() - 1) != *L"\\") + it->append(L"\\"); if (std::wstring::npos == it->find(L":\\")) { std::wcout << "A \":\" is required after the drive name of " << *it << '\n'; return 3; @@ -334,7 +360,7 @@ INT check_drives(std::vector& vDrives, printInfoStruct& printInfo) return -1; } -BOOL getFreeAndCap(drive& drive, const Bunit& unit) +static BOOL getFreeAndCap(drive& drive, const Bunit& unit) { if (debug) std::wcout << "Getting free disk space for drive " << drive.name << '\n'; @@ -347,10 +373,16 @@ BOOL getFreeAndCap(drive& drive, const Bunit& unit) drive.cap = round((tempTotal.QuadPart / pow(1024.0, unit))); if (debug) std::wcout << "\tAfter conversion: " << drive.cap << '\n' - << "\tfree: " << tempFree.QuadPart << '\n'; + << "\tfree: " << tempFree.QuadPart << '\n'; drive.free = round((tempFree.QuadPart / pow(1024.0, unit))); if (debug) std::wcout << "\tAfter conversion: " << drive.free << '\n' << '\n'; return TRUE; } + +static bool checkName(const drive& d, const std::wstring& s) +{ + return (s == d.name); +} + diff --git a/plugins/check_disk.h b/plugins/check_disk.h index 6f0a4dccf..c49199e0b 100644 --- a/plugins/check_disk.h +++ b/plugins/check_disk.h @@ -39,13 +39,14 @@ struct drive struct printInfoStruct { threshold warn, crit; - std::vector drives; + std::vector drives, exclude_drives; Bunit unit; }; -INT parseArguments(int, wchar_t **, boost::program_options::variables_map&, printInfoStruct&); -INT printOutput(printInfoStruct&, std::vector&); -INT check_drives(std::vector&); -INT check_drives(std::vector&, printInfoStruct&); -BOOL getFreeAndCap(drive&, const Bunit&); +static INT parseArguments(int, wchar_t **, boost::program_options::variables_map&, printInfoStruct&); +static INT printOutput(printInfoStruct&, std::vector&); +static INT check_drives(std::vector&, std::vector&); +static INT check_drives(std::vector&, printInfoStruct&); +static BOOL getFreeAndCap(drive&, const Bunit&); +static bool checkName(const drive& d, const std::wstring& s); #endif /*CHECK_DISK_H*/ \ No newline at end of file