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

   1 /******************************************************************************
   2  * Icinga 2                                                                   *
   3  * Copyright (C) 2012-2015 Icinga Development Team (http://www.icinga.org)    *
   4  *                                                                            *
   5  * This program is free software; you can redistribute it and/or              *
   6  * modify it under the terms of the GNU General Public License                *
   7  * as published by the Free Software Foundation; either version 2             *
   8  * of the License, or (at your option) any later version.                     *
   9  *                                                                            *
  10  * This program is distributed in the hope that it will be useful,            *
  11  * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
  12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
  13  * GNU General Public License for more details.                               *
  14  *                                                                            *
  15  * You should have received a copy of the GNU General Public License          *
  16  * along with this program; if not, write to the Free Software Foundation     *
  17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.             *
  18  ******************************************************************************/
  19 #include <Windows.h>
  20 #include <set>
  21 #include <Shlwapi.h>
  22 #include <vector>
  23 #include <iostream>
  24 #include <math.h>
  25
  26 #include "thresholds.h"
  27
  28 #include "boost/program_options.hpp"
  29
  30 #define VERSION 1.1
  31
  32 namespace po = boost::program_options;
  33
  34 using std::cout; using std::endl; using std::set;
  35 using std::vector; using std::wstring; using std::wcout;
  36
  37 static BOOL debug = FALSE;
  38
  39 struct drive
  40 {
  41         wstring name;
  42         double cap, free;
  43         drive(wstring p)
  44                 : name(p)
  45         {}
  46 };
  47
  48 struct printInfoStruct
  49 {
  50         threshold warn, crit;
  51         vector<wstring> drives;
  52         Bunit unit;
  53 };
  54
  55 static int parseArguments(int, wchar_t **, po::variables_map&, printInfoStruct&);
  56 static int printOutput(printInfoStruct&, vector<drive>&);
  57 static int check_drives(vector<drive>&);
  58 static int check_drives(vector<drive>&, printInfoStruct&);
  59 static bool getFreeAndCap(drive&, const Bunit&);
  60
  61 int wmain(int argc, wchar_t **argv)
  62 {
  63         vector<drive> vDrives;
  64         printInfoStruct printInfo{ };
  65         po::variables_map vm;
  66
  67         int ret;
  68
  69         ret = parseArguments(argc, argv, vm, printInfo);
  70         if (ret != -1)
  71                 return ret;
  72
  73         printInfo.warn.legal = !printInfo.warn.legal;
  74         printInfo.crit.legal = !printInfo.crit.legal;
  75
  76         if (printInfo.drives.empty())
  77                 ret = check_drives(vDrives);
  78         else
  79                 ret = check_drives(vDrives, printInfo);
  80
  81         if (ret != -1)
  82                 return ret;
  83
  84         for (vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); ++it) {
  85                 if (!getFreeAndCap(*it, printInfo.unit)) {
  86                         wcout << L"Failed to access drive at " << it->name << endl;
  87                         return 3;
  88                 }
  89         }
  90
  91         return printOutput(printInfo, vDrives);
  92 }
  93
  94 int parseArguments(int ac, wchar_t **av, po::variables_map& vm, printInfoStruct& printInfo)
  95 {
  96         wchar_t namePath[MAX_PATH];
  97         GetModuleFileName(NULL, namePath, MAX_PATH);
  98         wchar_t *progName = PathFindFileName(namePath);
  99
 100         po::options_description desc("Options");
 101
 102         desc.add_options()
 103                 ("help,h", "print usage message and exit")
 104                 ("version,V", "print version and exit")
 105                 ("debug,d", "Verbose/Debug output")
 106                 ("warning,w", po::wvalue<wstring>(), "warning threshold")
 107                 ("critical,c", po::wvalue<wstring>(), "critical threshold")
 108                 ("path,p", po::wvalue<vector<std::wstring>>()->multitoken(), "declare explicitly which drives to check (default checks all)")
 109                 ("unit,u", po::wvalue<wstring>(), "assign unit possible are: B, kB, MB, GB, TB")
 110                 ;
 111
 112         po::basic_command_line_parser<wchar_t> parser(ac, av);
 113
 114         try {
 115                 po::options_description allDesc;
 116
 117                 allDesc.add(desc);
 118                 allDesc.add_options()
 119                         ("exclude-type,X", po::wvalue<vector<std::wstring>>()->multitoken(), "exclude partition types (ignored)")
 120                         ("megabytes,m", "use megabytes")
 121                         ;
 122
 123                 po::store(
 124                         parser
 125                         .options(allDesc)
 126                         .style(
 127                         po::command_line_style::unix_style |
 128                         po::command_line_style::allow_long_disguise)
 129                         .run(),
 130                         vm);
 131                 vm.notify();
 132         } catch (std::exception& e) {
 133                 cout << e.what() << endl << desc << endl;
 134                 return 3;
 135         }
 136
 137         if (vm.count("help")) {
 138                 wcout << progName << " Help\n\tVersion: " << VERSION << endl;
 139                 wprintf(
 140                         L"%s is a simple program to check a machines free disk space.\n"
 141                         L"You can use the following options to define its behaviour:\n\n", progName);
 142                 cout << desc;
 143                 wprintf(
 144                         L"\nIt will then output a string looking something like this:\n\n"
 145                         L"\tDISK WARNING 29GB | disk=29GB;50%%;5;0;120\n\n"
 146                         L"\"DISK\" being the type of the check, \"WARNING\" the returned status\n"
 147                         L"and \"23.8304%%\" is the returned value.\n"
 148                         L"The performance data is found behind the \"|\", in order:\n"
 149                         L"returned value, warning threshold, critical threshold, minimal value and,\n"
 150                         L"if applicable, the maximal value.\n"
 151                         L"This program will also print out additional performance data disk by disk\n\n"
 152                         L"%s' exit codes denote the following:\n\n"
 153                         L" 0\tOK,\n\tNo Thresholds were broken or the programs check part was not executed\n"
 154                         L" 1\tWARNING,\n\tThe warning, but not the critical threshold was broken\n"
 155                         L" 2\tCRITICAL,\n\tThe critical threshold was broken\n"
 156                         L" 3\tUNKNOWN, \n\tThe program experienced an internal or input error\n\n"
 157                         L"Threshold syntax:\n\n"
 158                         L"-w THRESHOLD\n"
 159                         L"warn if threshold is broken, which means VALUE < THRESHOLD\n\n"
 160                         L"-w !THRESHOLD\n"
 161                         L"inverts threshold check, VALUE > THRESHOLD (analogous to above)\n\n"
 162                         L"-w [THR1-THR2]\n"
 163                         L"warn is VALUE is inside the range spanned by THR1 and THR2\n\n"
 164                         L"-w ![THR1-THR2]\n"
 165                         L"warn if VALUE is outside the range spanned by THR1 and THR2\n\n"
 166                         L"-w THRESHOLD%%\n"
 167                         L"if the plugin accepts percentage based thresholds those will be used.\n"
 168                         L"Does nothing if the plugin does not accept percentages, or only uses\n"
 169                         L"percentage thresholds. Ranges can be used with \"%%\", but both range values need\n"
 170                         L"to end with a percentage sign.\n\n"
 171                         L"All of these options work with the critical threshold \"-c\" too."
 172                         , progName);
 173                 cout << endl;
 174                 return 0;
 175         }
 176
 177         if (vm.count("version"))
 178                 cout << "Version: " << VERSION << endl;
 179
 180         if (vm.count("warning")) {
 181                 try {
 182                         printInfo.warn = threshold(vm["warning"].as<wstring>());
 183                 } catch (std::invalid_argument& e) {
 184                         cout << e.what() << endl;
 185                         return 3;
 186                 }
 187         }
 188         if (vm.count("critical")) {
 189                 try {
 190                         printInfo.crit = threshold(vm["critical"].as<wstring>());
 191                 } catch (std::invalid_argument& e) {
 192                         cout << e.what() << endl;
 193                         return 3;
 194                 }
 195         }
 196
 197         if (vm.count("path"))
 198                 printInfo.drives = vm["path"].as<vector<wstring>>();
 199
 200         if (vm.count("unit")) {
 201                 try {
 202                         printInfo.unit = parseBUnit(vm["unit"].as<wstring>());
 203                 } catch (std::invalid_argument) {
 204                         wcout << L"Unknown unit Type " << vm["unit"].as<wstring>() << endl;
 205                         return 3;
 206                 }
 207         } else
 208                 printInfo.unit = BunitB;
 209
 210         if (vm.count("debug"))
 211                 debug = TRUE;
 212
 213         return -1;
 214 }
 215
 216 int printOutput(printInfoStruct& printInfo, vector<drive>& vDrives)
 217 {
 218         if (debug)
 219                 wcout << L"Constructing output string" << endl;
 220
 221         vector<wstring> wsDrives, wsPerf;
 222         wstring unit = BunitStr(printInfo.unit);
 223
 224         state state = OK;
 225         wstring output = L"DISK OK - free space:";
 226
 227         double tCap = 0, tFree = 0;
 228         for (vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); it++) {
 229                 tCap += it->cap;
 230                 tFree += it->free;
 231                 wsDrives.push_back(it->name + L" " + removeZero(it->free) + L" " + unit + L" (" +
 232                                                    removeZero(std::round(it->free/it->cap * 100.0)) + L"%); ");
 233                 wsPerf.push_back(L" " + it->name + L"=" + removeZero(it->free) + unit + L";" +
 234                                                  printInfo.warn.pString(it->cap) + L";" + printInfo.crit.pString(it->cap) + L";0;"
 235                                                  + removeZero(it->cap));
 236         }
 237
 238         if (printInfo.warn.rend(tFree, tCap)) {
 239                 state = WARNING;
 240                 output = L"DISK WARNING - free space:";
 241         }
 242
 243         if (printInfo.crit.rend(tFree, tCap)) {
 244                 state = CRITICAL;
 245                 output = L"DISK CRITICAL - free space:";
 246         }
 247
 248         wcout << output;
 249         if (vDrives.size() > 1)
 250                 wcout << L"Total " << tFree << unit << L" (" << removeZero(std::round(tFree/tCap * 100.0)) << L"%); ";
 251
 252         for (vector<wstring>::const_iterator it = wsDrives.begin(); it != wsDrives.end(); it++)
 253                 wcout << *it;
 254         wcout << L"|";
 255
 256         for (vector<wstring>::const_iterator it = wsPerf.begin(); it != wsPerf.end(); it++)
 257                 wcout << *it;
 258
 259         wcout << endl;
 260         return state;
 261 }
 262
 263 int check_drives(vector<drive>& vDrives)
 264 {
 265         DWORD dwResult, dwSize = 0, dwVolumePathNamesLen = MAX_PATH + 1;
 266         wchar_t szLogicalDrives[1024], szVolumeName[MAX_PATH], *szVolumePathNames = NULL;
 267         HANDLE hVolume = NULL;
 268         wstring wsLogicalDrives;
 269         size_t volumeNameEnd = 0;
 270
 271         set<wstring> sDrives;
 272
 273         if (debug)
 274                 wcout << L"Getting logic drive string (includes network drives)" << endl;
 275
 276         dwResult = GetLogicalDriveStrings(MAX_PATH, szLogicalDrives);
 277         if (dwResult > MAX_PATH)
 278                 goto die;
 279         if (debug)
 280                 wcout << L"Splitting string into single drive names" << endl;
 281
 282         LPTSTR szSingleDrive = szLogicalDrives;
 283         while (*szSingleDrive) {
 284                 wstring drname = szSingleDrive;
 285                 sDrives.insert(drname);
 286                 szSingleDrive += wcslen(szSingleDrive) + 1;
 287                 if (debug)
 288                         wcout << "Got: " << drname << endl;
 289         }
 290
 291         if (debug)
 292                 wcout << L"Getting volume mountpoints (includes NTFS folders)" << endl
 293                 << L"Getting first volume" << endl;
 294
 295         hVolume = FindFirstVolume(szVolumeName, MAX_PATH);
 296         if (hVolume == INVALID_HANDLE_VALUE)
 297                 goto die;
 298
 299         if (debug)
 300                 wcout << L"Traversing through list of drives" << endl;
 301
 302         while (GetLastError() != ERROR_NO_MORE_FILES) {
 303                 if (debug)
 304                         wcout << L"Path name for " << szVolumeName << L"= \"";
 305                 volumeNameEnd = wcslen(szVolumeName) - 1;
 306                 szVolumePathNames = reinterpret_cast<wchar_t*>(new WCHAR[dwVolumePathNamesLen]);
 307
 308                 while (!GetVolumePathNamesForVolumeName(szVolumeName, szVolumePathNames, dwVolumePathNamesLen, &dwVolumePathNamesLen)) {
 309                         if (GetLastError() != ERROR_MORE_DATA)
 310                                 break;
 311                         delete[] reinterpret_cast<wchar_t*>(szVolumePathNames);
 312                         szVolumePathNames = reinterpret_cast<wchar_t*>(new WCHAR[dwVolumePathNamesLen]);
 313
 314                 }
 315                 if (debug)
 316                         wcout << szVolumePathNames << L"\"" << endl;
 317
 318                 //.insert() does the dublicate checking
 319                 sDrives.insert(wstring(szVolumePathNames));
 320                 FindNextVolume(hVolume, szVolumeName, MAX_PATH);
 321         }
 322         if (debug)
 323                 wcout << L"Creating vector from found volumes, ignoring cd drives etc.:" << endl;
 324         for (set<wstring>::iterator it = sDrives.begin(); it != sDrives.end(); ++it) {
 325                 UINT type = GetDriveType(it->c_str());
 326                 if (type == DRIVE_FIXED || type == DRIVE_REMOTE) {
 327                         if (debug)
 328                                 wcout << L"\t" << *it << endl;
 329                         vDrives.push_back(drive(*it));
 330                 }
 331         }
 332
 333         FindVolumeClose(hVolume);
 334         if (szVolumePathNames)
 335                 delete[] reinterpret_cast<wchar_t*>(szVolumePathNames);
 336         return -1;
 337
 338 die:
 339         if (hVolume)
 340                 FindVolumeClose(hVolume);
 341         die();
 342         return 3;
 343 }
 344
 345 int check_drives(vector<drive>& vDrives, printInfoStruct& printInfo)
 346 {
 347         wchar_t *slash = L"\\";
 348
 349         if (debug)
 350                 wcout << L"Parsing user input drive names" << endl;
 351
 352         for (vector<wstring>::iterator it = printInfo.drives.begin();
 353                         it != printInfo.drives.end(); ++it) {
 354                 if (it->at(it->length() - 1) != *slash)
 355                         it->append(slash);
 356                 if (std::wstring::npos == it->find(L":\\")) {
 357                         wcout << "A \":\" is required after the drive name of " << *it << endl;
 358                         return 3;
 359                 }
 360                 if (debug)
 361                         wcout << L"Added " << *it << endl;
 362                 vDrives.push_back(drive(*it));
 363         }
 364         return -1;
 365 }
 366
 367 bool getFreeAndCap(drive& drive, const Bunit& unit)
 368 {
 369         if (debug)
 370                 wcout << L"Getting free disk space for drive " << drive.name << endl;
 371         ULARGE_INTEGER tempFree, tempTotal;
 372         if (!GetDiskFreeSpaceEx(drive.name.c_str(), NULL, &tempTotal, &tempFree)) {
 373                 return FALSE;
 374         }
 375         if (debug)
 376                 wcout << L"\tcap: " << tempFree.QuadPart << endl;
 377         drive.cap = round((tempTotal.QuadPart / pow(1024.0, unit)));
 378         if (debug)
 379                 wcout << L"\tAfter conversion: " << drive.cap << endl
 380                 << L"\tfree: " << tempFree.QuadPart << endl;
 381         drive.free = round((tempFree.QuadPart / pow(1024.0, unit)));
 382         if (debug)
 383                 wcout << L"\tAfter conversion: " << drive.free << endl << endl;
 384
 385         return TRUE;
 386 }