maubot/maubot/management/api/log.py

171 lines
4.9 KiB
Python
Raw Normal View History

2018-11-11 17:27:09 +00:00
# maubot - A plugin-based Matrix bot system.
2022-03-25 12:22:37 +00:00
# Copyright (C) 2022 Tulir Asokan
2018-11-11 17:27:09 +00:00
#
# 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 <https://www.gnu.org/licenses/>.
2022-03-25 12:22:37 +00:00
from __future__ import annotations
from collections import deque
2022-03-25 12:22:37 +00:00
from datetime import datetime
2018-11-11 17:27:09 +00:00
import asyncio
2022-03-25 12:22:37 +00:00
import logging
2018-11-11 17:27:09 +00:00
2022-03-25 12:22:37 +00:00
from aiohttp import web, web_ws
2018-11-11 17:27:09 +00:00
from .auth import is_valid_token
from .base import routes
2022-03-25 12:22:37 +00:00
BUILTIN_ATTRS = {
"args",
"asctime",
"created",
"exc_info",
"exc_text",
"filename",
"funcName",
"levelname",
"levelno",
"lineno",
"module",
"msecs",
"message",
"msg",
"name",
"pathname",
"process",
"processName",
"relativeCreated",
"stack_info",
"thread",
"threadName",
}
INCLUDE_ATTRS = {
"filename",
"funcName",
"levelname",
"levelno",
"lineno",
"module",
"name",
"pathname",
}
2018-11-11 17:27:09 +00:00
EXCLUDE_ATTRS = BUILTIN_ATTRS - INCLUDE_ATTRS
MAX_LINES = 2048
2018-11-11 17:27:09 +00:00
class LogCollector(logging.Handler):
2022-03-25 12:22:37 +00:00
lines: deque[dict]
formatter: logging.Formatter
2022-03-25 12:22:37 +00:00
listeners: list[web.WebSocketResponse]
2020-04-05 13:39:33 +00:00
loop: asyncio.AbstractEventLoop
def __init__(self, level=logging.NOTSET) -> None:
2018-11-11 17:27:09 +00:00
super().__init__(level)
self.lines = deque(maxlen=MAX_LINES)
2018-11-11 17:27:09 +00:00
self.formatter = logging.Formatter()
self.listeners = []
2018-11-11 17:27:09 +00:00
def emit(self, record: logging.LogRecord) -> None:
2018-11-11 22:00:37 +00:00
try:
self._emit(record)
except Exception as e:
print("Logging error:", e)
def _emit(self, record: logging.LogRecord) -> None:
2018-11-11 17:27:09 +00:00
# JSON conversion based on Marsel Mavletkulov's json-log-formatter (MIT license)
# https://github.com/marselester/json-log-formatter
content = {
2022-03-25 12:22:37 +00:00
name: value for name, value in record.__dict__.items() if name not in EXCLUDE_ATTRS
2018-11-11 17:27:09 +00:00
}
content["id"] = str(record.relativeCreated)
2018-11-11 17:27:09 +00:00
content["msg"] = record.getMessage()
content["time"] = datetime.fromtimestamp(record.created)
2018-11-11 17:27:09 +00:00
if record.exc_info:
content["exc_info"] = self.formatter.formatException(record.exc_info)
for name, value in content.items():
if isinstance(value, datetime):
content[name] = value.astimezone().isoformat()
2020-04-05 13:39:33 +00:00
asyncio.run_coroutine_threadsafe(self.send(content), loop=self.loop)
self.lines.append(content)
2018-11-11 17:27:09 +00:00
async def send(self, record: dict) -> None:
for ws in self.listeners:
try:
await ws.send_json(record)
except Exception as e:
print("Log sending error:", e)
2018-11-11 17:27:09 +00:00
handler = LogCollector()
2018-11-11 17:27:09 +00:00
log = logging.getLogger("maubot.server.websocket")
sockets = []
2020-04-05 13:39:33 +00:00
def init(loop: asyncio.AbstractEventLoop) -> None:
logging.root.addHandler(handler)
2020-04-05 13:39:33 +00:00
handler.loop = loop
2018-11-11 17:27:09 +00:00
async def stop_all() -> None:
2021-11-19 13:22:54 +00:00
log.debug("Closing log listener websockets")
logging.root.removeHandler(handler)
2018-11-11 17:27:09 +00:00
for socket in sockets:
try:
await socket.close(code=1012)
2018-11-11 17:27:09 +00:00
except Exception:
pass
@routes.get("/logs")
async def log_websocket(request: web.Request) -> web.WebSocketResponse:
ws = web.WebSocketResponse()
await ws.prepare(request)
sockets.append(ws)
log.debug(f"Connection from {request.remote} opened")
authenticated = False
async def close_if_not_authenticated():
await asyncio.sleep(5)
2018-11-11 17:27:09 +00:00
if not authenticated:
await ws.close(code=4000)
log.debug(f"Connection from {request.remote} terminated due to no authentication")
asyncio.create_task(close_if_not_authenticated())
2018-11-11 17:27:09 +00:00
try:
2022-03-25 12:22:37 +00:00
msg: web_ws.WSMessage
2018-11-11 17:27:09 +00:00
async for msg in ws:
if msg.type != web.WSMsgType.TEXT:
continue
if is_valid_token(msg.data):
await ws.send_json({"auth_success": True})
await ws.send_json({"history": list(handler.lines)})
2018-11-11 17:27:09 +00:00
if not authenticated:
log.debug(f"Connection from {request.remote} authenticated")
handler.listeners.append(ws)
2018-11-11 17:27:09 +00:00
authenticated = True
elif not authenticated:
await ws.send_json({"auth_success": False})
except Exception:
try:
await ws.close()
except Exception:
pass
if authenticated:
handler.listeners.remove(ws)
2018-11-11 17:27:09 +00:00
log.debug(f"Connection from {request.remote} closed")
sockets.remove(ws)
return ws