granicus.if.org Git - icinga2/blob - plugins/check_disk.cpp

   1 /* Icinga 2 | (c) 2012 Icinga GmbH | GPLv2+ */
   2
   3 #include "plugins/thresholds.hpp"
   4 #include <boost/program_options.hpp>
   5 #include <vector>
   6 #include <windows.h>
   7 #include <set>
   8 #include <iostream>
   9 #include <functional>
  10 #include <shlwapi.h>
  11 #include <math.h>
  12
  13 using namespace std::placeholders;
  14
  15 #define VERSION 1.1
  16
  17 namespace po = boost::program_options;
  18
  19 struct drive
  20 {
  21         std::wstring name;
  22         double cap;
  23         double free;
  24         double used;
  25
  26         drive(std::wstring p)
  27                 : name(p)
  28         { }
  29 };
  30
  31 struct printInfoStruct
  32 {
  33         threshold warn;
  34         threshold crit;
  35         std::vector<std::wstring> drives;
  36         std::vector<std::wstring> exclude_drives;
  37         Bunit unit;
  38         bool showUsed{false};
  39 };
  40
  41 static bool l_Debug;
  42
  43 static int check_drives(std::vector<drive>& vDrives, std::vector<std::wstring>& vExclude_Drives)
  44 {
  45         DWORD dwResult, dwSize = 0, dwVolumePathNamesLen = MAX_PATH + 1;
  46         WCHAR szLogicalDrives[1024], szVolumeName[MAX_PATH], *szVolumePathNames = NULL;
  47         HANDLE hVolume = NULL;
  48         std::wstring wsLogicalDrives;
  49         size_t volumeNameEnd = 0;
  50
  51         std::set<std::wstring> sDrives;
  52
  53         if (l_Debug)
  54                 std::wcout << "Getting logic drive string (includes network drives)\n";
  55
  56         dwResult = GetLogicalDriveStrings(MAX_PATH, szLogicalDrives);
  57         if (dwResult > MAX_PATH)
  58                 goto die;
  59         if (l_Debug)
  60                 std::wcout << "Splitting string into single drive names\n";
  61
  62         LPTSTR szSingleDrive = szLogicalDrives;
  63         while (*szSingleDrive) {
  64                 std::wstring drname = szSingleDrive;
  65                 sDrives.insert(drname);
  66                 szSingleDrive += wcslen(szSingleDrive) + 1;
  67                 if (l_Debug)
  68                         std::wcout << "Got: " << drname << '\n';
  69         }
  70
  71         if (l_Debug)
  72                 std::wcout << "Getting volume mountpoints (includes NTFS folders)\n"
  73                 << "Getting first volume\n";
  74
  75         hVolume = FindFirstVolume(szVolumeName, MAX_PATH);
  76         if (hVolume == INVALID_HANDLE_VALUE)
  77                 goto die;
  78
  79         if (l_Debug)
  80                 std::wcout << "Traversing through list of drives\n";
  81
  82         while (GetLastError() != ERROR_NO_MORE_FILES) {
  83                 if (l_Debug)
  84                         std::wcout << "Path name for " << szVolumeName << "= \"";
  85                 volumeNameEnd = wcslen(szVolumeName) - 1;
  86                 szVolumePathNames = new WCHAR[dwVolumePathNamesLen];
  87
  88                 while (!GetVolumePathNamesForVolumeName(szVolumeName, szVolumePathNames, dwVolumePathNamesLen,
  89                         &dwVolumePathNamesLen)) {
  90                         if (GetLastError() != ERROR_MORE_DATA)
  91                                 break;
  92                         delete[] szVolumePathNames;
  93                         szVolumePathNames = new WCHAR[dwVolumePathNamesLen];
  94
  95                 }
  96                 if (l_Debug)
  97                         std::wcout << szVolumePathNames << "\"\n";
  98
  99                 sDrives.insert(std::wstring(szVolumePathNames));
 100                 FindNextVolume(hVolume, szVolumeName, MAX_PATH);
 101         }
 102         if (l_Debug)
 103                 std::wcout << "Creating vector from found volumes, ignoring CD drives etc.:\n";
 104         for (const auto& driveName : sDrives) {
 105                 unsigned int type = GetDriveType(driveName.c_str());
 106                 if (type == DRIVE_FIXED || type == DRIVE_REMOTE) {
 107                         if (l_Debug)
 108                                 std::wcout << "\t" << driveName << '\n';
 109                         vDrives.push_back(drive(driveName));
 110                 }
 111         }
 112
 113         FindVolumeClose(hVolume);
 114         if (szVolumePathNames)
 115                 delete[] szVolumePathNames;
 116
 117         if (l_Debug)
 118                 std::wcout << "Removing excluded drives\n";
 119
 120         for (const auto& driveName : vExclude_Drives) {
 121                 vDrives.erase(std::remove_if(vDrives.begin(), vDrives.end(),
 122                         [&driveName](const drive& d) { return d.name == driveName + L'\\'; }),
 123                         vDrives.end());
 124         }
 125
 126         return -1;
 127
 128 die:
 129         if (hVolume)
 130                 FindVolumeClose(hVolume);
 131         printErrorInfo();
 132         return 3;
 133 }
 134
 135 static int check_drives(std::vector<drive>& vDrives, printInfoStruct& printInfo)
 136 {
 137         if (l_Debug)
 138                 std::wcout << "Removing excluded drives from user input drives\n";
 139
 140         for (const auto& driveName : printInfo.exclude_drives) {
 141                 printInfo.drives.erase(std::remove(printInfo.drives.begin(), printInfo.drives.end(), driveName),
 142                         printInfo.drives.end());
 143         }
 144
 145         if (l_Debug)
 146                 std::wcout << "Parsing user input drive names\n";
 147
 148         for (auto& driveName : printInfo.drives) {
 149                 if (driveName.at(driveName.length() - 1) != *L"\\")
 150                         driveName.append(L"\\");
 151
 152                 if (std::wstring::npos == driveName.find(L":\\")) {
 153                         std::wcout << "A \":\" is required after the drive name of " << driveName << '\n';
 154                         return 3;
 155                 }
 156
 157                 if (l_Debug)
 158                         std::wcout << "Added " << driveName << '\n';
 159
 160                 vDrives.emplace_back(driveName);
 161         }
 162
 163         return -1;
 164 }
 165
 166 static bool getDriveSpaceValues(drive& drive, const Bunit& unit)
 167 {
 168         if (l_Debug)
 169                 std::wcout << "Getting free and used disk space for drive " << drive.name << '\n';
 170
 171         ULARGE_INTEGER tempFree, tempTotal;
 172         if (!GetDiskFreeSpaceEx(drive.name.c_str(), NULL, &tempTotal, &tempFree))
 173                 return false;
 174
 175         ULARGE_INTEGER tempUsed;
 176         tempUsed.QuadPart = tempTotal.QuadPart - tempFree.QuadPart;
 177
 178         if (l_Debug)
 179                 std::wcout << "\tcap: " << tempFree.QuadPart << '\n';
 180
 181         drive.cap = round((tempTotal.QuadPart / pow(1024.0, unit)));
 182
 183         if (l_Debug)
 184                 std::wcout << "\tAfter conversion: " << drive.cap << '\n'
 185                 << "\tfree: " << tempFree.QuadPart << '\n';
 186
 187         drive.free = round((tempFree.QuadPart / pow(1024.0, unit)));
 188
 189         if (l_Debug)
 190                 std::wcout << "\tAfter conversion: " << drive.free << '\n'
 191                 << "\tused: " << tempUsed.QuadPart << '\n';
 192
 193         drive.used = round((tempUsed.QuadPart / pow(1024.0, unit)));
 194
 195         if (l_Debug)
 196                 std::wcout << "\tAfter conversion: " << drive.used << '\n' << '\n';
 197
 198         return true;
 199 }
 200
 201 static int parseArguments(int ac, WCHAR **av, po::variables_map& vm, printInfoStruct& printInfo)
 202 {
 203         WCHAR namePath[MAX_PATH];
 204         GetModuleFileName(NULL, namePath, MAX_PATH);
 205         WCHAR *progName = PathFindFileName(namePath);
 206
 207         po::options_description desc("Options");
 208
 209         desc.add_options()
 210                 ("help,h", "Print usage message and exit")
 211                 ("version,V", "Print version and exit")
 212                 ("debug,d", "Verbose/Debug output")
 213                 ("warning,w", po::wvalue<std::wstring>(), "Warning threshold")
 214                 ("critical,c", po::wvalue<std::wstring>(), "Critical threshold")
 215                 ("path,p", po::wvalue<std::vector<std::wstring>>()->multitoken(), "Declare explicitly which drives to check (default checks all)")
 216                 ("exclude_device,x", po::wvalue<std::vector<std::wstring>>()->multitoken(), "Exclude these drives from check")
 217                 ("exclude-type,X", po::wvalue<std::vector<std::wstring>>()->multitoken(), "Exclude partition types (ignored)")
 218                 ("iwarning,W", po::wvalue<std::wstring>(), "Warning threshold for inodes (ignored)")
 219                 ("icritical,K", po::wvalue<std::wstring>(), "Critical threshold for inodes (ignored)")
 220                 ("unit,u", po::wvalue<std::wstring>(), "Assign unit possible are: B, kB, MB, GB, TB")
 221                 ("show-used,U", "Show used space instead of the free space")
 222                 ("megabytes,m", "use megabytes, overridden by -unit")
 223                 ;
 224
 225         po::wcommand_line_parser parser(ac, av);
 226
 227         try {
 228                 po::store(
 229                         parser
 230                         .options(desc)
 231                         .style(
 232                                 po::command_line_style::unix_style |
 233                                 po::command_line_style::allow_long_disguise)
 234                         .run(),
 235                         vm);
 236                 vm.notify();
 237         } catch (const std::exception& e) {
 238                 std::cout << e.what() << '\n' << desc << '\n';
 239                 return 3;
 240         }
 241
 242         if (vm.count("help")) {
 243                 std::wcout << progName << " Help\n\tVersion: " << VERSION << '\n';
 244                 wprintf(
 245                         L"%s is a simple program to check a machines disk space usage.\n"
 246                         L"You can use the following options to define its behaviour:\n\n", progName);
 247                 std::cout << desc;
 248                 wprintf(
 249                         L"\nIt will then output a string looking something like this:\n\n"
 250                         L"\tDISK WARNING 29GB | disk=29GB;50%%;5;0;120\n\n"
 251                         L"\"DISK\" being the type of the check, \"WARNING\" the returned status\n"
 252                         L"and \"23.8304%%\" is the returned value.\n"
 253                         L"The performance data is found behind the \"|\", in order:\n"
 254                         L"returned value, warning threshold, critical threshold, minimal value and,\n"
 255                         L"if applicable, the maximal value.\n"
 256                         L"This program will also print out additional performance data disk by disk\n\n"
 257                         L"%s' exit codes denote the following:\n\n"
 258                         L" 0\tOK,\n\tNo Thresholds were broken or the programs check part was not executed\n"
 259                         L" 1\tWARNING,\n\tThe warning, but not the critical threshold was broken\n"
 260                         L" 2\tCRITICAL,\n\tThe critical threshold was broken\n"
 261                         L" 3\tUNKNOWN, \n\tThe program experienced an internal or input error\n\n"
 262                         L"Threshold syntax:\n\n"
 263                         L"-w THRESHOLD\n"
 264                         L"warn if threshold is broken, which means VALUE < THRESHOLD\n\n"
 265                         L"-w !THRESHOLD\n"
 266                         L"inverts threshold check, VALUE > THRESHOLD (analogous to above)\n\n"
 267                         L"-w [THR1-THR2]\n"
 268                         L"warn is VALUE is inside the range spanned by THR1 and THR2\n\n"
 269                         L"-w ![THR1-THR2]\n"
 270                         L"warn if VALUE is outside the range spanned by THR1 and THR2\n\n"
 271                         L"-w THRESHOLD%%\n"
 272                         L"if the plugin accepts percentage based thresholds those will be used.\n"
 273                         L"Does nothing if the plugin does not accept percentages, or only uses\n"
 274                         L"percentage thresholds. Ranges can be used with \"%%\", but both range values need\n"
 275                         L"to end with a percentage sign.\n\n"
 276                         L"All of these options work with the critical threshold \"-c\" too."
 277                         , progName);
 278                 std::cout << '\n';
 279                 return 0;
 280         }
 281
 282         if (vm.count("version"))
 283                 std::cout << "Version: " << VERSION << '\n';
 284
 285         if (vm.count("warning")) {
 286                 try {
 287                         printInfo.warn = threshold(vm["warning"].as<std::wstring>());
 288                 } catch (const std::invalid_argument& e) {
 289                         std::cout << e.what() << '\n';
 290                         return 3;
 291                 }
 292         }
 293         if (vm.count("critical")) {
 294                 try {
 295                         printInfo.crit = threshold(vm["critical"].as<std::wstring>());
 296                 } catch (const std::invalid_argument& e) {
 297                         std::cout << e.what() << '\n';
 298                         return 3;
 299                 }
 300         }
 301
 302         if (vm.count("path"))
 303                 printInfo.drives = vm["path"].as<std::vector<std::wstring>>();
 304
 305         if (vm.count("exclude_device"))
 306                 printInfo.exclude_drives = vm["exclude_device"].as<std::vector<std::wstring>>();
 307
 308         if (vm.count("unit")) {
 309                 try {
 310                         printInfo.unit = parseBUnit(vm["unit"].as<std::wstring>());
 311                 } catch (const std::invalid_argument&) {
 312                         std::wcout << "Unknown unit Type " << vm["unit"].as<std::wstring>() << '\n';
 313                         return 3;
 314                 }
 315         } else {
 316                 if (vm.count("megabytes"))
 317                         printInfo.unit = BunitMB;
 318                 else
 319                         printInfo.unit = BunitB;
 320         }
 321
 322         printInfo.showUsed = vm.count("show-used") > 0;
 323
 324         l_Debug = vm.count("debug") > 0;
 325
 326         return -1;
 327 }
 328
 329 static int printOutput(printInfoStruct& printInfo, std::vector<drive>& vDrives)
 330 {
 331         if (l_Debug)
 332                 std::wcout << "Constructing output string\n";
 333
 334         std::vector<std::wstring> wsDrives, wsPerf;
 335         std::wstring unit = BunitStr(printInfo.unit);
 336
 337         state state = OK;
 338
 339         std::wstring output = L"DISK OK - free space:";
 340
 341         if (printInfo.showUsed)
 342                 output = L"DISK OK - used space:";
 343
 344         double tCap = 0, tFree = 0, tUsed = 0;
 345
 346         for (std::vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); it++) {
 347                 tCap += it->cap;
 348                 tFree += it->free;
 349                 tUsed += it->used;
 350
 351                 if (printInfo.showUsed) {
 352                         wsDrives.push_back(it->name + L" " + removeZero(it->used) + L" " + unit + L" (" +
 353                                 removeZero(std::round(it->used / it->cap * 100.0)) + L"%); ");
 354
 355                         wsPerf.push_back(L" " + it->name + L"=" + removeZero(it->used) + unit + L";" +
 356                                 printInfo.warn.pString(it->cap) + L";" + printInfo.crit.pString(it->cap) + L";0;"
 357                                 + removeZero(it->cap));
 358
 359                         if (printInfo.crit.set && !printInfo.crit.rend(it->used, it->cap))
 360                                 state = CRITICAL;
 361
 362                         if (state == OK && printInfo.warn.set && !printInfo.warn.rend(it->used, it->cap))
 363                                 state = WARNING;
 364                 } else {
 365                         wsDrives.push_back(it->name + L" " + removeZero(it->free) + L" " + unit + L" (" +
 366                                 removeZero(std::round(it->free / it->cap * 100.0)) + L"%); ");
 367
 368                         wsPerf.push_back(L" '" + it->name + L"'=" + removeZero(it->free) + unit + L";" +
 369                                 printInfo.warn.pString(it->cap) + L";" + printInfo.crit.pString(it->cap) + L";0;"
 370                                 + removeZero(it->cap));
 371
 372                         if (printInfo.crit.rend(it->free, it->cap))
 373                                 state = CRITICAL;
 374
 375                         if (state == OK && printInfo.warn.rend(it->free, it->cap))
 376                                 state = WARNING;
 377                 }
 378         }
 379
 380         if (state == WARNING) {
 381                 output = L"DISK WARNING - free space:";
 382
 383                 if (printInfo.showUsed)
 384                         output = L"DISK WARNING - used space:";
 385         }
 386
 387         if (state == CRITICAL) {
 388                 output = L"DISK CRITICAL - free space:";
 389
 390                 if (printInfo.showUsed)
 391                         output = L"DISK CRITICAL - used space:";
 392         }
 393
 394         std::wcout << output;
 395
 396         if (vDrives.size() > 1 && printInfo.showUsed) {
 397                 std::wcout << "Total " << (printInfo.showUsed ? tUsed : tFree) << unit
 398                         << " (" << removeZero(std::round(tUsed / tCap * 100.0)) << "%); ";
 399         }
 400
 401         for (const auto& driveName : wsDrives)
 402                 std::wcout << driveName;
 403
 404         std::wcout << "|";
 405
 406         for (const auto& perf : wsPerf)
 407                 std::wcout << perf;
 408
 409         std::wcout << '\n';
 410
 411         return state;
 412 }
 413
 414 int wmain(int argc, WCHAR **argv)
 415 {
 416         std::vector<drive> vDrives;
 417         printInfoStruct printInfo;
 418         po::variables_map vm;
 419
 420         int ret;
 421
 422         ret = parseArguments(argc, argv, vm, printInfo);
 423         if (ret != -1)
 424                 return ret;
 425
 426         printInfo.warn.legal = !printInfo.warn.legal;
 427         printInfo.crit.legal = !printInfo.crit.legal;
 428
 429         if (printInfo.drives.empty())
 430                 ret = check_drives(vDrives, printInfo.exclude_drives);
 431         else
 432                 ret = check_drives(vDrives, printInfo);
 433
 434         if (ret != -1)
 435                 return ret;
 436
 437         for (std::vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); ++it) {
 438                 if (!getDriveSpaceValues(*it, printInfo.unit)) {
 439                         std::wcout << "Failed to access drive at " << it->name << '\n';
 440                         return 3;
 441                 }
 442         }
 443
 444         return printOutput(printInfo, vDrives);
 445 }