From ffdd0429b3048404a43c7e9c140695a5a8b48d1d Mon Sep 17 00:00:00 2001 From: Alex Tau Date: Wed, 30 Apr 2025 00:16:27 +0300 Subject: [PATCH] some actual alerts and telegram client --- flake.nix | 6 +++ modules/default.nix | 28 ++++++---- pyproject.toml | 1 + src/lego_monitoring/__init__.py | 54 +++++++++++++++++-- src/lego_monitoring/alerting/alerts.py | 73 ++++++++++++++++++++++++++ src/lego_monitoring/alerting/enum.py | 23 ++++++++ src/lego_monitoring/core/checkers.py | 53 +++++++++++++++++++ src/lego_monitoring/core/config.py | 43 +++++++++------ src/lego_monitoring/core/cvars.py | 8 +++ uv.lock | 58 ++++++++++++++++++-- 10 files changed, 314 insertions(+), 33 deletions(-) create mode 100644 src/lego_monitoring/alerting/alerts.py create mode 100644 src/lego_monitoring/alerting/enum.py create mode 100644 src/lego_monitoring/core/checkers.py create mode 100644 src/lego_monitoring/core/cvars.py diff --git a/flake.nix b/flake.nix index 292a023..1d8e21e 100644 --- a/flake.nix +++ b/flake.nix @@ -59,6 +59,12 @@ # Implement build fixups here. # Note that uv2nix is _not_ using Nixpkgs buildPythonPackage. # It's using https://pyproject-nix.github.io/pyproject.nix/build.html + + pyaes = _prev.pyaes.overrideAttrs ( + old: { + buildInputs = old.buildInputs or [ ] ++ [ _prev.setuptools ]; + } + ); }; # This example is only using x86_64-linux diff --git a/modules/default.nix b/modules/default.nix index d5d9d59..8808560 100644 --- a/modules/default.nix +++ b/modules/default.nix @@ -11,16 +11,21 @@ package: options.services.lego-monitoring = { enable = lib.mkEnableOption "lego-monitoring service."; - nonSecretConfigOption = lib.mkOption { - type = lib.types.str; - default = "defaultValue"; - description = "An example non-secret config option."; + enabledCheckerSets = lib.mkOption { + type = lib.types.listOf lib.types.str; + default = [ ]; + description = "List of enabled checker sets. Each checker set is a module which checks something and generates alerts based on check results."; }; - configOptionSecretPath = lib.mkOption { - type = lib.types.nullOr lib.types.str; - default = null; - description = "Path to an example secret config option."; + telegram = { + credsSecretPath = lib.mkOption { + type = lib.types.str; + description = "Path to a file containing Telegram api_id, api_hash, and bot token, separated by the `,` character."; + }; + roomId = lib.mkOption { + type = lib.types.int; + description = "ID of chat where to send alerts."; + }; }; }; @@ -28,8 +33,11 @@ package: cfg = config.services.lego-monitoring; json = pkgs.formats.json {}; serviceConfigFile = json.generate "config.json" { - non_secret_config_option = cfg.nonSecretConfigOption; - config_option_secret_path = cfg.configOptionSecretPath; + enabled_checker_sets = cfg.enabledCheckerSets; + telegram = with cfg.telegram; { + creds_secret_path = credsSecretPath; + room_id = roomId; + }; }; in lib.mkIf cfg.enable { systemd.services.lego-monitoring = { diff --git a/pyproject.toml b/pyproject.toml index ed0e578..11bdaf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ readme = "README.md" requires-python = ">=3.12" dependencies = [ "alt-utils>=0.0.6", + "telethon>=1.40.0", ] [project.scripts] diff --git a/src/lego_monitoring/__init__.py b/src/lego_monitoring/__init__.py index 15bd42e..bd1d8f0 100644 --- a/src/lego_monitoring/__init__.py +++ b/src/lego_monitoring/__init__.py @@ -1,19 +1,65 @@ import argparse +import asyncio +import logging +import signal import time +from .alerting import alerts +from .core import cvars from .core.config import load_config +stopping = False + + +def stop_gracefully(signum, frame): + global stopping + stopping = True + def main() -> None: + logging.basicConfig(level=logging.INFO) + + asyncio.run(async_main()) + + +async def async_main(): parser = argparse.ArgumentParser( prog="lego-monitoring", description="Lego-monitoring service", ) - parser.add_argument('-c', '--config', required=True) + parser.add_argument("-c", "--config", required=True) config_path = parser.parse_args().config config = load_config(config_path) + cvars.config.set(config) - while True: - print(f"service running... opt 1 is {config.non_secret_config_option}, opt 2 is secret, but if you really wanna know, it's {config.config_option}", flush=True) - time.sleep(300) + tg_client = await alerts.get_client() + cvars.tg_client.set(tg_client) + + checker_sets = { + "start": [ + alerts.send_start_alert(), + ], + "stop": [], # this is checked later + } + + checkers = [] + for enabled_set in config.enabled_checker_sets: + for checker in checker_sets[enabled_set]: + checkers.append(checker) + + signal.signal(signal.SIGTERM, stop_gracefully) + + 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: + if "stop" in config.enabled_checker_sets: + await alerts.send_stop_alert() + await tg_client.disconnect() + raise SystemExit + else: + await asyncio.sleep(3) diff --git a/src/lego_monitoring/alerting/alerts.py b/src/lego_monitoring/alerting/alerts.py new file mode 100644 index 0000000..e5b4ccf --- /dev/null +++ b/src/lego_monitoring/alerting/alerts.py @@ -0,0 +1,73 @@ +from dataclasses import dataclass + +from telethon import TelegramClient +from telethon.sessions import MemorySession + +from ..core import cvars +from .enum import AlertType, Severity + + +@dataclass +class Alert: + alert_type: AlertType + message: str + severity: Severity + + +async def get_client() -> TelegramClient: + config = cvars.config.get() + api_id, api_hash, bot_token = config.telegram.creds.split(",") + client = await TelegramClient(MemorySession(), api_id, api_hash).start(bot_token=bot_token) + client.parse_mode = "html" + return client + + +def format_message(alert: Alert) -> str: + match alert.severity: + case Severity.INFO: + severity_emoji = "ℹī¸" + case Severity.WARNING: + severity_emoji = "⚠ī¸" + case Severity.CRITICAL: + severity_emoji = "🆘" + message = f"{severity_emoji} {alert.alert_type} Alert\n{alert.message}" + return message + + +async def send_alert(alert: Alert) -> None: + try: + client = cvars.tg_client.get() + except LookupError: # being called standalone + # cvars.config.set(get_config()) + # temp_client = True + # client = await get_client() + # cvars.matrix_client.set(client) + raise NotImplementedError # TODO + else: + ... # temp_client = False + room_id = cvars.config.get().telegram.room_id + message = format_message(alert) + await client.send_message(entity=room_id, message=message) + # if temp_client: + # await client.close() + + +async def send_start_alert() -> None: + config = cvars.config.get() + await send_alert( + Alert( + alert_type=AlertType.BOOT, + message=f"Service running with enabled checkers: {', '.join(config.enabled_checker_sets)}", + severity=Severity.INFO, + ) + ) + + +async def send_stop_alert() -> None: + await send_alert( + Alert( + alert_type=AlertType.BOOT, + message="Service stopping.", + severity=Severity.INFO, + ) + ) diff --git a/src/lego_monitoring/alerting/enum.py b/src/lego_monitoring/alerting/enum.py new file mode 100644 index 0000000..1cb6a08 --- /dev/null +++ b/src/lego_monitoring/alerting/enum.py @@ -0,0 +1,23 @@ +from enum import StrEnum + + +class AlertType(StrEnum): + BOOT = "BOOT" + TEST = "TEST" + # ERROR = "ERROR" + # RAM = "RAM" + # CPU = "CPU" + # TEMP = "TEMP" + # VULN = "VULN" + # LOGIN = "LOGIN" + # SMART = "SMART" # TODO + # RAID = "RAID" + # DISKS = "DISKS" + # UPS = "UPS" + # UPDATE = "UPDATE" + + +class Severity(StrEnum): + INFO = "INFO" + WARNING = "WARNING" + CRITICAL = "CRITICAL" diff --git a/src/lego_monitoring/core/checkers.py b/src/lego_monitoring/core/checkers.py new file mode 100644 index 0000000..12db2b8 --- /dev/null +++ b/src/lego_monitoring/core/checkers.py @@ -0,0 +1,53 @@ +import asyncio +import datetime +import logging +from typing import Callable, Coroutine + +from ..alerting import alerts + + +async def _call_check(check: Callable | Coroutine, *args, **kwargs) -> list[alerts.Alert]: + 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") + return result + + +async def interval_checker(check: Callable | Coroutine, interval: datetime.timedelta, *args, **kwargs): + interval_secs = interval.total_seconds() + while True: + logging.info(f"Calling {check.__name__}") + result = await _call_check(check, *args, **kwargs) + logging.info(f"Got {len(result)} alerts") + for alert in result: + await alerts.send_alert(alert) + await asyncio.sleep(interval_secs) + + +async def scheduled_checker( + check: Callable | Coroutine, period: datetime.timedelta, when: datetime.time, *args, **kwargs +): + match period: + case datetime.timedelta(days=1): + while True: + now = datetime.datetime.now() + next_datetime = datetime.datetime.combine(datetime.date.today(), when) + if next_datetime < now: + next_datetime += datetime.timedelta(days=1) + logging.info(f"Scheduled to call {check.__name__} at {next_datetime.isoformat()}") + await asyncio.sleep( + (next_datetime - now).total_seconds() + ) # might be negative at this point, asyncio doesn't care + + logging.info(f"Calling {check.__name__}") + result = await _call_check(check, *args, **kwargs) + logging.info(f"Got {len(result)} alerts") + for alert in result: + await alerts.send_alert(alert) + case _: + raise NotImplementedError diff --git a/src/lego_monitoring/core/config.py b/src/lego_monitoring/core/config.py index 25566e8..71cc766 100644 --- a/src/lego_monitoring/core/config.py +++ b/src/lego_monitoring/core/config.py @@ -1,32 +1,43 @@ import json from dataclasses import dataclass -from typing import Optional from alt_utils import NestedDeserializableDataclass +@dataclass +class TelegramConfig: + creds: str + room_id: int + + @dataclass class Config(NestedDeserializableDataclass): - non_secret_config_option: str - config_option: Optional[str] + enabled_checker_sets: list[str] + telegram: TelegramConfig + def load_config(filepath: str) -> Config: + def load_secrets(d: dict) -> dict: + new_d = {} + for k in d: + if k.endswith("_secret_path"): + actual_opt_key = k[:-12] + secret_path = d[k] + if secret_path is None: + new_d[actual_opt_key] = None + else: + with open(secret_path) as sf: + new_d[actual_opt_key] = sf.read().rstrip() + elif type(d[k]) == dict: + new_d[k] = load_secrets(d[k]) + else: + new_d[k] = d[k] + return new_d + with open(filepath) as f: cfg_dict = json.load(f) - # load secrets from paths - new_cfg_dict = {} - for k in cfg_dict: - if k.endswith('_secret_path'): - actual_opt_key = k[:-12] - secret_path = cfg_dict[k] - if secret_path is None: - new_cfg_dict[actual_opt_key] = None - else: - with open(secret_path) as sf: - new_cfg_dict[actual_opt_key] = sf.read().rstrip() - else: - new_cfg_dict[k] = cfg_dict[k] + new_cfg_dict = load_secrets(cfg_dict) cfg = Config.from_dict(new_cfg_dict) return cfg diff --git a/src/lego_monitoring/core/cvars.py b/src/lego_monitoring/core/cvars.py new file mode 100644 index 0000000..b1f7b50 --- /dev/null +++ b/src/lego_monitoring/core/cvars.py @@ -0,0 +1,8 @@ +from contextvars import ContextVar + +from telethon import TelegramClient + +from .config import Config + +config: ContextVar[Config] = ContextVar("config") +tg_client: ContextVar[TelegramClient] = ContextVar("tg_client") diff --git a/uv.lock b/uv.lock index 7934399..41ce1c9 100644 --- a/uv.lock +++ b/uv.lock @@ -16,10 +16,62 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "alt-utils" }, + { name = "setuptools" }, + { name = "telethon" }, ] [package.metadata] -requires-dist = [{ name = "alt-utils", specifier = ">=0.0.6" }] +requires-dist = [ + { name = "alt-utils", specifier = ">=0.0.6" }, + { name = "setuptools", specifier = ">=80.0.0" }, + { name = "telethon", specifier = ">=1.40.0" }, +] -[package.metadata.requires-dev] -dev = [] +[[package]] +name = "pyaes" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/66/2c17bae31c906613795711fc78045c285048168919ace2220daa372c7d72/pyaes-1.6.1.tar.gz", hash = "sha256:02c1b1405c38d3c370b085fb952dd8bea3fadcee6411ad99f312cc129c536d8f", size = 28536 } + +[[package]] +name = "pyasn1" +version = "0.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/e9/01f1a64245b89f039897cb0130016d79f77d52669aae6ee7b159a6c4c018/pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034", size = 145322 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/f1/d6a797abb14f6283c0ddff96bbdd46937f64122b8c925cab503dd37f8214/pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629", size = 83135 }, +] + +[[package]] +name = "rsa" +version = "4.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyasn1" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/da/8a/22b7beea3ee0d44b1916c0c1cb0ee3af23b700b6da9f04991899d0c555d4/rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75", size = 29034 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696 }, +] + +[[package]] +name = "setuptools" +version = "80.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/80/97e25f0f1e4067677806084b7382a6ff9979f3d15119375c475c288db9d7/setuptools-80.0.0.tar.gz", hash = "sha256:c40a5b3729d58dd749c0f08f1a07d134fb8a0a3d7f87dc33e7c5e1f762138650", size = 1354221 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/63/5517029d6696ddf2bd378d46f63f479be001c31b462303170a1da57650cb/setuptools-80.0.0-py3-none-any.whl", hash = "sha256:a38f898dcd6e5380f4da4381a87ec90bd0a7eec23d204a5552e80ee3cab6bd27", size = 1240907 }, +] + +[[package]] +name = "telethon" +version = "1.40.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyaes" }, + { name = "rsa" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/af/9b7111e3f63fffe8e55b7ceb8bda023173e2052f420b6debcb25fd2fbc15/telethon-1.40.0.tar.gz", hash = "sha256:40e83326877a2e68b754d4b6d0d1ca5ac924110045b039e02660f2d67add97db", size = 646723 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/5a/c5370edb3215d19a6e858f4169b8eec725ba55f9d39df0f557508048c037/Telethon-1.40.0-py3-none-any.whl", hash = "sha256:146fd4cb2a7afa66bc67f9c2167756096a37b930f65711a3e7399ec9874dcfa7", size = 722013 }, +]