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)