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