creds file -> config file

This commit is contained in:
Alex 2024-10-27 14:47:02 +03:00
parent c653055134
commit f199292aba
5 changed files with 70 additions and 49 deletions

View file

@ -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,
)

View file

@ -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()

View file

@ -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}")
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())

6
misc/cvars.py Normal file
View file

@ -0,0 +1,6 @@
from contextvars import ContextVar
import nio
config: ContextVar[dict] = ContextVar("config")
matrix_client: ContextVar[nio.AsyncClient] = ContextVar("matrix_client")

View file

@ -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()