power/power.py

import time
import config
import logging
import smtplib
import asyncio
import discord
from nut2 import PyNUTClient
from proxmoxer import ProxmoxAPI
from email.mime.text import MIMEText
from logging.handlers import TimedRotatingFileHandler

ups_id = config.UPS_ID
gmail_user = config.GMAIL_USER
gmail_password = config.GMAIL_PASS
recipient_email = config.GMAIL_ADDR

bot_token = config.BOT_TOKEN
user_id = config.USER_ID

pve_host = config.PVE_HOST
pve_user = config.PVE_USER
pve_pass = config.PVE_PASS
pve_nodes = ['pve', 'c530']

def setup_logging():
    logger = logging.getLogger('powerlog')
    logger.setLevel(logging.INFO)
    handler = TimedRotatingFileHandler(
        'power.log',
        when='midnight',
        interval=1,
        backupCount=7,
        delay=True
    )
    formatter = logging.Formatter(
        '%(asctime)s [%(levelname)s] %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger

logger = setup_logging()

proxmox = ProxmoxAPI(pve_host, user=pve_user, password=pve_pass, verify_ssl=False)

def shutdown_pve(node):
    try:
        proxmox.nodes(node).status.post(command='shutdown')
        logger.warning("Node {node} is shutting down.")
    except Exception as e:
        logger.error(e)

def monitor_ups():
    client = PyNUTClient()
    try:
        ups_status = client.list_vars('ups')
        return {
            'status': ups_status.get('ups.status'),
            'battery_charge': float(ups_status.get('battery.charge'))
            }
    except Exception as e:
        logger.error(f"Error communicating with NUT: {e}")
        return None

def send_email(subject, body):
    msg = MIMEText(body)
    msg['Subject'] = subject
    msg['From'] = gmail_user
    msg['To'] = recipient_email
    try:
        with smtplib.SMTP('smtp.gmail.com', 587) as server:
            server.starttls()
            server.login(gmail_user, gmail_password)
            server.send_message(msg)
            logger.info(f"Email message sent: {subject}")
    except Exception as e:
        logger.info(f"Failed to send email: {e}")

intents = discord.Intents.default()
client = discord.Client(intents=intents)

def send_discord(message):
    async def message_send():
        intents = discord.Intents.default()
        client = discord.Client(intents=intents)
        @client.event
        async def on_ready():
            logger.info(f"Client connected as {client.user}")
            try:
                user = await client.fetch_user(user_id)
                await user.send(message)
                logger.info(f"Discord message sent: {message}")
            except Exception as e:
                logger.error(e)
            finally:
                await client.close()
        await client.start(bot_token)
    asyncio.run((message_send))

def pwr_offline(battery):
    message = f"{ups_id} UPS is running on battery power! {battery}% charge remaining."
    logger.warning(message)
    send_email(f"{ups_id}: Power Outage Detected!", message)
    send_discord(message)

def pwr_online(battery):
    message = f"{ups_id} UPS power has been restored. {battery}% charge remaining."
    logger.warning(message)
    send_email(f"{ups_id}: Power On Line", message)
    send_discord(message)

def batt_low(battery):
    message = f"{ups_id} battery level low: {battery}% charge remaining."
    logger.warning(message)
    send_email(f"{ups_id} battery low.", message)
    send_discord(message)

def batt_crit(battery):
    message = f"{ups_id} battery level critial: {battery}% charge remaining, Shutting down PVE."
    logger.warning(message)
    send_email(f"{ups_id} battery critical!!", message)
    send_discord(message)

def main():
    logger.info("Starting UPS monitoring service.")
    prev_status = None
    while True:
        ups_data = monitor_ups()
        if ups_data:
            status = ups_data['status']
            battery = ups_data['battery_charge']
            if status and status != prev_status:
                if status == "OB DISCHRG":
                    pwr_offline(battery)
                elif status == "OL":
                    pwr_online(battery)
                else:
                    logger.info(f"UPS status changed to: {status}")
                prev_status = status
            elif status == "OB DISCHRG" and battery < 50:
                if battery < 25:
                    batt_crit(battery)
                    for node in pve_nodes:
                        shutdown_pve(node)
                        time.sleep(300)
                else:
                    batt_low(battery)
                    time.sleep(60)
        time.sleep(5)

if __name__ == "__main__":
    main()