From 6f06eec3cc509cc32a247f688fbf8e47777d2adb Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 14 Dec 2018 00:33:33 +0200 Subject: [PATCH] Add log viewer to mbc --- maubot/cli/commands/__init__.py | 2 +- maubot/cli/commands/build.py | 9 +-- maubot/cli/commands/logs.py | 110 ++++++++++++++++++++++++++++++++ maubot/cli/commands/upload.py | 15 ++--- maubot/cli/config.py | 22 +++++++ 5 files changed, 141 insertions(+), 17 deletions(-) create mode 100644 maubot/cli/commands/logs.py diff --git a/maubot/cli/commands/__init__.py b/maubot/cli/commands/__init__.py index e94091c..4eba273 100644 --- a/maubot/cli/commands/__init__.py +++ b/maubot/cli/commands/__init__.py @@ -1 +1 @@ -from . import upload, build, login, init +from . import upload, build, login, init, logs diff --git a/maubot/cli/commands/build.py b/maubot/cli/commands/build.py index 3e90f6a..0f3201c 100644 --- a/maubot/cli/commands/build.py +++ b/maubot/cli/commands/build.py @@ -27,7 +27,7 @@ import click from ...loader import PluginMeta from ..cliq.validators import PathValidator from ..base import app -from ..config import config +from ..config import get_default_server from .upload import upload_file yaml = YAML() @@ -98,11 +98,8 @@ def write_plugin(meta: PluginMeta, output: Union[str, IO]) -> None: def upload_plugin(output: Union[str, IO]) -> None: - try: - server = config["default_server"] - token = config["servers"][server] - except KeyError: - print(Fore.RED + "Default server not configured." + Fore.RESET) + server, token = get_default_server() + if not token: return if isinstance(output, str): with open(output, "rb") as file: diff --git a/maubot/cli/commands/logs.py b/maubot/cli/commands/logs.py new file mode 100644 index 0000000..3705d47 --- /dev/null +++ b/maubot/cli/commands/logs.py @@ -0,0 +1,110 @@ +# maubot - A plugin-based Matrix bot system. +# Copyright (C) 2018 Tulir Asokan +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +from datetime import datetime +import asyncio + +from colorama import Fore +from aiohttp import WSMsgType, WSMessage, ClientSession +from mautrix.client.api.types.util import Obj +import click + +from ..config import get_token, get_default_server +from ..base import app + +history_count: int = 10 + + +@app.command(help="View the logs of a server") +@click.argument("server", required=False) +@click.option("-t", "--tail", default=10, help="Maximum number of old log lines to display") +def logs(server: str, tail: int) -> None: + if not server: + server, token = get_default_server() + else: + token = get_token(server) + if not token: + return + global history_count + history_count = tail + loop = asyncio.get_event_loop() + future = asyncio.ensure_future(view_logs(server, token), loop=loop) + try: + loop.run_until_complete(future) + except KeyboardInterrupt: + future.cancel() + loop.run_until_complete(future) + loop.close() + + +def parsedate(entry: Obj) -> None: + i = entry.time.index("+") + i = entry.time.index(":", i) + entry.time = entry.time[:i] + entry.time[i + 1:] + entry.time = datetime.strptime(entry.time, "%Y-%m-%dT%H:%M:%S.%f%z") + + +levelcolors = { + "DEBUG": "", + "INFO": Fore.CYAN, + "WARNING": Fore.YELLOW, + "ERROR": Fore.RED, + "FATAL": Fore.MAGENTA, +} + + +def print_entry(entry: dict) -> None: + entry = Obj(**entry) + parsedate(entry) + print("{levelcolor}[{date}] [{level}@{logger}] {message}{resetcolor}" + .format(date=entry.time.strftime("%Y-%m-%d %H:%M:%S"), + level=entry.levelname, + levelcolor=levelcolors.get(entry.levelname, ""), + resetcolor=Fore.RESET, + logger=entry.name, + message=entry.msg)) + + +def handle_msg(data: dict) -> bool: + if "auth_success" in data: + if data["auth_success"]: + print(Fore.GREEN + "Connected to log websocket" + Fore.RESET) + else: + print(Fore.RED + "Failed to authenticate to log websocket" + Fore.RESET) + return False + elif "history" in data: + for entry in data["history"][-history_count:]: + print_entry(entry) + else: + print_entry(data) + return True + + +async def view_logs(server: str, token: str) -> None: + async with ClientSession() as session: + async with session.ws_connect(f"{server}/_matrix/maubot/v1/logs") as ws: + await ws.send_str(token) + try: + msg: WSMessage + async for msg in ws: + if msg.type == WSMsgType.TEXT: + if not handle_msg(msg.json()): + break + elif msg.type == WSMsgType.ERROR: + print(Fore.YELLOW + "Connection error: " + msg.data + Fore.RESET) + elif msg.type == WSMsgType.CLOSE: + print(Fore.YELLOW + "Server closed connection" + Fore.RESET) + except asyncio.CancelledError: + pass diff --git a/maubot/cli/commands/upload.py b/maubot/cli/commands/upload.py index 8da3558..428c1a4 100644 --- a/maubot/cli/commands/upload.py +++ b/maubot/cli/commands/upload.py @@ -22,7 +22,7 @@ from colorama import Fore import click from ..base import app -from ..config import config +from ..config import get_default_server, get_token class UploadError(Exception): @@ -34,15 +34,10 @@ class UploadError(Exception): @click.option("-s", "--server", help="The maubot instance to upload the plugin to") def upload(path: str, server: str) -> None: if not server: - try: - server = config["default_server"] - except KeyError: - print(Fore.RED + "Default server not configured" + Fore.RESET) - return - try: - token = config["servers"][server] - except KeyError: - print(Fore.RED + "Server not found" + Fore.RESET) + server, token = get_default_server() + else: + token = get_token(server) + if not token: return with open(path, "rb") as file: upload_file(file, server, token) diff --git a/maubot/cli/config.py b/maubot/cli/config.py index 653ffef..70336f8 100644 --- a/maubot/cli/config.py +++ b/maubot/cli/config.py @@ -13,9 +13,12 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . +from typing import Tuple, Optional import json import os +from colorama import Fore + config = { "servers": {}, "default_server": None, @@ -23,6 +26,25 @@ config = { configdir = os.environ.get("XDG_CONFIG_HOME", os.path.join(os.environ.get("HOME"), ".config")) +def get_default_server() -> Tuple[Optional[str], Optional[str]]: + try: + server: str = config["default_server"] + except KeyError: + server = None + if server is None: + print(f"{Fore.RED}Default server not configured.{Fore.RESET}") + return None, None + return server, get_token(server) + + +def get_token(server: str) -> Optional[str]: + try: + return config["servers"][server] + except KeyError: + print(f"{Fore.RED}No access token saved for {server}.{Fore.RESET}") + return None + + def save_config() -> None: with open(f"{configdir}/maubot-cli.json", "w") as file: json.dump(config, file)