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

   1 /******************************************************************************
   2  * Icinga 2                                                                   *
   3  * Copyright (C) 2012-2014 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.0
  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 struct drive
  38 {
  39         wstring name;
  40         double cap, free;
  41         drive(wstring p)
  42                 : name(p)
  43         {}
  44 };
  45
  46 struct printInfoStruct
  47 {
  48         threshold warn, crit;
  49         vector<wstring> drives;
  50         Bunit unit;
  51 };
  52
  53 static int parseArguments(int, wchar_t **, po::variables_map&, printInfoStruct&);
  54 static int printOutput(printInfoStruct&, vector<drive>&);
  55 static int check_drives(vector<drive>&);
  56 static int check_drives(vector<drive>&, printInfoStruct&);
  57 static bool getFreeAndCap(drive&, const Bunit&);
  58
  59 int wmain(int argc, wchar_t **argv)
  60 {
  61         vector<drive> vDrives;
  62         printInfoStruct printInfo{ };
  63         po::variables_map vm;
  64
  65         int ret;
  66
  67         ret = parseArguments(argc, argv, vm, printInfo);
  68         if (ret != -1)
  69                 return ret;
  70
  71         printInfo.warn.legal = !printInfo.warn.legal;
  72         printInfo.crit.legal = !printInfo.crit.legal;
  73
  74         if (printInfo.drives.empty())
  75                 ret = check_drives(vDrives);
  76         else
  77                 ret = check_drives(vDrives, printInfo);
  78
  79         if (ret != -1)
  80                 return ret;
  81
  82         for (vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); ++it) {
  83                 if (!getFreeAndCap(*it, printInfo.unit)) {
  84                         wcout << L"Failed to access drive at " << it->name << endl;
  85                         return 3;
  86                 }
  87         }
  88
  89         return printOutput(printInfo, vDrives);
  90 }
  91
  92 int parseArguments(int ac, wchar_t **av, po::variables_map& vm, printInfoStruct& printInfo)
  93 {
  94         wchar_t namePath[MAX_PATH];
  95         GetModuleFileName(NULL, namePath, MAX_PATH);
  96         wchar_t *progName = PathFindFileName(namePath);
  97
  98         po::options_description desc("Options");
  99
 100         desc.add_options()
 101                 ("help,h", "print usage message and exit")
 102                 ("version,V", "print version and exit")
 103                 ("warning,w", po::wvalue<wstring>(), "warning threshold")
 104                 ("critical,c", po::wvalue<wstring>(), "critical threshold")
 105                 ("path,p", po::wvalue<vector<std::wstring>>()->multitoken(), "declare explicitly which drives to check (default checks all)")
 106                 ("unit,u", po::wvalue<wstring>(), "assign unit possible are: B, kB, MB, GB, TB")
 107                 ;
 108
 109         po::basic_command_line_parser<wchar_t> parser(ac, av);
 110
 111         try {
 112                 po::store(
 113                         parser
 114                         .options(desc)
 115                         .style(
 116                         po::command_line_style::unix_style |
 117                         po::command_line_style::allow_long_disguise)
 118                         .run(),
 119                         vm);
 120                 vm.notify();
 121         } catch (std::exception& e) {
 122                 cout << e.what() << endl << desc << endl;
 123                 return 3;
 124         }
 125
 126         if (vm.count("help")) {
 127                 wcout << progName << " Help\n\tVersion: " << VERSION << endl;
 128                 wprintf(
 129                         L"%s is a simple program to check a machines free disk space.\n"
 130                         L"You can use the following options to define its behaviour:\n\n", progName);
 131                 cout << desc;
 132                 wprintf(
 133                         L"\nIt will then output a string looking something like this:\n\n"
 134                         L"\tDISK WARNING 29GB|disk=29GB;50%%;5;0;120\n\n"
 135                         L"\"DISK\" being the type of the check, \"WARNING\" the returned status\n"
 136                         L"and \"23.8304%%\" is the returned value.\n"
 137                         L"The performance data is found behind the \"|\", in order:\n"
 138                         L"returned value, warning threshold, critical threshold, minimal value and,\n"
 139                         L"if applicable, the maximal value.\n"
 140                         L"This program will also print out additional performance data disk by disk\n\n"
 141                         L"%s' exit codes denote the following:\n\n"
 142                         L" 0\tOK,\n\tNo Thresholds were broken or the programs check part was not executed\n"
 143                         L" 1\tWARNING,\n\tThe warning, but not the critical threshold was broken\n"
 144                         L" 2\tCRITICAL,\n\tThe critical threshold was broken\n"
 145                         L" 3\tUNKNOWN, \n\tThe program experienced an internal or input error\n\n"
 146                         L"Threshold syntax:\n\n"
 147                         L"-w THRESHOLD\n"
 148                         L"warn if threshold is broken, which means VALUE < THRESHOLD\n\n"
 149                         L"-w !THRESHOLD\n"
 150                         L"inverts threshold check, VALUE > THRESHOLD (analogous to above)\n\n"
 151                         L"-w [THR1-THR2]\n"
 152                         L"warn is VALUE is inside the range spanned by THR1 and THR2\n\n"
 153                         L"-w ![THR1-THR2]\n"
 154                         L"warn if VALUE is outside the range spanned by THR1 and THR2\n\n"
 155                         L"-w THRESHOLD%%\n"
 156                         L"if the plugin accepts percentage based thresholds those will be used.\n"
 157                         L"Does nothing if the plugin does not accept percentages, or only uses\n"
 158                         L"percentage thresholds. Ranges can be used with \"%%\", but both range values need\n"
 159                         L"to end with a percentage sign.\n\n"
 160                         L"All of these options work with the critical threshold \"-c\" too."
 161                         , progName);
 162                 cout << endl;
 163                 return 0;
 164         }
 165
 166         if (vm.count("version"))
 167                 cout << "Version: " << VERSION << endl;
 168
 169         if (vm.count("warning")) {
 170                 try {
 171                         printInfo.warn = threshold(vm["warning"].as<wstring>());
 172                 } catch (std::invalid_argument& e) {
 173                         cout << e.what() << endl;
 174                         return 3;
 175                 }
 176         }
 177         if (vm.count("critical")) {
 178                 try {
 179                         printInfo.crit = threshold(vm["critical"].as<wstring>());
 180                 } catch (std::invalid_argument& e) {
 181                         cout << e.what() << endl;
 182                         return 3;
 183                 }
 184         }
 185
 186         if (vm.count("path"))
 187                 printInfo.drives = vm["path"].as<vector<wstring>>();
 188
 189         if (vm.count("unit")) {
 190                 try {
 191                         printInfo.unit = parseBUnit(vm["unit"].as<wstring>());
 192                 } catch (std::invalid_argument) {
 193                         wcout << L"Unknown unit Type " << vm["unit"].as<wstring>() << endl;
 194                         return 3;
 195                 }
 196         } else
 197                 printInfo.unit = BunitB;
 198
 199         return -1;
 200 }
 201
 202 int printOutput(printInfoStruct& printInfo, vector<drive>& vDrives)
 203 {
 204         state state = OK;
 205         double tCap = 0, tFree = 0;
 206         std::wstringstream perf, prePerf;
 207         wstring unit = BunitStr(printInfo.unit);
 208
 209         for (vector<drive>::iterator it = vDrives.begin(); it != vDrives.end(); ++it) {
 210                 tCap += it->cap; tFree += it->free;
 211                 perf << std::fixed << L" " << it->name << L"=" << removeZero(it->free) << unit << L";"
 212                         << printInfo.warn.pString() << L";" << printInfo.crit.pString() << L";0;" << removeZero(tCap);
 213         }
 214
 215         prePerf << L" | disk=" << removeZero(tFree) << unit << L";" << printInfo.warn.pString() << L";"
 216                 << printInfo.crit.pString() << L";0;" << removeZero(tCap);
 217
 218         if (printInfo.warn.perc) {
 219                 if (printInfo.warn.rend((tFree / tCap) * 100.0))
 220                         state = WARNING;
 221         } else {
 222                 if (printInfo.warn.rend(tFree))
 223                         state = WARNING;
 224         }
 225
 226         if (printInfo.crit.perc) {
 227                 if (printInfo.crit.rend((tFree / tCap) * 100.0))
 228                         state = CRITICAL;
 229         } else {
 230                 if (printInfo.crit.rend(tFree))
 231                         state = CRITICAL;
 232         }
 233
 234         switch (state) {
 235         case OK:
 236                 wcout << L"DISK OK " << tFree << unit << prePerf.str() << perf.str() << endl;
 237                 break;
 238         case WARNING:
 239                 wcout << L"DISK WARNING " << tFree << unit << prePerf.str() << perf.str() << endl;
 240                 break;
 241         case CRITICAL:
 242                 wcout << L"DISK CRITICAL " << tFree << unit << prePerf.str() << perf.str() << endl;
 243                 break;
 244         }
 245
 246         return state;
 247 }
 248
 249 int check_drives(vector<drive>& vDrives)
 250 {
 251         DWORD dwResult, dwSize = 0, dwVolumePathNamesLen = MAX_PATH + 1;
 252         wchar_t szLogicalDrives[MAX_PATH], szVolumeName[MAX_PATH], *szVolumePathNames;
 253         HANDLE hVolume;
 254         wstring wsLogicalDrives;
 255         size_t volumeNameEnd = 0;
 256
 257         set<wstring> sDrives;
 258
 259         dwResult = GetLogicalDriveStrings(MAX_PATH, szLogicalDrives);
 260         if (dwResult < 0 || dwResult > MAX_PATH)
 261                 goto die;
 262
 263         LPTSTR szSingleDrive = szLogicalDrives;
 264         while (*szSingleDrive) {
 265                 wstring drname = szSingleDrive;
 266                 sDrives.insert(drname);
 267                 szSingleDrive += wcslen(szSingleDrive) + 1;
 268         }
 269
 270         hVolume = FindFirstVolume(szVolumeName, MAX_PATH);
 271         if (hVolume == INVALID_HANDLE_VALUE)
 272                 goto die;
 273
 274         while (GetLastError() != ERROR_NO_MORE_FILES) {
 275                 volumeNameEnd = wcslen(szVolumeName) - 1;
 276                 szVolumePathNames = reinterpret_cast<wchar_t*>(new WCHAR[dwVolumePathNamesLen]);
 277
 278                 while (!GetVolumePathNamesForVolumeName(szVolumeName, szVolumePathNames, dwVolumePathNamesLen, &dwVolumePathNamesLen)) {
 279                         if (GetLastError() != ERROR_MORE_DATA)
 280                                 break;
 281                         delete[] reinterpret_cast<wchar_t*>(szVolumePathNames);
 282                         szVolumePathNames = reinterpret_cast<wchar_t*>(new WCHAR[dwVolumePathNamesLen]);
 283
 284                 }
 285
 286                 sDrives.insert(wstring(szVolumePathNames));
 287                 FindNextVolume(hVolume, szVolumeName, MAX_PATH);
 288         }
 289
 290         for (set<wstring>::iterator it = sDrives.begin(); it != sDrives.end(); ++it) {
 291                 UINT type = GetDriveType(it->c_str());
 292                 if (type == DRIVE_FIXED || type == DRIVE_REMOTE) {
 293                         vDrives.push_back(drive(*it));
 294                 }
 295         }
 296         return -1;
 297
 298 die:
 299         if (hVolume)
 300                 FindVolumeClose(hVolume);
 301         die();
 302         return 3;
 303 }
 304
 305 int check_drives(vector<drive>& vDrives, printInfoStruct& printInfo)
 306 {
 307         wchar_t *slash = L"\\";
 308
 309         for (vector<wstring>::iterator it = printInfo.drives.begin();
 310                         it != printInfo.drives.end(); ++it) {
 311                 if (it->at(it->length() - 1) != *slash)
 312                         it->append(slash);
 313
 314                 if (std::wstring::npos == it->find(L":\\")) {
 315                         wcout << "A \":\" is required after the drive name of " << *it << endl;
 316                         return 3;
 317                 }
 318                 vDrives.push_back(drive(*it));
 319         }
 320         return -1;
 321 }
 322
 323 bool getFreeAndCap(drive& drive, const Bunit& unit)
 324 {
 325         ULARGE_INTEGER tempFree, tempTotal;
 326         if (!GetDiskFreeSpaceEx(drive.name.c_str(), NULL, &tempTotal, &tempFree)) {
 327                 return FALSE;
 328         }
 329
 330         drive.cap = round((tempTotal.QuadPart / pow(1024.0, unit)));
 331         drive.free = round((tempFree.QuadPart / pow(1024.0, unit)));
 332
 333         return TRUE;
 334 }