From d4553107a2a283c2cccd205f10e7a1cbdf6edf9d Mon Sep 17 00:00:00 2001
From: vvzvlad <git@vvzvlad.xyz>
Date: Tue, 13 Aug 2024 04:26:50 +0300
Subject: [PATCH] update

---
 shutdown_ups.py      | 54 ++++++++++++++++++++++++++------------------
 ups_shutdown.service |  2 +-
 2 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/shutdown_ups.py b/shutdown_ups.py
index 2ef7712..e9aa4b3 100644
--- a/shutdown_ups.py
+++ b/shutdown_ups.py
@@ -3,10 +3,7 @@ import time
 import requests
 import socket
 import os
-
-ups_ip = "10.31.41.46"
-time_threshold = 4
-interval = 20
+import sys
 
 def parse_time(time_str):
     try:
@@ -38,26 +35,39 @@ def run_command(command):
     result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
     return result.stdout.strip()
 
-while True:
-    oid_on_battery_time = "1.3.6.1.4.1.318.1.1.1.2.2.3.0"
-    result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_on_battery_time}')
-    time_on_battery = parse_time(result)
+def main(ups_ip, time_threshold, interval):
+    print("Starting UPS monitoring...")
+    while True:
+        oid_on_battery_time = "1.3.6.1.4.1.318.1.1.1.2.2.3.0"
+        result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_on_battery_time}')
+        time_on_battery = parse_time(result)
 
-    oid_states = "1.3.6.1.4.1.318.1.1.1.11.1.1.0"
-    result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_states}')
-    state = parse_states(result.strip('"'))
+        oid_states = "1.3.6.1.4.1.318.1.1.1.11.1.1.0"
+        result = run_command(f'/usr/bin/snmpget -v2c -O v -O q -c public {ups_ip} {oid_states}')
+        state = parse_states(result.strip('"'))
 
-    print(f"Time on battery: {time_on_battery} minutes, State: {state}")
+        print(f"Time on battery: {time_on_battery} minutes, State: {state}")
+
+        if time_on_battery != False and time_on_battery < time_threshold:
+            if state == "on-battery":
+                shutdown_status = run_command('shutdown --show 2>&1')
+                if "No scheduled shutdown" in shutdown_status:
+                    print("UPS is failing, shutting down in 60 seconds...")
+                    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")
+                    run_command('sudo shutdown -h 1 "Shutting down due to UPS failure"')
 
-    if time_on_battery != False and time_on_battery < time_threshold:
         if state == "on-battery":
-            shutdown_status = run_command('shutdown --show 2>&1')
-            if "No scheduled shutdown" in shutdown_status:
-                print("UPS is failing, shutting down in 60 seconds...")
-                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")
-                run_command('sudo shutdown -h 1 "Shutting down due to UPS failure"')
+            time.sleep(5)
+        else:
+            time.sleep(interval)
 
-    if state == "on-battery":
-        time.sleep(5)
-    else:
-        time.sleep(interval)
+if __name__ == "__main__":
+    if len(sys.argv) != 4:
+        print("Usage: script.py <ups_ip> <time_threshold> <interval>")
+        sys.exit(1)
+
+    ups_ip = sys.argv[1]
+    time_threshold = int(sys.argv[2])
+    interval = int(sys.argv[3])
+
+    main(ups_ip, time_threshold, interval)
diff --git a/ups_shutdown.service b/ups_shutdown.service
index a4b8bfd..26ebffe 100644
--- a/ups_shutdown.service
+++ b/ups_shutdown.service
@@ -3,7 +3,7 @@ Description=Service to manage UPS shutdown script
 After=network.target
 
 [Service]
-ExecStart=/usr/bin/python3 /usr/local/bin/shutdown_ups.py
+ExecStart=/usr/bin/python3 /usr/local/bin/shutdown_ups.py 10.31.41.46 4 20
 Restart=always
 RestartSec=5s