#!/usr/bin/env python3

import subprocess
import time
import socket
import os

UPS_IP = "10.31.41.46"
TIME_THRESHOLD = 6
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])
        return (days * 24 * 60) + (hours * 60) + minutes + (int(seconds // 60))
    except Exception as e:
        print(f"Error parsing time: {e}")
        return None

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):
    states_str = states_str.strip('"')
    if states_str == "": return "Invalid input"
    if len(states_str) < 64: return "Invalid input" 
    if states_str[1] == '0': return "on-line"
    if states_str[1] == '1': return "on-battery"
    return "Invalid input" 

def run_command(command):
    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:
        if state == "on-battery": time.sleep(INTERVAL_BATTERY)
        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)

        result = snmp_get(OID_STATES)
        if result is None: continue
        state = parse_states(result)

        result = snmp_get(OID_BATTERY_PERCENT)
        if result is None: continue
        battery_percent = int(result)

        print(f"Batterry {battery_percent}%, {state}, estimated time on battery: {time_on_battery} minutes", 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):
            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")
            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('shutdown -c')

if __name__ == "__main__":
    print("Starting UPS monitoring...", flush=True)
    main()