[WIP] refactor for adding new flags

.pylintrc

@@ -7,6 +7,7 @@ disable=fixme,
         too-many-locals,
         too-many-statements,
         too-many-branches,
+        too-many-return-statements,
         bare-except,
         missing-module-docstring,
         missing-function-docstring,

check_brevisone

@@ -42,7 +42,7 @@
 # works based on those contributions, and sublicense and distribute
 # those contributions and any derivatives thereof.
 
-from argparse import ArgumentParser
+from argparse import ArgumentParser, ArgumentTypeError
 import urllib.request
 from urllib.parse import urljoin
 import ssl
@@ -60,48 +60,30 @@ WARNING = 1
 CRITICAL = 2
 UNKNOWN = 3
 
-__version__ = '3.0.0'
+__version__ = '3.1.0'
 
 
-def read_int(string):
-    """
-    Get and sanitize integer data
-    """
-    try:
-        return int(string.split(':')[1])
-    except ValueError:
-        return 0
-
-
-def read_signal(string):
-    """
-    Get and sanitize the signal data
-    """
-    sig = ':'.join(string.split(':')[1:]).strip()
-    sig = sig.replace('db', '').replace('dBm', '').strip()
-
-    try:
-        return int(sig)
-    except ValueError:
-        return 0
-
-
-def generate_output(status='UNKNOWN', lines=None, perfdata=None):
+def generate_output(status='UNKNOWN', perfdata=None):
     """
     Generate plugin data output with status and perfdata
     """
 
     pluginoutput = str(status)
 
     # All lines we received from the Endpoint
-    if lines:
-        pluginoutput += ' - ' + ' '.join(lines)
+    if perfdata:
+        pluginoutput += ' - '
+        for k,v in perfdata.items():
+            pluginoutput += '{}: {} '.format(k,v)
 
     # Perfdata we explicitly extracted from the data
     # The perfdata var should be a dict, we normalize the keys and
     # transform everything into the expected output 'key 1': 3 > key_1=3
     if perfdata:
-        pluginoutput += '|' + ' '.join([key.lower().replace(" ", "_").replace(",", "") + '=' + str(value) for key, value in perfdata.items()])
+        pluginoutput += '|'
+        for k,v in perfdata.items():
+            if (k and v) and isinstance(v, int):
+                pluginoutput += k.lower().replace(" ", "_").replace(",", "") + '=' + str(v) + ' '
 
     print(pluginoutput)
 
@@ -130,6 +112,14 @@ def commandline(args):
                         help='The critical threshold for the minimum signal strength (in db, default -107)',
                         default=-107,
                         type=int)
+    parser.add_argument('--disk-warning',
+                        help='The warning threshold for the disk space (in bytes)',
+                        required=False,
+                        type=int)
+    parser.add_argument('--disk-critical',
+                        help='The critical threshold for the disk space (in bytes)',
+                        required=False,
+                        type=int)
     parser.add_argument('--ssl-insecure',
                         dest='insecure',
                         action='store_true',
@@ -140,7 +130,17 @@ def commandline(args):
                         default='https',
                         help='HTTP protocol, use one of http or https (default https)')
 
-    return parser.parse_args(args)
+    args = parser.parse_args(args)
+
+    if args.disk_warning and not args.disk_critical:
+        parser.print_help()
+        raise ArgumentTypeError("%s: error: --disk-warning requires --disk-critical" % parser.prog)
+
+    if args.disk_critical and not args.disk_warning:
+        parser.print_help()
+        raise ArgumentTypeError("%s: error: --disk-critical requires --disk-warning" % parser.prog)
+
+    return args
 
 
 def get_data(base_url, timeout, insecure):
@@ -152,7 +152,7 @@ def get_data(base_url, timeout, insecure):
     HTTP/1.0 200 OK
     que: foo
     failed: 0
-    signal_strength: 15 db
+    signal: 15 db
     total: 25
     """
 
@@ -177,21 +177,31 @@ def get_data(base_url, timeout, insecure):
 
     return resp
 
-def extract_perfdata(lines):
+def parse_data(data):
     """
-    Safely extract perfdata
+    Safely extract data from the APIs reponse
     """
-    if len(lines) < 4:
-        return None
+    lines = [str(i).strip() for i in data.split("\n") if i]
+
+    parsed_data = {}
+
+    for l in lines:
+        d = l.split(":")
+        if len(d) == 2:
+            key = d[0].strip()
+            value = d[1].strip()
+
+            # Remove the db string from the signal value
+            if key == "signal":
+                value = value.replace('db', '').replace('dBm', '').strip()
 
-    perfdata = {
-        'que': read_int(lines[0]),
-        'failed': read_int(lines[1]),
-        'signal': read_signal(lines[2]),
-        'total': read_int(lines[3])
-    }
+            # Parse integer value to be actual integers
+            if value.lstrip('-').isdigit():
+                value = int(value)
 
-    return perfdata
+            parsed_data[key] = value
+
+    return parsed_data
 
 def main(args):
     try:
@@ -203,25 +213,31 @@ def main(args):
         print('UNKNOWN - Could not connect to SMS Gateway', data_exc)
         return UNKNOWN
 
-    # Split up lines and clean up items
-    lines = [str(i).strip() for i in data.split("\n") if i]
     # Safely extract perfdata, which we'll also use to determine the final status
-    perfdata = extract_perfdata(lines)
+    perfdata = parse_data(data)
 
     if not perfdata:
         print('UNKNOWN - Could not determine status', perfdata)
         return UNKNOWN
 
+    if 'failed' not in perfdata:
+        print('UNKNOWN - Could not determine failed status', perfdata)
+        return UNKNOWN
+
+    if 'que' not in perfdata:
+        print('UNKNOWN - Could not determine que status', perfdata)
+        return UNKNOWN
+
     # Determine the final status
     if perfdata['failed'] >= args.fail or perfdata['signal'] <= args.signal_critical:
-        generate_output(status='CRITICAL', lines=lines, perfdata=perfdata)
+        generate_output(status='CRITICAL', perfdata=perfdata)
         return CRITICAL
 
     if perfdata['que'] >= args.queue or perfdata['signal'] <= args.signal_warning:
-        generate_output(status='WARNING', lines=lines, perfdata=perfdata)
+        generate_output(status='WARNING', perfdata=perfdata)
         return WARNING
 
-    generate_output(status='OK', lines=lines, perfdata=perfdata)
+    generate_output(status='OK', perfdata=perfdata)
     return OK
 
 

test_check_brevisone.py

@@ -18,11 +18,11 @@ def test_output(self, mock_print):
         generate_output(status='CRITICAL')
         mock_print.assert_called_with("CRITICAL")
 
-        generate_output(status='OK', lines=['1: 2', '3: 4'], perfdata={'1': '2', '3': 4})
-        mock_print.assert_called_with("OK - 1: 2 3: 4|1=2 3=4")
+        generate_output(status='OK', perfdata={'1': '2', 'foo': 4})
+        mock_print.assert_called_with("OK - 1: 2 foo: 4 |foo=4 ")
 
-        generate_output(status='CRITICAL', lines=['que: foo', 'foo bar: 1'], perfdata={'que': 'foo', 'foo bar': 1})
-        mock_print.assert_called_with("CRITICAL - que: foo foo bar: 1|que=foo foo_bar=1")
+        generate_output(status='CRITICAL', perfdata={'que': '0', 'foo bar': 1})
+        mock_print.assert_called_with("CRITICAL - que: 0 foo bar: 1 |foo_bar=1 ")
 
 class CLITesting(unittest.TestCase):
 
@@ -88,9 +88,9 @@ def test_main_unknown(self, mock_data):
     @mock.patch('check_brevisone.get_data')
     def test_main_ok(self, mock_data):
         d = """
-        que: foo
+        que: 0
         failed: 0
-        signal_strength: 15 db
+        signal: 15 db
         total: 25
         """
         mock_data.return_value = d
@@ -102,9 +102,9 @@ def test_main_ok(self, mock_data):
     @mock.patch('check_brevisone.get_data')
     def test_main_warn(self, mock_data):
         d = """
-        que: foo
+        que: 1
         failed: 0
-        signal_strength: -91 db
+        signal: -91 db
         total: 25
         """
         mock_data.return_value = d
@@ -116,7 +116,7 @@ def test_main_warn(self, mock_data):
     @mock.patch('check_brevisone.get_data')
     def test_main_crit(self, mock_data):
         d = """
-        que: foo
+        que: 0
         failed: 10
         signal_strength: -91 db
         total: 25