power/power.py

import time
import config
import logging
import smtplib
import asyncio
import discord
import paramiko
from ping3 import ping
from nut2 import PyNUTClient
from proxmoxer import ProxmoxAPI
from email.mime.text import MIMEText
from wakeonlan import send_magic_packet
from logging.handlers import TimedRotatingFileHandler

# Config
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
servers = ['pve', 'c530']
nas_ip = config.HOST1_IP

wol_targets = {
    'host1': {'mac': config.HOST1_MAC, 'ip': config.HOST1_IP},
    'host2': {'mac': config.HOST2_MAC, 'ip': config.HOST2_IP},
    'host3': {'mac': config.HOST3_MAC, 'ip': config.HOST3_IP}
}

# Log Setup
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()

# Email
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)

# Discord
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())

# UPS Monitor
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

# Server Power Management
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(f"Node {node} is shutting down.")
    except Exception as e:
        logger.error(e)

def shutdown_nas():
    try:
        logger.warning("Nas is shutting down.")
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(nas_ip)
        stdin, stdout, stderr = client.exec_command('sudo shutdown')
        for line in stdout:
            print(line.strip())
        client.close()
    except Exception as e:
        logger.error(e)

def wake_up(battery):
        for target, addr in wol_targets.items():
            mac = addr['mac']
            ip = addr['ip']
            while True:
                response = ping(ip)
                if response:
                    logger.info(response)
                    logger.info(f"Host {mac} at {ip} is up. Ping {response}.")
                    time.sleep(5)
                    break
                else:
                    send_magic_packet(mac)
                    logger.warning(f"Host {mac} at {ip} is offline. Sending magic packet.")
                    time.sleep(10)

# Status Communications
def pwr_online(battery):
    message = f"{ups_id} UPS power has been restored. {battery}% charge remaining."
    logger.info(message)
    send_email(f"{ups_id}: Power On Line", message)
    send_discord(message)
    if ups_id == "Server" and battery > 75:
        wake_up(battery)

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 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)
    if ups_id == "Server":
        for node in servers:
            shutdown_pve(node)
    time.sleep(120)
    shutdown_nas()

# Main Loop
def main():
    logger.info("Starting UPS monitoring service.")
    prev_status = "OL"
    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" or status == "OL CHRG":
                    pwr_online(battery)                       
                else:
                    logger.info(f"UPS status changed to: {status}")
                prev_status = status
            elif status == "OB DISCHRG" and battery < 70:
                if battery < 65:
                    batt_crit(battery)
                    if ups_id == "Server":
                        for node in servers:
                            shutdown_pve(node)
                    time.sleep(300)
                else:
                    batt_low(battery)
                    time.sleep(60)
        time.sleep(5)

if __name__ == "__main__":
    main()