diff --git a/alerting/alerts.py b/alerting/alerts.py index e403c5a..8827c2a 100644 --- a/alerting/alerts.py +++ b/alerting/alerts.py @@ -6,7 +6,8 @@ from typing import Optional import aiofiles import nio -from alerting.common import CREDS_FILE, ROOM_ID +from alerting.common import CONFIG_FILE +from misc import cvars class AlertType(StrEnum): @@ -41,13 +42,11 @@ async def get_client() -> nio.AsyncClient: Returns a Matrix client. It is better to call get_client once and use it for multiple send_alert calls """ - async with aiofiles.open(CREDS_FILE) as f: - contents = await f.read() - creds = json.loads(contents) - client = nio.AsyncClient(creds["homeserver"]) - client.access_token = creds["access_token"] - client.user_id = creds["user_id"] - client.device_id = creds["device_id"] + matrix_cfg = cvars.config.get()["matrix"] + client = nio.AsyncClient(matrix_cfg["homeserver"]) + client.access_token = matrix_cfg["access_token"] + client.user_id = matrix_cfg["user_id"] + client.device_id = matrix_cfg["device_id"] return client @@ -67,12 +66,19 @@ def format_message(alert: Alert) -> str: return message, None -async def send_alert(alert: Alert, client: Optional[nio.AsyncClient] = None) -> None: - if client is None: +async def send_alert(alert: Alert) -> None: + try: + client = cvars.matrix_client.get() + except LookupError: # being called standalone + async with aiofiles.open(CONFIG_FILE) as f: + contents = await f.read() + cvars.config.set(json.loads(contents)) temp_client = True client = await get_client() + cvars.matrix_client.set(client) else: temp_client = False + room_id = cvars.config.get()["matrix"]["room_id"] message, html_message = format_message(alert) content = { "msgtype": "m.text", @@ -82,7 +88,7 @@ async def send_alert(alert: Alert, client: Optional[nio.AsyncClient] = None) -> content["format"] = "org.matrix.custom.html" content["formatted_body"] = html_message await client.room_send( - room_id=ROOM_ID, + room_id=room_id, message_type="m.room.message", content=content, ) diff --git a/alerting/common.py b/alerting/common.py index 443358b..a9ed133 100644 --- a/alerting/common.py +++ b/alerting/common.py @@ -1,8 +1,4 @@ import os from pathlib import Path -CREDS_FILE = (Path(os.path.dirname(os.path.realpath(__file__))) / "credentials.json").resolve() -HOMESERVER = "https://matrix.altau.su" -USER_ID = "@alertbot:altau.su" -DEVICE_NAME = "lego" -ROOM_ID = "!lheAkNnuTYwITzBxiW:altau.su" +CONFIG_FILE = (Path(os.path.dirname(os.path.realpath(__file__))) / "config.json").resolve() diff --git a/alerting/login.py b/alerting/login.py index bd20dd2..2b36c0e 100755 --- a/alerting/login.py +++ b/alerting/login.py @@ -1,41 +1,45 @@ #!/usr/bin/env python3 -import asyncio import getpass import json import os -import stat -from common import CREDS_FILE, DEVICE_NAME, HOMESERVER, USER_ID +from common import CONFIG_FILE from nio import AsyncClient, LoginResponse async def main() -> None: - if os.path.exists(CREDS_FILE): - print(f"Creds already configured in {CREDS_FILE}") - raise SystemExit + try: + with open(CONFIG_FILE) as f: + cfg = json.load(f) + if "matrix" in cfg: + print(f"Creds already configured in {CONFIG_FILE}") + raise SystemExit + cfg["matrix"] = {} + except (FileNotFoundError, json.JSONDecodeError): + open(CONFIG_FILE, "w").close() + os.chmod(CONFIG_FILE, 0o600) + cfg = {"matrix": {}} - client = AsyncClient(HOMESERVER, USER_ID) + homeserver = input("Homeserver: ") + user_id = input("User ID: ") + device_name = input("Device name: ") + room_id = input("Room ID: ") password = getpass.getpass() - resp = await client.login(password, device_name=DEVICE_NAME) + + client = AsyncClient(homeserver, user_id) + resp = await client.login(password, device_name=device_name) await client.close() + if isinstance(resp, LoginResponse): - open(CREDS_FILE, "w").close() - os.chmod(CREDS_FILE, stat.S_IRUSR | stat.S_IWUSR) - with open(CREDS_FILE, "w") as f: - json.dump( - { - "homeserver": HOMESERVER, - "user_id": resp.user_id, - "device_id": resp.device_id, - "access_token": resp.access_token, - }, - f, - ) - print(f"Logged in as {resp.user_id}. Credentials saved to {CREDS_FILE}") + cfg["matrix"]["homeserver"] = homeserver + cfg["matrix"]["user_id"] = resp.user_id + cfg["matrix"]["device_id"] = resp.device_id + cfg["matrix"]["access_token"] = resp.access_token + cfg["matrix"]["room_id"] = room_id + + with open(CONFIG_FILE, "w") as f: + json.dump(cfg, f, indent=2) + print(f"Logged in as {resp.user_id}. Credentials saved to {CONFIG_FILE}") else: raise Exception(f"Failed to log in: {resp}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/misc/cvars.py b/misc/cvars.py new file mode 100644 index 0000000..b4cfff2 --- /dev/null +++ b/misc/cvars.py @@ -0,0 +1,6 @@ +from contextvars import ContextVar + +import nio + +config: ContextVar[dict] = ContextVar("config") +matrix_client: ContextVar[nio.AsyncClient] = ContextVar("matrix_client") diff --git a/service.py b/service.py index fb63030..60572ae 100755 --- a/service.py +++ b/service.py @@ -1,13 +1,16 @@ #!/usr/bin/env python3 import asyncio +import json import logging import signal from typing import Callable, Coroutine +import aiofiles import nio from alerting import alerts -from misc import checks +from alerting.common import CONFIG_FILE +from misc import checks, cvars logging.basicConfig(level=logging.INFO) @@ -19,7 +22,7 @@ def stop_gracefully(signum, frame): stopping = True -async def checker(check: Callable | Coroutine, interval_secs: int, client: nio.AsyncClient, *args, **kwargs): +async def checker(check: Callable | Coroutine, interval_secs: int, *args, **kwargs): while True: logging.info(f"Calling {check.__name__}") if isinstance(check, Callable): @@ -32,7 +35,7 @@ async def checker(check: Callable | Coroutine, interval_secs: int, client: nio.A 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 alerts.send_alert(alert) await asyncio.sleep(interval_secs) @@ -43,13 +46,19 @@ async def main(): WEEK = 7 * DAY signal.signal(signal.SIGTERM, stop_gracefully) + + async with aiofiles.open(CONFIG_FILE) as f: + contents = await f.read() + cvars.config.set(json.loads(contents)) + client = await alerts.get_client() + cvars.matrix_client.set(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), + checker(checks.temp_check, 5 * MINUTE), + checker(checks.cpu_check, 5 * MINUTE), + checker(checks.ups_check, 5 * MINUTE), + checker(checks.ram_check, 1 * MINUTE), + checker(checks.vuln_check, 1 * DAY), ) async with asyncio.TaskGroup() as tg: checker_tasks: set[asyncio.Task] = set()