lego-monitoring/service.py
2024-08-17 00:00:16 +03:00

68 lines
1.9 KiB
Python
Executable file

#!/usr/bin/env python3
import asyncio
import logging
import signal
from typing import Callable, Coroutine
import nio
from alerting import alerts
from misc import checks
logging.basicConfig(level=logging.INFO)
stopping = False
def stop_gracefully(signum, frame):
global stopping
stopping = True
async def checker(check: Callable | Coroutine, interval_secs: int, client: nio.AsyncClient, *args, **kwargs):
while True:
logging.info(f"Calling {check.__name__}")
if isinstance(check, Callable):
result = check(*args, **kwargs)
if isinstance(result, Coroutine):
result = await result
elif isinstance(check, Coroutine):
result = await check
else:
raise TypeError(f"check is {type(check)}, neither function nor coroutine")
logging.info(f"Got {len(result)} alerts")
for alert in result:
await alerts.send_alert(alert, client)
await asyncio.sleep(interval_secs)
async def main():
MINUTE = 60
HOUR = 60 * MINUTE
DAY = 24 * HOUR
WEEK = 7 * DAY
signal.signal(signal.SIGTERM, stop_gracefully)
client = await alerts.get_client()
checkers = (
checker(checks.temp_check, 5 * MINUTE, client),
checker(checks.cpu_check, 5 * MINUTE, client),
checker(checks.ups_check, 5 * MINUTE, client),
checker(checks.ram_check, 1 * MINUTE, client),
checker(checks.vuln_check, 1 * DAY, client),
)
async with asyncio.TaskGroup() as tg:
checker_tasks: set[asyncio.Task] = set()
for c in checkers:
task = tg.create_task(c)
checker_tasks.add(task)
while True:
if stopping:
await client.close()
raise SystemExit
else:
await asyncio.sleep(3)
if __name__ == "__main__":
asyncio.run(main())