From 88861f3880ea75b73058a79fb0c78016639b62cc Mon Sep 17 00:00:00 2001
From: Conatum <cona@thewisewolf.dev>
Date: Thu, 3 Nov 2022 15:35:30 +0200
Subject: [PATCH] rewrite: New live-logger.

---
 bot/meta/logger.py | 135 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 133 insertions(+), 2 deletions(-)

diff --git a/bot/meta/logger.py b/bot/meta/logger.py
index b4ef3402..8a1bf653 100644
--- a/bot/meta/logger.py
+++ b/bot/meta/logger.py
@@ -1,15 +1,37 @@
 import sys
 import logging
 import asyncio
+from logging.handlers import QueueListener, QueueHandler
+from queue import SimpleQueue
+from contextlib import contextmanager
+
 from contextvars import ContextVar
-from discord import AllowedMentions
+from discord import AllowedMentions, Webhook
+import aiohttp
 
 from .config import conf
 from . import sharding
+from utils.lib import split_text, utc_now
 
 
 log_context: ContextVar[str] = ContextVar('logging_context', default='CTX: ROOT CONTEXT')
 log_action: ContextVar[str] = ContextVar('logging_action', default='UNKNOWN ACTION')
+log_app: ContextVar[str] = ContextVar('logging_shard', default="SHARD {:03}".format(sharding.shard_number))
+
+
+@contextmanager
+def logging_context(context=None, action=None):
+    if context is not None:
+        context_t = log_context.set(context)
+    if action is not None:
+        action_t = log_action.set(action)
+    try:
+        yield
+    finally:
+        if context is not None:
+            log_context.reset(context_t)
+        if action is not None:
+            log_action.reset(action_t)
 
 
 RESET_SEQ = "\033[0m"
@@ -37,7 +59,7 @@ def colour_escape(fmt: str) -> str:
 
 
 log_format = ('[%(green)%(asctime)-19s%(reset)][%(red)%(levelname)-8s%(reset)]' +
-              '[%(cyan)SHARD {:02}%(reset)]'.format(sharding.shard_number) +
+              '[%(cyan)%(app)-15s%(reset)]' +
               '[%(cyan)%(context)-22s%(reset)]' +
               '[%(cyan)%(action)-22s%(reset)]' +
               ' %(bold)%(cyan)%(name)s:%(reset)' +
@@ -70,6 +92,7 @@ class ContextInjection(logging.Filter):
             record.context = log_context.get()
         if not hasattr(record, 'action'):
             record.action = log_action.get()
+        record.app = log_app.get()
         return True
 
 
@@ -86,6 +109,114 @@ logging_handler_err.setFormatter(log_fmt)
 logging_handler_err.addFilter(ContextInjection())
 logger.addHandler(logging_handler_err)
 
+
+class LocalQueueHandler(QueueHandler):
+    def emit(self, record: logging.LogRecord) -> None:
+        # Removed the call to self.prepare(), handle task cancellation
+        try:
+            self.enqueue(record)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            self.handleError(record)
+
+
+class WebHookHandler(logging.StreamHandler):
+    def __init__(self, webhook_url, batch=False):
+        super().__init__(self)
+        self.webhook_url = webhook_url
+        self.batched = ""
+        self.batch = batch
+        self.loop = None
+
+    def get_loop(self):
+        if self.loop is None:
+            self.loop = asyncio.new_event_loop()
+        return self.loop
+
+    def emit(self, record):
+        self.get_loop().run_until_complete(self.post(record))
+
+    async def post(self, record):
+        try:
+            timestamp = utc_now().strftime("%d/%m/%Y, %H:%M:%S")
+            header = f"[{record.levelname}][{record.app}][{record.context}][{record.action}][{timestamp}]"
+            message = record.msg
+
+            # TODO: Maybe send file instead of splitting?
+            # TODO: Reformat header a little
+            if len(message) > 1900:
+                blocks = split_text(message, blocksize=1900, code=False)
+            else:
+                blocks = [message]
+
+            if len(blocks) > 1:
+                blocks = [
+                    "```md\n{}[{}/{}]\n{}\n```".format(header, i+1, len(blocks), block) for i, block in enumerate(blocks)
+                ]
+            else:
+                blocks = ["```md\n{}\n{}\n```".format(header, blocks[0])]
+
+            # Post the log message(s)
+            if self.batch:
+                if len(message) > 500:
+                    await self._send_batched()
+                    await self._send(*blocks)
+                elif len(self.batched) + len(blocks[0]) > 500:
+                    self.batched += blocks[0]
+                    await self._send_batched()
+                else:
+                    self.batched += blocks[0]
+            else:
+                await self._send(*blocks)
+        except Exception as ex:
+            print(ex)
+
+    async def _send_batched(self):
+        if self.batched:
+            batched = self.batched
+            self.batched = ""
+            await self._send(batched)
+
+    async def _send(self, *blocks):
+        async with aiohttp.ClientSession() as session:
+            webhook = Webhook.from_url(self.webhook_url, session=session)
+            for block in blocks:
+                await webhook.send(block)
+
+
+handlers = []
+if webhook := conf.logging['general_log']:
+    handler = WebHookHandler(webhook, batch=True)
+    handlers.append(handler)
+
+if webhook := conf.logging['error_log']:
+    handler = WebHookHandler(webhook, batch=False)
+    handler.setLevel(logging.ERROR)
+    handlers.append(handler)
+
+if webhook := conf.logging['critical_log']:
+    handler = WebHookHandler(webhook, batch=False)
+    handler.setLevel(logging.CRITICAL)
+    handlers.append(handler)
+
+if handlers:
+    queue: SimpleQueue[logging.LogRecord] = SimpleQueue()
+
+    handler = QueueHandler(queue)
+    handler.setLevel(logging.INFO)
+    handler.addFilter(ContextInjection())
+    logger.addHandler(handler)
+
+    listener = QueueListener(
+        queue, *handlers, respect_handler_level=True
+    )
+    listener.start()
+
+
+# QueueHandler to feed entries to a Queue
+# On the other end of the Queue, feed to the webhook
+
 # TODO: Add an async handler for posting
 # Subclass this, create a DiscordChannelHandler, taking a Client and a channel as an argument
 # Then we can handle error channels etc differently