vps/monitoring.py

#!/usr/bin/python3

import datetime
import json
import os
import requests
import psutil
import time


# CPU usage threshold in percent
CPU_THRESHOLD = 50
# Memory usage threshold in GB
MEMORY_THRESHOLD = 3.25
# Disk usage threshold in GB
DISK_LOCAL_THRESHOLD = 25
DISK_SMB_THRESHOLD = 120

# Time in seconds between each check
MONITORING_INTERVAL = 30

# Discord Webhook URL
WEBHOOK_URL = os.environ["WEBHOOK_URL"]
# Minimum delay in seconds between alerts for a resource
WEBHOOK_INTERVAL = 600


def send_webhook(title, description):
    print(f'Sending Webhook with title "{title}".')
    data = {
        "username": "System Monitoring",
        "embeds": [
            {
                "title": title,
                "description": description,
                "timestamp": datetime.datetime.now(datetime.UTC).isoformat(),
            }
        ],
    }
    requests.post(
        WEBHOOK_URL, data=json.dumps(data), headers={"Content-Type": "application/json"}
    )


def check_cpu():
    cpu_percent = psutil.cpu_percent(interval=None)

    if (
        round(cpu_percent) > CPU_THRESHOLD
        and (time.time() - check_cpu.last_alert) > WEBHOOK_INTERVAL
    ):
        check_cpu.last_alert = time.time()
        send_webhook(
            "CPU Usage Alert",
            f"Average CPU usage over last {MONITORING_INTERVAL} seconds is past threshold of {CPU_THRESHOLD}%.\nCurrent usage: {cpu_percent}%.",
        )

    return f"{cpu_percent}%"


check_cpu.last_alert = 0


def check_memory():
    memory_info = psutil.virtual_memory()

    memory_used_bytes = memory_info.used
    memory_total_bytes = memory_info.total

    memory_used_gb = round(memory_used_bytes / (1024 * 1024 * 1024), 2)
    memory_total_gb = round(memory_total_bytes / (1024 * 1024 * 1024), 2)

    if (
        memory_used_gb > MEMORY_THRESHOLD
        and (time.time() - check_memory.last_alert) > WEBHOOK_INTERVAL
    ):
        check_memory.last_alert = time.time()
        send_webhook(
            "Memory Usage Alert",
            f"Memory usage is past threshold of {MEMORY_THRESHOLD} GB.\nCurrent usage: {memory_used_gb} GB / {memory_total_gb} GB.",
        )

    return f"{memory_used_gb} GB / {memory_total_gb} GB"


check_memory.last_alert = 0


def check_local_disk():
    disk_info = psutil.disk_usage("/")

    disk_used_bytes = disk_info.used
    disk_total_bytes = disk_info.total

    disk_used_gb = round(disk_used_bytes / (1024 * 1024 * 1024), 2)
    disk_total_gb = round(disk_total_bytes / (1024 * 1024 * 1024), 2)

    if (
        disk_used_gb > DISK_LOCAL_THRESHOLD
        and (time.time() - check_local_disk.last_alert) > WEBHOOK_INTERVAL
    ):
        check_local_disk.last_alert = time.time()
        send_webhook(
            "Local Disk Usage Alert",
            f"Local disk used space is past threshold of {DISK_LOCAL_THRESHOLD} GB.\nCurrent used space: {disk_used_gb} GB / {disk_total_gb} GB.",
        )

    return f"{disk_used_gb} GB / {disk_total_gb} GB"


check_local_disk.last_alert = 0


def check_smb_disk():
    disk_info = psutil.disk_usage("/mnt/storagebox")

    disk_used_bytes = disk_info.used
    disk_total_bytes = disk_info.total

    disk_used_gb = round(disk_used_bytes / (1024 * 1024 * 1024), 2)
    disk_total_gb = round(disk_total_bytes / (1024 * 1024 * 1024), 2)

    if (
        disk_used_gb > DISK_SMB_THRESHOLD
        and (time.time() - check_smb_disk.last_alert) > WEBHOOK_INTERVAL
    ):
        check_smb_disk.last_alert = time.time()
        send_webhook(
            "SMB Disk Usage Alert",
            f"SMB disk used space is past threshold of {DISK_SMB_THRESHOLD} GB.\nCurrent used space: {disk_used_gb} GB / {disk_total_gb} GB.",
        )

    return f"{disk_used_gb} GB / {disk_total_gb} GB"


check_smb_disk.last_alert = 0


def monitor():
    print(f"Refreshing every {MONITORING_INTERVAL} seconds.")

    while True:
        cpu_usage = check_cpu()
        memory_usage = check_memory()
        disk_local_usage = check_local_disk()
        disk_smb_usage = check_smb_disk()

        print(
            f"CPU usage: {cpu_usage}, Memory usage: {memory_usage}, Disk usage: Local = {disk_local_usage}, SMB = {disk_smb_usage}\r",
            end="\r",
        )

        time.sleep(MONITORING_INTERVAL)


if __name__ == "__main__":
    monitor()