From e30779cc1365dc756958dc7f325e7ebae3f863c6 Mon Sep 17 00:00:00 2001
From: dece <shgck@pistache.land>
Date: Sat, 3 Sep 2022 20:43:56 +0200
Subject: [PATCH] shrlok: update and clean plugin (and taxref)

---
 edmond/plugins/shrlok.py | 20 +++++++++-----------
 edmond/plugins/taxref.py | 13 +++++++------
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/edmond/plugins/shrlok.py b/edmond/plugins/shrlok.py
index 7c5f3a4..381e769 100644
--- a/edmond/plugins/shrlok.py
+++ b/edmond/plugins/shrlok.py
@@ -1,5 +1,7 @@
+import json
 import socket
 
+from edmond.bot import Bot
 from edmond.plugin import Plugin
 
 
@@ -10,7 +12,7 @@ class ShrlokPlugin(Plugin):
     plugins to share longer texts through a Web portal.
     """
 
-    def __init__(self, bot):
+    def __init__(self, bot: Bot):
         super().__init__(bot)
         self.socket = ""
         self.file_root = ""
@@ -29,22 +31,18 @@ class ShrlokPlugin(Plugin):
             self.bot.log_d("No socket path specified, shrlok plugin disabled.")
             self.is_ready = False
 
-    def post(self, content_type, content_data):
-        header = '{{"type":"{}"}}'.format(content_type).encode() + b"\0"
+    def post(self, header: dict, data: bytes):
+        encoded_header = json.dumps(header).encode()
         try:
             with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
                 sock.connect(self.socket)
-                data = header + content_data
-                sock.sendall(str(len(data)).encode() + b"\0" + data)
+                header_and_data = encoded_header + b"\0" + data
+                length = len(header_and_data)
+                packet = str(length).encode() + b"\0" + header_and_data
+                sock.sendall(packet)
                 response = sock.recv(4096)
         except OSError as exc:
             self.bot.log_e(f"Can't post data: {exc}")
             return None
         url = response.decode().replace(self.file_root, self.url_root, 1)
         return url
-
-    def post_text(self, text):
-        return self.post("txt", text.encode())
-
-    def post_html(self, html):
-        return self.post("html", html.encode())
diff --git a/edmond/plugins/taxref.py b/edmond/plugins/taxref.py
index 053a8c6..e1d10ca 100644
--- a/edmond/plugins/taxref.py
+++ b/edmond/plugins/taxref.py
@@ -1,10 +1,11 @@
 import random
 import urllib.parse
-from typing import Optional
+from typing import cast, Optional
 
 import requests
 
 from edmond.plugin import Plugin
+from edmond.plugins.shrlok import ShrlokPlugin
 
 BASE_URL = "https://taxref.mnhn.fr/api"
 IMG_FETCH_HTML = """\
@@ -164,13 +165,13 @@ class TaxrefPlugin(Plugin):
         def get_img_url(item) -> Optional[str]:
             return item.get("_links", {}).get("file", {}).get("href")
 
-        if shrlok := self.bot.get_plugin("shrlok"):
+        if shrlok := cast(ShrlokPlugin, self.bot.get_plugin("shrlok")):
             if len(items) > 10:
                 items = random.sample(items, 10)
             urls = map(get_img_url, items)
             urls_text = ",".join(map(lambda url: f'"{url}"', urls))
-            html = IMG_FETCH_HTML.format(urls_text)
-            link = shrlok.post_html(html)  # type: ignore
+            html = IMG_FETCH_HTML.format(urls_text).encode()
+            link = shrlok.post({"type": "raw", "ext": "html"}, html)
             if not link:
                 self.bot.log_d("shrlok plugin returned an empty string.")
         else:
@@ -207,7 +208,7 @@ class TaxrefPlugin(Plugin):
         else:
             # More than one result? For simplicity sake, use the shrlok plugin
             # if available or just show an ambiguous response.
-            if shrlok := self.bot.get_plugin("shrlok"):
+            if shrlok := cast(ShrlokPlugin, self.bot.get_plugin("shrlok")):
                 text = (
                     "\n".join(
                         (
@@ -219,7 +220,7 @@ class TaxrefPlugin(Plugin):
                     )
                     + "\n"
                 )
-                reply = shrlok.post_text(text)  # type: ignore
+                reply = shrlok.post({"type": "txt"}, text.encode())
             else:
                 reply = self.get_ambiguous_reply(items)