mem_checker/mem_checker/mem_checker.py

"""
A script for monitoring available RAM and free disk space.
It uses the free and df utilites.
When the threshold is exceeded, it sends a post request to the api.
The request contains the following json format:
    [
        {
            type_msg: str,
            message: str
            detail: str
        }
    ]

:copyright: Sergey Vanyushkin <pi3c@yandex.ru>
:git: https://git.pi3c.ru/pi3c/mem_checker.git
:license: MIT
2024г.
"""

import json
import logging
import subprocess
import sys
import urllib.parse
import urllib.request
from abc import ABC, abstractmethod
from dataclasses import dataclass
from time import sleep

ENDPOINT_URL = "http://127.0.0.1:8000/user"  # endpoint for send alerts
FREE_IN_MB = 500  # free: the value of the trigger in megabytes
FREE_IN_GB = 0.5  # free: the value of the trigger in gigabytes
DF_IN_PERCENT = 50  # df: the value of the trigger in percents
SLEEPTIME = 5  # main_loop sleep time per seconds


logger = logging.getLogger("mem_checker_logger")
logger.setLevel(logging.DEBUG)
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
console_handler.setFormatter(formatter)

# Добавляем обработчик в логгер
logger.addHandler(console_handler)


@dataclass()
class MessageDTO:
    type_msg: str
    message: str
    detail: str


class Command(ABC):
    @property
    def command(self):
        return self.__command

    @command.setter
    def command(self, cmd: list):
        self.__command = cmd

    @abstractmethod
    def _stdout_parser(self, stdout: str) -> list[MessageDTO]:
        raise NotImplementedError

    @abstractmethod
    def _filter_alerts(self, stat: list[MessageDTO]) -> list[MessageDTO]:
        raise NotImplementedError

    def _execute(self, command: list) -> tuple[str, str]:
        try:
            result = subprocess.run(
                command,
                capture_output=True,
                check=True,
                text=True,
            )
        except Exception as e:
            return "", str(e)

        return result.stdout, result.stderr

    def get_alerts(self) -> list[MessageDTO]:
        stdout, stderr = self._execute(command=self.command)
        if stderr:
            return [
                MessageDTO(
                    type_msg="Error",
                    message="Free command execution error",
                    detail=stderr,
                )
            ]
        stat = self._stdout_parser(stdout=stdout)
        return self._filter_alerts(stat)


class FreeMem(Command):
    def __init__(
        self,
    ) -> None:
        self.command = ["free", "-h"]

    def _filter_alerts(self, stat: list[MessageDTO]) -> list[MessageDTO]:
        filtered_alerts = []
        for obj in stat:
            if obj.type_msg == "Error":
                filtered_alerts.append(obj)
            elif obj.type_msg == "Status":
                val, unit = obj.detail.split()
                if (float(val) < FREE_IN_MB and unit == "Mi") or (float(val) < FREE_IN_GB and unit == "Gi"):
                    obj.type_msg = "Alert"
                    filtered_alerts.append(obj)
        return filtered_alerts

    def _stdout_parser(self, stdout: str) -> list[MessageDTO]:
        try:
            memory = stdout.split()[12]
            memory_units = memory[-2:]
            memory = memory[:-2].replace(",", ".")
            memory_value = float(memory) if "." in memory else int(memory)
        except Exception as e:
            logger.error(e)
            return [
                MessageDTO(
                    type_msg="ParserError",
                    message="Stdout parsing runtime error",
                    detail=str(e),
                ),
            ]
        return [
            MessageDTO(
                type_msg="Status",
                message="free mem available",
                detail=f"{memory_value} {memory_units}",
            )
        ]


class FreeHdd(Command):
    def __init__(self) -> None:
        self.command = ["df", "-h"]

    def _filter_alerts(self, stat: list[MessageDTO]) -> list[MessageDTO]:
        filtered_alerts = list()
        for obj in stat:
            persents = int(obj.detail.split()[4][:-1])
            if DF_IN_PERCENT < persents:
                obj.type_msg = "Alert"
                filtered_alerts.append(obj)
        return filtered_alerts

    def _stdout_parser(self, stdout) -> list[MessageDTO]:
        def row_to_dto(row: list):
            return MessageDTO(type_msg="Status", message="hdd mem available", detail=" ".join(row))

        stat = [row.split() for row in stdout.split("\n")]
        stat = list(map(row_to_dto, stat[1:-1]))
        return self._filter_alerts(stat)


def send_alert(alerts: list[MessageDTO]) -> None:
    data = [{"type_msg": item.type_msg, "message": item.message, "detail": item.detail} for item in alerts]

    json_data = json.dumps(data).encode("utf-8")

    req = urllib.request.Request(
        ENDPOINT_URL,
        data=json_data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req) as response:
            response.read().decode("utf-8")
    except Exception as e:
        logger.error(e)


def main_loop() -> None:
    free_mem = FreeMem()
    hdd_mem = FreeHdd()

    alerts: list[MessageDTO] = list()

    alerts.extend(free_mem.get_alerts())
    alerts.extend(hdd_mem.get_alerts())

    if alerts:
        send_alert(alerts=alerts)


if __name__ == "__main__":
    if "--test-mode" in sys.argv:
        while True:
            main_loop()
            sleep(SLEEPTIME)
    main_loop()