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 aiofiles
import nio import nio
from alerting.common import CREDS_FILE, ROOM_ID from alerting.common import CONFIG_FILE
from misc import cvars
class AlertType(StrEnum): class AlertType(StrEnum):
@ -41,13 +42,11 @@ async def get_client() -> nio.AsyncClient:
Returns a Matrix client. Returns a Matrix client.
It is better to call get_client once and use it for multiple send_alert calls It is better to call get_client once and use it for multiple send_alert calls
""" """
async with aiofiles.open(CREDS_FILE) as f: matrix_cfg = cvars.config.get()["matrix"]
contents = await f.read() client = nio.AsyncClient(matrix_cfg["homeserver"])
creds = json.loads(contents) client.access_token = matrix_cfg["access_token"]
client = nio.AsyncClient(creds["homeserver"]) client.user_id = matrix_cfg["user_id"]
client.access_token = creds["access_token"] client.device_id = matrix_cfg["device_id"]
client.user_id = creds["user_id"]
client.device_id = creds["device_id"]
return client return client
@ -67,12 +66,19 @@ def format_message(alert: Alert) -> str:
return message, None return message, None
async def send_alert(alert: Alert, client: Optional[nio.AsyncClient] = None) -> None: async def send_alert(alert: Alert) -> None:
if client is 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 temp_client = True
client = await get_client() client = await get_client()
cvars.matrix_client.set(client)
else: else:
temp_client = False temp_client = False
room_id = cvars.config.get()["matrix"]["room_id"]
message, html_message = format_message(alert) message, html_message = format_message(alert)
content = { content = {
"msgtype": "m.text", "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["format"] = "org.matrix.custom.html"
content["formatted_body"] = html_message content["formatted_body"] = html_message
await client.room_send( await client.room_send(
room_id=ROOM_ID, room_id=room_id,
message_type="m.room.message", message_type="m.room.message",
content=content, content=content,
) )

View file

@ -1,8 +1,4 @@
import os import os
from pathlib import Path from pathlib import Path
CREDS_FILE = (Path(os.path.dirname(os.path.realpath(__file__))) / "credentials.json").resolve() CONFIG_FILE = (Path(os.path.dirname(os.path.realpath(__file__))) / "config.json").resolve()
HOMESERVER = "https://matrix.altau.su"
USER_ID = "@alertbot:altau.su"
DEVICE_NAME = "lego"
ROOM_ID = "!lheAkNnuTYwITzBxiW:altau.su"

View file

@ -1,41 +1,45 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
import asyncio
import getpass import getpass
import json import json
import os import os
import stat
from common import CREDS_FILE, DEVICE_NAME, HOMESERVER, USER_ID from common import CONFIG_FILE
from nio import AsyncClient, LoginResponse from nio import AsyncClient, LoginResponse
async def main() -> None: async def main() -> None:
if os.path.exists(CREDS_FILE): try:
print(f"Creds already configured in {CREDS_FILE}") with open(CONFIG_FILE) as f:
raise SystemExit 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() 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() await client.close()
if isinstance(resp, LoginResponse): if isinstance(resp, LoginResponse):
open(CREDS_FILE, "w").close() cfg["matrix"]["homeserver"] = homeserver
os.chmod(CREDS_FILE, stat.S_IRUSR | stat.S_IWUSR) cfg["matrix"]["user_id"] = resp.user_id
with open(CREDS_FILE, "w") as f: cfg["matrix"]["device_id"] = resp.device_id
json.dump( cfg["matrix"]["access_token"] = resp.access_token
{ cfg["matrix"]["room_id"] = room_id
"homeserver": HOMESERVER,
"user_id": resp.user_id, with open(CONFIG_FILE, "w") as f:
"device_id": resp.device_id, json.dump(cfg, f, indent=2)
"access_token": resp.access_token, print(f"Logged in as {resp.user_id}. Credentials saved to {CONFIG_FILE}")
},
f,
)
print(f"Logged in as {resp.user_id}. Credentials saved to {CREDS_FILE}")
else: else:
raise Exception(f"Failed to log in: {resp}") 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 #!/usr/bin/env python3
import asyncio import asyncio
import json
import logging import logging
import signal import signal
from typing import Callable, Coroutine from typing import Callable, Coroutine
import aiofiles
import nio import nio
from alerting import alerts from alerting import alerts
from misc import checks from alerting.common import CONFIG_FILE
from misc import checks, cvars
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
@ -19,7 +22,7 @@ def stop_gracefully(signum, frame):
stopping = True 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: while True:
logging.info(f"Calling {check.__name__}") logging.info(f"Calling {check.__name__}")
if isinstance(check, Callable): 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") raise TypeError(f"check is {type(check)}, neither function nor coroutine")
logging.info(f"Got {len(result)} alerts") logging.info(f"Got {len(result)} alerts")
for alert in result: for alert in result:
await alerts.send_alert(alert, client) await alerts.send_alert(alert)
await asyncio.sleep(interval_secs) await asyncio.sleep(interval_secs)
@ -43,13 +46,19 @@ async def main():
WEEK = 7 * DAY WEEK = 7 * DAY
signal.signal(signal.SIGTERM, stop_gracefully) 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() client = await alerts.get_client()
cvars.matrix_client.set(client)
checkers = ( checkers = (
checker(checks.temp_check, 5 * MINUTE, client), checker(checks.temp_check, 5 * MINUTE),
checker(checks.cpu_check, 5 * MINUTE, client), checker(checks.cpu_check, 5 * MINUTE),
checker(checks.ups_check, 5 * MINUTE, client), checker(checks.ups_check, 5 * MINUTE),
checker(checks.ram_check, 1 * MINUTE, client), checker(checks.ram_check, 1 * MINUTE),
checker(checks.vuln_check, 1 * DAY, client), checker(checks.vuln_check, 1 * DAY),
) )
async with asyncio.TaskGroup() as tg: async with asyncio.TaskGroup() as tg:
checker_tasks: set[asyncio.Task] = set() checker_tasks: set[asyncio.Task] = set()