import subprocess
import time
import requests
import socket
import os

ups_ip = "10.31.41.46"
time_threshold = 4
interval = 20

def parse_time(time_str):
    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 total_minutes
    except:
        return False

def send_email(recipient, subject, body, sender):
    email_message = f"Subject: {subject}\nFrom: {sender}\nTo: {recipient}\n\n{body}"
    process = os.popen('/usr/sbin/sendmail -t', 'w')
    process.write(email_message)
    process.close()


def parse_states(states_str):
    if len(states_str) < 2: return "Invalid input" 
    second_char = states_str[1]
    if second_char == '0': return "on-line"
    elif second_char == '1': return "on-battery"
    else: return "Invalid input" 

def run_command(command):
    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return result.stdout.strip()

def main():
    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('"'))

        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 state == "on-battery":
            time.sleep(5)
        else:
            time.sleep(interval)

if __name__ == "__main__":
    main()