refactoring

shutdown_ups.py

@@ -1,25 +1,31 @@
 #!/usr/bin/env python3
+
 import subprocess
 import time
-import requests
 import socket
 import os
 
-ups_ip = "10.31.41.46"
+UPS_IP = "10.31.41.46"
-time_threshold = 6
+TIME_THRESHOLD = 6
-interval = 60
+INTERVAL_LINE = 60
+INTERVAL_BATTERY = 5
+OID_ON_BATTERY_TIME = "1.3.6.1.4.1.318.1.1.1.2.2.3.0"
+OID_STATES = "1.3.6.1.4.1.318.1.1.1.11.1.1.0"
+OID_BATTERY_PERCENT = "1.3.6.1.4.1.318.1.1.1.2.2.1.0"
 
 def parse_time(time_str):
+    if time_str == "" or time_str is None:
+        return None
     try:
         time_parts = time_str.split(":")[-4:]
         days = int(time_parts[0])
         hours = int(time_parts[1])
         minutes = int(time_parts[2])
         seconds = float(time_parts[3])
-        total_minutes = (days * 24 * 60) + (hours * 60) + minutes + (int(seconds // 60))
+        return (days * 24 * 60) + (hours * 60) + minutes + (int(seconds // 60))
-        return total_minutes
+    except Exception as e:
-    except:
+        print(f"Error parsing time: {e}")
-        return False
+        return None
 
 def send_email(recipient, subject, body, sender):
     email_message = f"Subject: {subject}\nFrom: {sender}\nTo: {recipient}\n\n{body}"
@@ -27,44 +33,75 @@ def send_email(recipient, subject, body, sender):
     process.write(email_message)
     process.close()
 
-
 def parse_states(states_str):
-    if len(states_str) < 2: return "Invalid input" 
+    states_str = states_str.strip('"')
-    second_char = states_str[1]
+    if states_str == "": return "Invalid input"
-    if second_char == '0': return "on-line"
+    if len(states_str) < 64: return "Invalid input" 
-    elif second_char == '1': return "on-battery"
+    if states_str[1] == '0': return "on-line"
-    else: return "Invalid input" 
+    if states_str[1] == '1': return "on-battery"
+    return "Invalid input" 
 
 def run_command(command):
-    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    result = subprocess.run(f"{command}", shell=True,
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.STDOUT,
+                            text=True, check=False)
+    data = result.stdout.strip()
+    if "No such file or directory" in data:
+        print(f"Error: {data}", flush=True)
+        return None
     return result.stdout.strip()
 
+def snmp_get(oid):
+    result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {UPS_IP} {oid}')
+    if result is None or result == "":
+        return None
+    if "Timeout" in result:
+        print(f"UPS {UPS_IP} OID {oid} is unreachable, skipping...", flush=True)
+        return None
+    if "No Such Instance" in result:
+        print(f"UPS {UPS_IP} OID {oid} not found, skipping...", flush=True)
+        return None
+    if "No Such Object" in result:
+        print(f"UPS {UPS_IP} OID {oid} not found, skipping...", flush=True)
+        return None
+    if "No response" in result:
+        print(f"UPS {UPS_IP} OID {oid} not response, skipping...", flush=True)
+        return None
+    return result
+
 def main():
+    state = "on-line"
     while True:
-        oid_on_battery_time = "1.3.6.1.4.1.318.1.1.1.2.2.3.0"
+        if state == "on-battery": time.sleep(INTERVAL_BATTERY)
-        result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_on_battery_time}')
+        if state == "on-line": time.sleep(INTERVAL_LINE)
+
+        result = snmp_get(OID_ON_BATTERY_TIME)
+        if result is None: continue
         time_on_battery = parse_time(result)
 
-        oid_states = "1.3.6.1.4.1.318.1.1.1.11.1.1.0"
+        result = snmp_get(OID_STATES)
-        result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_states}')
+        if result is None: continue
-        state = parse_states(result.strip('"'))
+        state = parse_states(result)
 
-        print(f"Time on battery: {time_on_battery} minutes, State: {state}", flush=True)
+        result = snmp_get(OID_BATTERY_PERCENT)
+        if result is None: continue
+        battery_percent = int(result)
 
-        shutdown_status = run_command('shutdown --show 2>&1')
+        print(f"UPS percent: {battery_percent}, time on battery: {time_on_battery} minutes, state: {state}", flush=True)
+
+        shutdown_status = run_command('shutdown --show')
         
-        if (state == "on-battery") and (time_on_battery > time_threshold) and ("No scheduled shutdown" in shutdown_status):
+        if (state == "on-battery") and (time_on_battery > TIME_THRESHOLD) and ("No scheduled shutdown" in shutdown_status):
             print("UPS is failing, shutting down in 60 seconds...", flush=True)
-            send_email(f"{socket.gethostname()}@vvzvlad.xyz", "UPS is failing, shutting down in 60 seconds...", f"Proxmox node '{socket.gethostname()}': UPS is failing, shutting down in 60 seconds...", f"{socket.gethostname()}@vvzvlad.xyz")
+            send_email(f"{socket.gethostname()}@vvzvlad.xyz", "UPS is failing, shutting down in 60 seconds...",
-            run_command('sudo shutdown -h 3 "Shutting down due to UPS failure"')
+                        f"Proxmox node '{socket.gethostname()}': UPS is failing, shutting down in 60 seconds...",
+                        f"{socket.gethostname()}@vvzvlad.xyz")
+            run_command('shutdown -h 3 "Shutting down due to UPS failure"')
+
         if (state == "on-line") and ("Shutdown scheduled" in shutdown_status):
             print("UPS is back online, cancelling shutdown...", flush=True)
-            run_command('sudo shutdown -c')
+            run_command('shutdown -c')
-
-        if state == "on-battery":
-            time.sleep(5)
-        else:
-            time.sleep(interval)
 
 if __name__ == "__main__":
     print("Starting UPS monitoring...", flush=True)