From 8232fc2472557b8adb093bbad8fd5d8fd77c5f69 Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Fri, 21 Apr 2023 23:54:18 -0600
Subject: [PATCH] error handling

---
 check_bandwidth.py | 110 ++++++++++++++++++++++++++-------------------
 1 file changed, 64 insertions(+), 46 deletions(-)

diff --git a/check_bandwidth.py b/check_bandwidth.py
index 778f460..7c15d9a 100755
--- a/check_bandwidth.py
+++ b/check_bandwidth.py
@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 import argparse
+import sys
+import traceback
 
 import psutil
 
+import checker.nagios as nagios
 from checker.markdown import list_to_markdown_table
 
 # Parse the bandwidth limit argument
@@ -12,58 +15,73 @@ parser.add_argument('--critical', type=int, default=75, help='Critical if percen
 parser.add_argument('--warn', type=int, default=50, help='Warning if percent of bandwidth usage is greater than or equal to this.')
 args = parser.parse_args()
 
-data = []
-warn_value = (args.bandwidth * args.warn / 100)
-crit_value = (args.bandwidth * args.critical / 100)
 
-# Get network interface statistics
-net_io_counters = psutil.net_io_counters(pernic=True)
+def main():
+    data = []
+    warn_value = (args.bandwidth * args.warn / 100)
+    crit_value = (args.bandwidth * args.critical / 100)
 
-# Calculate bandwidth utilization for each interface
-for interface, stats in net_io_counters.items():
-    # Get the number of bytes sent and received
-    bytes_sent = stats.bytes_sent
-    bytes_recv = stats.bytes_recv
+    # Get network interface statistics
+    net_io_counters = psutil.net_io_counters(pernic=True)
 
-    # Wait for 1 second
-    psutil.cpu_percent(interval=1)
+    # Calculate bandwidth utilization for each interface
+    for interface, stats in net_io_counters.items():
+        # Get the number of bytes sent and received
+        bytes_sent = stats.bytes_sent
+        bytes_recv = stats.bytes_recv
 
-    # Get the number of bytes sent and received after 1 second
-    new_stats = psutil.net_io_counters(pernic=True)[interface]
-    new_bytes_sent = new_stats.bytes_sent
-    new_bytes_recv = new_stats.bytes_recv
+        # Wait for 1 second
+        psutil.cpu_percent(interval=1)
 
-    # Calculate the bandwidth utilization in bits per second
-    bandwidth_utilization = (8 * (new_bytes_sent - bytes_sent + new_bytes_recv - bytes_recv)) / (1 * 1000 * 1000)
-    data.append([interface, bandwidth_utilization, 'none'])
+        # Get the number of bytes sent and received after 1 second
+        new_stats = psutil.net_io_counters(pernic=True)[interface]
+        new_bytes_sent = new_stats.bytes_sent
+        new_bytes_recv = new_stats.bytes_recv
 
-critical = []
-warn = []
-ok = []
-perf_data = []
-for i in range(len(data)):
-    interface = data[i][0]
-    bandwidth_utilization = data[i][1]
-    if bandwidth_utilization >= crit_value:
-        critical.append(interface)
-        data[i][2] = 'critical'
-    elif bandwidth_utilization >= warn_value:
-        warn.append(interface)
-        data[i][2] = 'warning'
-    else:
-        ok.append(interface)
-        data[i][2] = 'ok'
-    perf_data.append(f'{interface}={round(bandwidth_utilization, 2)}MB;{warn_value};{crit_value}; ')
+        # Calculate the bandwidth utilization in bits per second
+        bandwidth_utilization = (8 * (new_bytes_sent - bytes_sent + new_bytes_recv - bytes_recv)) / (1 * 1000 * 1000)
+        data.append([interface, bandwidth_utilization, 'none'])
 
-if len(ok):
-    print(f'OK: {", ".join(ok)}')
-if len(warn):
-    print(f'WARNING: {", ".join(warn)}')
-if len(critical):
-    print(f'CRITICAL: {", ".join(critical)}')
+    exit_code = nagios.OK
+    critical = []
+    warn = []
+    ok = []
+    perf_data = []
+    for i in range(len(data)):
+        interface = data[i][0]
+        bandwidth_utilization = data[i][1]
+        if bandwidth_utilization >= crit_value:
+            critical.append(interface)
+            data[i][2] = 'critical'
+            exit_code = nagios.CRITICAL
+        elif bandwidth_utilization >= warn_value:
+            warn.append(interface)
+            data[i][2] = 'warning'
+            exit_code = nagios.WARNING
+        else:
+            ok.append(interface)
+            data[i][2] = 'ok'
+        perf_data.append(f'{interface}={round(bandwidth_utilization, 2)}MB;{warn_value};{crit_value}; ')
 
-data = [(x[0], f'{round(x[1], 2)} Mbps', x[2]) for x in data]
-data.insert(0, ('Interface', 'Bandwidth', 'State'))
-print(list_to_markdown_table(data, align='left'))
+    if len(ok):
+        print(f'OK: {", ".join(ok)}')
+    if len(warn):
+        print(f'WARNING: {", ".join(warn)}')
+    if len(critical):
+        print(f'CRITICAL: {", ".join(critical)}')
 
-print(f' |{"".join(perf_data)}')
+    data = [(x[0], f'{round(x[1], 2)} Mbps', x[2]) for x in data]
+    data.insert(0, ('Interface', 'Bandwidth', 'State'))
+    print(list_to_markdown_table(data, align='left'))
+
+    print(f' |{"".join(perf_data)}')
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        print(f'UNKNOWN: exception "{e}"')
+        print(traceback.format_exc())
+        sys.exit(nagios.UNKNOWN)