From b9e635f20bb993004e822743a51bf44df1f356d5 Mon Sep 17 00:00:00 2001 From: nexy7574 Date: Sun, 15 Sep 2024 23:03:24 +0100 Subject: [PATCH] Add msc command --- Pipfile | 1 + Pipfile.lock | 19 ++++-- app/__main__.py | 17 +++-- app/main.py | 42 ++++++++---- app/models.py | 8 --- app/modules/ai.py | 37 ++++++----- app/modules/evaluation.py | 67 +++++++++++++------ app/modules/fediverse_preview.py | 28 +++++--- app/modules/info.py | 18 ++---- app/modules/latency.py | 24 +++---- app/modules/msc_getter.py | 108 ++++++++++++++++++++++++++++++- app/modules/ts_transcode.py | 86 ++++++++++++++---------- app/modules/version.py | 12 +--- app/modules/yd_dl.py | 71 +++++++++++++++----- 14 files changed, 378 insertions(+), 160 deletions(-) delete mode 100644 app/models.py diff --git a/Pipfile b/Pipfile index 494e48a..6a3333d 100644 --- a/Pipfile +++ b/Pipfile @@ -12,6 +12,7 @@ httpx = "*" ollama = "*" nio-bot = {git = "git+https://github.com/nexy7574/nio-bot.git"} thefuzz = "*" +redis = "*" [dev-packages] ruff = "*" diff --git a/Pipfile.lock b/Pipfile.lock index 498864e..3fbac39 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "6da0a4d5a8f4ae40d08db21d7dae5b6c579fb9b750f3ad379f7f286bebd6ba99" + "sha256": "d3f9962ca8bf1f79917897ba8cfb00bce82aacb60a7443816695d01c9b1b0d31" }, "pipfile-spec": 6, "requires": { @@ -671,11 +671,11 @@ }, "idna": { "hashes": [ - "sha256:69297d5da0cc9281c77efffb4e730254dd45943f45bbfb461de5991713989b1e", - "sha256:e5c5dafde284f26e9e0f28f6ea2d6400abd5ca099864a67f576f3981c6476124" + "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", + "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3" ], "markers": "python_version >= '3.6'", - "version": "==3.9" + "version": "==3.10" }, "iso8601": { "hashes": [ @@ -826,7 +826,7 @@ "nio-bot": { "git": "git+https://github.com/nexy7574/nio-bot.git", "markers": "python_version < '3.13' and python_version >= '3.9'", - "ref": "46821599e36868d97d1961aee2ea3d0108de81cf" + "ref": "6aec2074c24bfe07ffd44239ecd4c66b4079ab56" }, "ollama": { "hashes": [ @@ -1257,6 +1257,15 @@ "markers": "python_version >= '3.8'", "version": "==3.9.7" }, + "redis": { + "hashes": [ + "sha256:0c5b10d387568dfe0698c6fad6615750c24170e548ca2deac10c649d463e9870", + "sha256:56134ee08ea909106090934adc36f65c9bcbbaecea5b21ba704ba6fb561f8eb4" + ], + "index": "pypi", + "markers": "python_version >= '3.7'", + "version": "==5.0.8" + }, "referencing": { "hashes": [ "sha256:25b42124a6c8b632a425174f24087783efb348a6f1e0008e63cd4466fedf703c", diff --git a/app/__main__.py b/app/__main__.py index ccb2915..0a94eea 100644 --- a/app/__main__.py +++ b/app/__main__.py @@ -10,10 +10,17 @@ log = logging.getLogger(__name__) @click.option( "--log-level", default="INFO", - type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False), - help="Set the log level." + type=click.Choice( + ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False + ), + help="Set the log level.", +) +@click.option( + "--log-file", + default=None, + type=click.Path(), + help="Log to a file as well as stdout", ) -@click.option("--log-file", default=None, type=click.Path(), help="Log to a file as well as stdout") def cli(log_level: str, log_file: str | None): logging.basicConfig( level=log_level.upper(), @@ -23,7 +30,9 @@ def cli(log_level: str, log_file: str | None): if log_file is not None: file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.DEBUG) - file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")) + file_handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) logging.getLogger().addHandler(file_handler) diff --git a/app/main.py b/app/main.py index 3e37dc8..3674ce6 100644 --- a/app/main.py +++ b/app/main.py @@ -1,14 +1,16 @@ # import sys -import asyncio +import hashlib import sys import niobot import tomllib import platform import logging -import tortoise from pathlib import Path +import redis.asyncio as redis + + log = logging.getLogger(__name__) logging.getLogger("nio.rooms").setLevel(logging.WARNING) logging.getLogger("nio.crypto.log").setLevel(logging.WARNING) @@ -23,17 +25,31 @@ with open("config.toml", "rb") as fd: (store := (Path.cwd() / "store")).mkdir(parents=True, exist_ok=True) -class TortoiseIntegratedBot(niobot.NioBot): +class NonsenseBot(niobot.NioBot): cfg: dict + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.redis: redis.Redis | None = None + + @staticmethod + def redis_key(*parts: str) -> str: + """Returns a redis key for the given info""" + h = hashlib.md5() + for part in parts: + h.update(part.encode()) + return h.hexdigest() + + async def close(self): + await self.redis.aclose() + async def start(self, **kwargs): url = config["database"].get("uri") - if not url: - log.critical("No database URI specified in config.toml - using ephemeral sqlite.") - url = "sqlite://:memory:" + if url: + self.redis = redis.Redis.from_url(url) + assert await self.redis.ping() + sys.path.extend(("..", ".")) - await tortoise.Tortoise.init(db_url=url, modules={"models": ["app.models"]}) - await tortoise.Tortoise.generate_schemas() for file in (Path(__file__).parent / "./modules").glob("*.py"): if file.name.startswith("__"): @@ -50,14 +66,14 @@ class TortoiseIntegratedBot(niobot.NioBot): await super().start(**kwargs) -bot = TortoiseIntegratedBot( +bot = NonsenseBot( homeserver=config["bot"]["homeserver"], user_id=config["bot"]["user_id"], device_id=config["bot"].get("device_id", platform.node()), store_path=str(store.resolve()), command_prefix=config["bot"].get("prefix", "h!"), owner_id=config["bot"].get("owner_id") or "@nex:nexy7574.co.uk", - auto_read_messages=False + auto_read_messages=False, ) bot.cfg = config @@ -76,11 +92,9 @@ async def on_command(ctx: niobot.Context): async def on_command_error(ctx: niobot.Context, exc: Exception): exc = getattr(exc, "original", exc) or exc log.error("Command %s failed.", ctx.command, exc_info=exc) - text = "\N{warning sign} Command failed: `%s`" % exc + text = "\N{WARNING SIGN} Command failed: `%s`" % exc text = text.replace(bot.access_token, "REDACTED") - await ctx.respond( - text - ) + await ctx.respond(text) @bot.on_event("command_complete") diff --git a/app/models.py b/app/models.py deleted file mode 100644 index 476d3a4..0000000 --- a/app/models.py +++ /dev/null @@ -1,8 +0,0 @@ -from tortoise import Model, fields - - -class Account(Model): - mxid: str = fields.CharField(max_length=255, primary_key=True, generated=False) - current_balance: int = fields.IntField(default=0) - bank_balance: int = fields.IntField(default=0) - currency: str = fields.CharField(max_length=3, default="GBP") diff --git a/app/modules/ai.py b/app/modules/ai.py index 22ae2d0..bc2031f 100644 --- a/app/modules/ai.py +++ b/app/modules/ai.py @@ -3,6 +3,7 @@ This module is an old mockup of an AI chatbot system using ollama and a set of s It is not actively used or maintained anymore. """ + import asyncio import json import logging @@ -17,7 +18,7 @@ from ollama import AsyncClient import niobot if typing.TYPE_CHECKING: - from ..main import TortoiseIntegratedBot + from ..main import NonsenseBot @asynccontextmanager @@ -28,7 +29,7 @@ async def ollama_client(url: str) -> AsyncClient: class AIModule(niobot.Module): - bot: "TortoiseIntegratedBot" + bot: "NonsenseBot" log = logging.getLogger(__name__) async def find_server(self, gpu_only: bool = True) -> dict[str, str | bool] | None: @@ -65,7 +66,9 @@ class AIModule(niobot.Module): @niobot.command("whitelist.add") @niobot.is_owner() - async def whitelist_add(self, ctx: niobot.Context, user_id: str, model: str = "llama3:latest"): + async def whitelist_add( + self, ctx: niobot.Context, user_id: str, model: str = "llama3:latest" + ): """[Owner] Adds a user to the whitelist.""" users = self.read_users() users[user_id] = model @@ -104,7 +107,9 @@ class AIModule(niobot.Module): return users[ctx.message.sender] = model self.write_users(users) - await ctx.respond(f"Set model to {model}. Don't forget to pull it with `h!ollama.pull`.") + await ctx.respond( + f"Set model to {model}. Don't forget to pull it with `h!ollama.pull`." + ) @niobot.command("ollama.pull") @niobot.from_homeserver("nexy7574.co.uk", "nicroxio.co.uk", "shronk.net") @@ -112,7 +117,9 @@ class AIModule(niobot.Module): """Pulls the model you set.""" users = self.read_users() if ctx.message.sender not in users: - await ctx.respond("You need to set a model first. See: `h!help ollama.set-model`") + await ctx.respond( + "You need to set a model first. See: `h!help ollama.set-model`" + ) return model = users[ctx.message.sender] server = await self.find_server(gpu_only=False) @@ -137,7 +144,9 @@ class AIModule(niobot.Module): message = " ".join(ctx.args) users = self.read_users() if ctx.message.sender not in users: - await ctx.respond("You need to set a model first. See: `h!help ollama.set-model`") + await ctx.respond( + "You need to set a model first. See: `h!help ollama.set-model`" + ) return model = users[ctx.message.sender] res = await ctx.respond("Finding server...") @@ -148,7 +157,9 @@ class AIModule(niobot.Module): async with ollama_client(server["url"]) as client: await res.edit(content=f"Generating response...") try: - response = await client.chat(model, [{"role": "user", "content": message}]) + response = await client.chat( + model, [{"role": "user", "content": message}] + ) except httpx.HTTPError as e: response = {"message": {"content": f"Error: {e}"}} await res.edit(content=response["message"]["content"]) @@ -162,16 +173,12 @@ class AIModule(niobot.Module): """Checks which servers are online.""" lines: dict[str, dict[str, str | None | bool]] = {} for name, cfg in self.bot.cfg["ollama"].items(): - lines[name] = { - "url": cfg["url"], - "gpu": cfg["gpu"], - "online": None - } + lines[name] = {"url": cfg["url"], "gpu": cfg["gpu"], "online": None} emojis = { - True: "\N{white heavy check mark}", - False: "\N{cross mark}", - None: "\N{hourglass with flowing sand}" + True: "\N{WHITE HEAVY CHECK MARK}", + False: "\N{CROSS MARK}", + None: "\N{HOURGLASS WITH FLOWING SAND}", } def get_lines(): diff --git a/app/modules/evaluation.py b/app/modules/evaluation.py index bf13b08..fc27d0a 100644 --- a/app/modules/evaluation.py +++ b/app/modules/evaluation.py @@ -2,6 +2,7 @@ Module that provides code-evaluation commands. This entire module is locked to NioBot.owner. """ + import logging import shlex import tempfile @@ -21,14 +22,16 @@ import functools class EvalModule(niobot.Module): def __init__(self, bot: niobot.NioBot): if not bot.owner_id: - raise RuntimeError("No owner ID set in niobot. Refusing to load for security reasons.") + raise RuntimeError( + "No owner ID set in niobot. Refusing to load for security reasons." + ) super().__init__(bot) async def owner_check(self, ctx: niobot.Context) -> bool: """Checks if the current user is the owner, and if not, gives them a very informative message.""" if not self.bot.is_owner(ctx.message.sender): await ctx.respond( - "\N{cross mark} Only the owner of this bot can run evaluation commands. Nice try, though!" + "\N{CROSS MARK} Only the owner of this bot can run evaluation commands. Nice try, though!" ) return False return True @@ -38,7 +41,14 @@ class EvalModule(niobot.Module): """Removes any code block syntax from the given string.""" code = code.strip() lines = code.splitlines(False) - if len(lines[0]) == 2 or lines[0] in ("py", "python", "python3", "sh", "shell", "bash"): + if len(lines[0]) == 2 or lines[0] in ( + "py", + "python", + "python3", + "sh", + "shell", + "bash", + ): # likely a codeblock language identifier lines = lines[1:] return "\n".join(lines) @@ -80,11 +90,11 @@ class EvalModule(niobot.Module): "niobot": niobot, "pprint": pprint.pprint, } - code = code.replace("\u00A0", " ") # 00A0 is   + code = code.replace("\u00a0", " ") # 00A0 is   code = textwrap.indent(code, " ") code = f"async def __eval():\n{code}" msg = await ctx.respond(f"Evaluating:\n```py\n{code}\n```") - e = await self.client.add_reaction(ctx.room, ctx.message, "\N{hammer}") + e = await self.client.add_reaction(ctx.room, ctx.message, "\N{HAMMER}") # noinspection PyBroadException try: start = time.time() * 1000 @@ -116,10 +126,12 @@ class EvalModule(niobot.Module): lines.append("Stdout:\n```\n" + stdout.getvalue() + "```") if stderr.getvalue(): lines.append("Stderr:\n```\n" + stderr.getvalue() + "```") - await ctx.client.add_reaction(ctx.room, ctx.message, "\N{white heavy check mark}") + await ctx.client.add_reaction( + ctx.room, ctx.message, "\N{WHITE HEAVY CHECK MARK}" + ) await msg.edit("\n".join(lines)) except Exception: - await ctx.client.add_reaction(ctx.room, ctx.message, "\N{cross mark}") + await ctx.client.add_reaction(ctx.room, ctx.message, "\N{CROSS MARK}") await msg.edit(f"Error:\n```py\n{traceback.format_exc()}```") @niobot.command("shell") @@ -136,7 +148,7 @@ class EvalModule(niobot.Module): msg = await ctx.respond(f"Running command: `{command}`") cmd, args = command.split(" ", 1) - await self.client.add_reaction(ctx.room, ctx.message, "\N{hammer}") + await self.client.add_reaction(ctx.room, ctx.message, "\N{HAMMER}") # noinspection PyBroadException try: with tempfile.TemporaryDirectory() as tmpdir: @@ -149,22 +161,34 @@ class EvalModule(niobot.Module): cwd=tmpdir, ) await proc.wait() - await ctx.client.add_reaction(ctx.room, ctx.message, "\N{white heavy check mark}") + await ctx.client.add_reaction( + ctx.room, ctx.message, "\N{WHITE HEAVY CHECK MARK}" + ) lines = [f"Input:\n```sh\n$ {cmd} {args}\n```\n"] - stdout = (await proc.stdout.read()).decode().replace("```", "`\u200b`\u200b`") - stderr = (await proc.stderr.read()).decode().replace("```", "`\u200b`\u200b`") + stdout = ( + (await proc.stdout.read()) + .decode() + .replace("```", "`\u200b`\u200b`") + ) + stderr = ( + (await proc.stderr.read()) + .decode() + .replace("```", "`\u200b`\u200b`") + ) if stdout: lines.append("Stdout:\n```\n" + stdout + "```\n") if stderr: lines.append("Stderr:\n```\n" + stderr + "```\n") await msg.edit("\n".join(lines)) except Exception: - await ctx.client.add_reaction(ctx.room, ctx.message, "\N{cross mark}") + await ctx.client.add_reaction(ctx.room, ctx.message, "\N{CROSS MARK}") await msg.edit(f"Error:\n```py\n{traceback.format_exc()}```") @niobot.command("thumbnail") @niobot.is_owner() - async def thumbnail(self, ctx: niobot.Context, url: str, width: int = 320, height: int = 240): + async def thumbnail( + self, ctx: niobot.Context, url: str, width: int = 320, height: int = 240 + ): """Get the thumbnail for a URL. URL can be an mxc, in which case C2S will be used to get a thumbnail. @@ -173,10 +197,7 @@ class EvalModule(niobot.Module): parsed = urlparse(url) media_id = parsed.path.lstrip("/") thumbnail = await ctx.bot.thumbnail( - parsed.netloc, - media_id, - width=width, - height=height + parsed.netloc, media_id, width=width, height=height ) if not isinstance(thumbnail, niobot.ThumbnailResponse): return await ctx.respond("Failed to thumbnail: %r" % thumbnail) @@ -188,7 +209,9 @@ class EvalModule(niobot.Module): file = await niobot.ImageAttachment.from_file(body, fn) return await ctx.respond("Downloaded from %r" % url, file=file) - async with aiohttp.ClientSession(headers={"User-Agent": niobot.__user_agent__}) as client: + async with aiohttp.ClientSession( + headers={"User-Agent": niobot.__user_agent__} + ) as client: async with client.get(url) as response: if response.status != 200: await ctx.respond("Error: %d" % response.status) @@ -200,7 +223,9 @@ class EvalModule(niobot.Module): data, ) thumb.save("file.webp", "webp") - attachment = await niobot.ImageAttachment.from_file("file.webp", generate_blurhash=True) + attachment = await niobot.ImageAttachment.from_file( + "file.webp", generate_blurhash=True + ) await ctx.respond("thumbnail.webp", file=attachment) @niobot.command() @@ -232,7 +257,9 @@ class EvalModule(niobot.Module): if room.startswith("#"): resolved = await self.bot.room_resolve_alias(room) if not isinstance(resolved, niobot.RoomResolveAliasResponse): - return await ctx.respond("Failed to resolve room ID for: %r", room) + return await ctx.respond( + "Failed to resolve room ID for: %r", room + ) room = resolved.room_id response = await self.bot.room_leave(room) return await ctx.respond(repr(response)) diff --git a/app/modules/fediverse_preview.py b/app/modules/fediverse_preview.py index cf97eb1..fcebb87 100644 --- a/app/modules/fediverse_preview.py +++ b/app/modules/fediverse_preview.py @@ -4,18 +4,20 @@ This module takes fediverse links (e.g.) and provides a preview in a reply. By default, this will only work with fedi.transgender.ing, but if you set `[fediverse_preview.urls]` in your config, you can add more. This module will use the Mastodon protocol, however it is designed to be used with misskey. """ + import logging import niobot import httpx import typing import textwrap from urllib.parse import urlparse + if typing.TYPE_CHECKING: - from ..main import TortoiseIntegratedBot + from ..main import NonsenseBot class FediversePreviewModule(niobot.Module): - bot: "TortoiseIntegratedBot" + bot: "NonsenseBot" log = logging.getLogger(__name__) @niobot.event("message") @@ -25,12 +27,18 @@ class FediversePreviewModule(niobot.Module): ignore = config.get("ignore", []) if not isinstance(event, niobot.RoomMessageText): return - if event.sender in ignore or room.room_id in ignore or event.sender == self.bot.user_id: + if ( + event.sender in ignore + or room.room_id in ignore + or event.sender == self.bot.user_id + ): return sent = [] async with httpx.AsyncClient( - headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:129.0) Gecko/20100101 Firefox/129.0"} + headers={ + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:129.0) Gecko/20100101 Firefox/129.0" + } ) as client: for item in event.body.split(): if not item.startswith(tuple(supported_prefixes)): @@ -43,7 +51,9 @@ class FediversePreviewModule(niobot.Module): elif len(sent) >= 5: self.log.info("Already sent 5 posts, stopping") break - resp = await client.get("https://%s/api/v1/statuses/%s" % (parsed.netloc, post_id)) + resp = await client.get( + "https://%s/api/v1/statuses/%s" % (parsed.netloc, post_id) + ) if resp.status_code != 200: self.log.error("Got HTTP %d from %s", resp.status_code, resp.url) continue @@ -57,19 +67,21 @@ class FediversePreviewModule(niobot.Module): self.log.info("Detected fediverse post %s: %r", post_id, text) rendered = await self.bot._markdown_to_html(text) text_body = "
%s
" % rendered - body = "@%s:
%s" % ( + body = '@%s:
%s' % ( "https://%s/@%s" % (parsed.netloc, username), username, text_body, ) if data.get("media_attachments"): - body += "

{:,} attachments

".format(len(data["media_attachments"])) + body += "

{:,} attachments

".format( + len(data["media_attachments"]) + ) self.log.info("Sending fediverse post %s", post_id) await self.bot.send_message( room, body, reply_to=event, content_type="html.raw", - override={"body": f"@{username}: {data['content']!r}"} + override={"body": f"@{username}: {data['content']!r}"}, ) sent.append(post_id) diff --git a/app/modules/info.py b/app/modules/info.py index ce09f65..3fa68a2 100644 --- a/app/modules/info.py +++ b/app/modules/info.py @@ -21,10 +21,7 @@ class InfoModule(niobot.Module): resolved: niobot.RoomResolveAliasError return await x.edit( "Sorry, I was unable to resolve that room alias.\n\n{!r}".format( - str(resolved).replace( - self.bot.access_token, - "[REDACTED]" - ) + str(resolved).replace(self.bot.access_token, "[REDACTED]") ) ) room = resolved.room_id @@ -42,10 +39,7 @@ class InfoModule(niobot.Module): self.log.debug("Ratio for %r: %r", room_obj.display_name, ratio) if ratio > 1: matches.append((room_obj, ratio)) - matches.sort( - key=lambda v: v[1], - reverse=True - ) + matches.sort(key=lambda v: v[1], reverse=True) if not matches: return await ctx.respond( "[room] was not a room ID or alias, and I do not know of any mutual rooms that have" @@ -54,9 +48,7 @@ class InfoModule(niobot.Module): room = matches[0][0].room_id if room not in ctx.bot.rooms: - return await ctx.respond( - "Unknown room %r. Am I in it?" % room - ) + return await ctx.respond("Unknown room %r. Am I in it?" % room) room_obj: niobot.MatrixRoom = ctx.bot.rooms[room] lines = [ @@ -70,11 +62,11 @@ class InfoModule(niobot.Module): "* History visibility: {obj.history_visibility!r}", "* User count: {user_count:,} ({invited_count:,} invited)", "* Avatar URL: `{obj.room_avatar_url!s}`", - "* Canonical alias: {obj.canonical_alias}" + "* Canonical alias: {obj.canonical_alias}", ] body = "\n".join(lines).format( obj=room_obj, user_count=room_obj.member_count, - invited_count=room_obj.invited_count + invited_count=room_obj.invited_count, ) return await ctx.respond(body) diff --git a/app/modules/latency.py b/app/modules/latency.py index a1ace6f..411fe44 100644 --- a/app/modules/latency.py +++ b/app/modules/latency.py @@ -3,6 +3,7 @@ This module simply returns the latency between a message being sent and the bot It also supports a rudimentary federation latency check, by sending a request to the target homeserver. """ + import urllib.parse import niobot @@ -11,7 +12,6 @@ import httpx class LatencyModule(niobot.Module): - @niobot.command("latency", aliases=["ping"]) async def latency(self, ctx: niobot.Context, homeserver: str = None): """ @@ -23,10 +23,7 @@ class LatencyModule(niobot.Module): response = await ctx.respond( "Latency: {:,.2f}ms{}".format( - latency, - " (pinging homeserver...)" - if homeserver - else '' + latency, " (pinging homeserver...)" if homeserver else "" ) ) if homeserver: @@ -43,16 +40,22 @@ class LatencyModule(niobot.Module): end_resolv = time.perf_counter() resolv = round(end_resolv - start_resolv, 2) except ValueError as e: - await ctx.respond("\N{WARNING SIGN} Failed to resolve homeserver: " + str(e)) + await ctx.respond( + "\N{WARNING SIGN} Failed to resolve homeserver: " + str(e) + ) homeserver = loc resolv = -1 - async with httpx.AsyncClient(headers={"User-Agent": niobot.__user_agent__}) as client: + async with httpx.AsyncClient( + headers={"User-Agent": niobot.__user_agent__} + ) as client: timings = [] for rnd in range(5): start = time.perf_counter() try: - res = await client.get(f"{homeserver}/_matrix/federation/v1/version") + res = await client.get( + f"{homeserver}/_matrix/federation/v1/version" + ) res.raise_for_status() except httpx.HTTPError: timings.append(-1) @@ -70,9 +73,6 @@ class LatencyModule(niobot.Module): hs_target = homeserver[8:] return await ctx.respond( "Latency: {:,.2f}ms ({:,.2f}ms to resolve homeserver, federation latency to {!r}: {:,.2f}ms)".format( - latency, - resolv * 1000, - hs_target, - fed_latency + latency, resolv * 1000, hs_target, fed_latency ) ) diff --git a/app/modules/msc_getter.py b/app/modules/msc_getter.py index f98c2d6..b5556f1 100644 --- a/app/modules/msc_getter.py +++ b/app/modules/msc_getter.py @@ -1,7 +1,78 @@ +import json +import re +import typing + +import httpx +from pathlib import Path import niobot +if typing.TYPE_CHECKING: + from ..main import NonsenseBot + class MSCGetter(niobot.Module): + bot: NonsenseBot + + def __init__(self, bot): + super().__init__(bot) + self.latest_msc = None + self.msc_cache = Path.cwd() / ".msc-cache" + self.msc_cache.mkdir(parents=True, exist_ok=True) + + async def get_msc_with_cache(self, number: int): + if number not in range(1, 10_000): + return {"error": "Invalid MSC ID"} + + file = self.msc_cache / ("%d.json" % number) + content = None + if file.exists(): + try: + content = json.loads(file.read_text("utf-8", "replace")) + except json.JSONDecodeError: + file.unlink() + + if content: + return content + + async with httpx.AsyncClient() as client: + response = await client.get( + "https://api.github.com/repos/matrix-org/matrix-spec-proposals/pulls/%d" + % number + ) + if response.status_code != 200: + return {"error": "HTTP %d" % response.status_code} + content = response.json() + file.write_text(json.dumps(content)) + return content + + @niobot.event("message") + async def on_message(self, room: niobot.MatrixRoom, message: niobot.RoomMessage): + if self.bot.is_old(message): + return + if message.sender == self.bot.user: + return + if not isinstance(message, niobot.RoomMessageText): + return + if "m.in_reply_to" in message.source.get("m.relates_to", []): + return + if await self.bot.redis.get( + self.bot.redis_key(room.room_id, "auto_msc.enabled") + ): + matches = re.finditer("[MmSsCc]\W?([0-9]{1,4})", message.body) + + lines = [] + for m in matches: + no = m.group(1) + if no: + data = await self.get_msc_with_cache(int(no)) + if data.get("error"): + continue + lines.append(f"[{data['title']}]({data['html_url']})") + if lines: + return await self.bot.send_message( + room, "\n".join((f"* {ln}" for ln in lines)), reply_to=message + ) + @niobot.command() async def msc(self, ctx: niobot.Context, number: str): """Fetches the given MSC""" @@ -9,8 +80,39 @@ class MSCGetter(niobot.Module): number = number[3:] elif number.startswith("#"): number = number[1:] - if not number.isdigit(): + if not number.isdigit() or len(number) != 4: return await ctx.respond("Invalid MXC number.") - return await ctx.respond( - "https://github.com/matrix-org/matrix-spec-proposals/pull/" + number + + data: dict = await self.get_msc_with_cache(int(number)) + if data.get("error"): + return await ctx.respond(data["error"]) + return await ctx.respond(f"[{data['title']}]({data['html_url']})") + + @niobot.command("automsc.enable") + @niobot.sender_has_power(50) + async def auto_msc_enable(self, ctx: niobot.Context): + """Automatically enables MSC linking. Requires a power level of at least 50.""" + key = self.bot.redis_key(ctx.room.room_id, "auto_msc.enabled") + exists = await self.bot.redis.get(key) + if exists: + return await ctx.respond("AutoMSC is already enabled in this room.") + await self.bot.redis.set(key, 1) + await self.bot.redis.save() + return await self.bot.add_reaction( + ctx.room, ctx.message, "\N{WHITE HEAVY CHECK MARK}" ) + + @niobot.command("automsc.disable") + @niobot.sender_has_power(50) + async def auto_msc_disable(self, ctx: niobot.Context): + """Disables automatic MSC linking. Requires a power level of at least 50.""" + key = self.bot.redis_key(ctx.room.room_id, "auto_msc.enabled") + exists = await self.bot.redis.get(key) + if exists: + await self.bot.redis.delete(key) + await self.bot.redis.save() + await self.bot.add_reaction( + ctx.room, ctx.message, "\N{WHITE HEAVY CHECK MARK}" + ) + else: + return await ctx.respond("AutoMSC is already disabled in this room.") diff --git a/app/modules/ts_transcode.py b/app/modules/ts_transcode.py index a889090..f6f6dca 100644 --- a/app/modules/ts_transcode.py +++ b/app/modules/ts_transcode.py @@ -7,6 +7,7 @@ clients cannot play. This can easily be adapted to other rooms and users, or even to other types of media, by changing the USER_IDS list. """ + import logging import shlex import shutil @@ -20,11 +21,7 @@ from pathlib import Path class TruthSocialTranscode(niobot.Module): - USER_IDS = { - "@trumpbot:shronk.net", - "@tatebot:shronk.net", - "@nex:nexy7574.co.uk" - } + USER_IDS = {"@trumpbot:shronk.net", "@tatebot:shronk.net", "@nex:nexy7574.co.uk"} def __init__(self, bot: niobot.NioBot): super().__init__(bot) @@ -43,7 +40,9 @@ class TruthSocialTranscode(niobot.Module): self.bot.event_callbacks.remove(callback) super().__teardown__() - async def on_message_internal(self, room: niobot.MatrixRoom, event: niobot.RoomMessageVideo) -> str: + async def on_message_internal( + self, room: niobot.MatrixRoom, event: niobot.RoomMessageVideo + ) -> str: self.log.info("Processing event: %r (%r)", event, room) # if not isinstance(event, niobot.RoomMessageVideo): # # We are not interested in non-videos. @@ -58,7 +57,9 @@ class TruthSocialTranscode(niobot.Module): # Videos sent in the truth room are not encrypted, so that makes this easier. http_url = await self.bot.mxc_to_http( event.url, - homeserver=event.sender.split(":", 1)[1] # use the sender homeserver for download + homeserver=event.sender.split(":", 1)[ + 1 + ], # use the sender homeserver for download ) http_url = "https://" + http_url name = event.flattened()["content.filename"] @@ -73,26 +74,35 @@ class TruthSocialTranscode(niobot.Module): await self.bot.send_message( room, "Error downloading video, unable to transcode: %r" % dl, - reply_to=event.event_id + reply_to=event.event_id, ) return "ERROR_DOWNLOAD" # First, probe the video to see if it has H265 content in it. probe = await asyncio.create_subprocess_exec( - "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", - "stream=codec_name", "-of", "default=noprint_wrappers=1:nokey=1", + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=codec_name", + "-of", + "default=noprint_wrappers=1:nokey=1", temp_fd.name, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await probe.communicate() if probe.returncode != 0: # Something went wrong, let's not continue. - self.log.error("Error probing video:\n%s", (stderr or b"EOF").decode("utf-8")) + self.log.error( + "Error probing video:\n%s", (stderr or b"EOF").decode("utf-8") + ) await self.bot.send_message( room, "Error probing video, unable to determine whether a transcode is required.", - reply_to=event.event_id + reply_to=event.event_id, ) return "ERROR_PROBE" @@ -107,12 +117,10 @@ class TruthSocialTranscode(niobot.Module): # compression. # We will also limit the bitrate to about 5 megabit, because Trump has a habit of posting raw footage # that ends up being something stupid like 20 megabit at 60fps. - with tempfile.NamedTemporaryFile(suffix="-transcoded-" + name + ".mp4") as out_fd: - self.log.info( - "transcoding %s -> %s", - temp_fd.name, - out_fd.name - ) + with tempfile.NamedTemporaryFile( + suffix="-transcoded-" + name + ".mp4" + ) as out_fd: + self.log.info("transcoding %s -> %s", temp_fd.name, out_fd.name) args = [ "ffmpeg", "-v", @@ -148,12 +156,12 @@ class TruthSocialTranscode(niobot.Module): # "vfr", "-f", "mp4", - out_fd.name + out_fd.name, ] transcode = await asyncio.create_subprocess_exec( *args, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + stderr=asyncio.subprocess.PIPE, ) self.log.info("Starting transcode...") self.log.debug("Running command: %r", shlex.join(args)) @@ -163,21 +171,20 @@ class TruthSocialTranscode(niobot.Module): self.log.error( "Error transcoding video:\n%s\n\n%s", stderr.decode("utf-8"), - stdout.decode("utf-8") + stdout.decode("utf-8"), ) await self.bot.send_message( room, - "Error transcoding video, unable to transcode: got return code %d" % transcode.returncode, - reply_to=event.event_id + "Error transcoding video, unable to transcode: got return code %d" + % transcode.returncode, + reply_to=event.event_id, ) return "ERROR_TRANSCODE" # We transcoded! Now we need to upload it. file = await niobot.VideoAttachment.from_file(out_fd.name) self.log.info("Uploading transcoded video") await self.bot.send_message( - room.room_id, - file=file, - reply_to=event.event_id + room.room_id, file=file, reply_to=event.event_id ) self.log.info("Uploaded transcoded video") return "OK" @@ -191,18 +198,31 @@ class TruthSocialTranscode(niobot.Module): @staticmethod def is_allowed(ctx: niobot.Context) -> bool: - return ctx.message.sender.split(":", 1)[1] in {"shronk.net", "nicroxio.co.uk", "nexy7574.co.uk"} + return ctx.message.sender.split(":", 1)[1] in { + "shronk.net", + "nicroxio.co.uk", + "nexy7574.co.uk", + } @niobot.command("transcode") - @niobot.from_homeserver("nexy7574.co.uk", "shronk.net", "nicroxio.co.uk", "transgender.ing") - async def do_transcode(self, ctx: niobot.Context, message_link: typing.Annotated[niobot.RoomMessage, niobot.EventParser("m.room.message")]): + @niobot.from_homeserver( + "nexy7574.co.uk", "shronk.net", "nicroxio.co.uk", "transgender.ing" + ) + async def do_transcode( + self, + ctx: niobot.Context, + message_link: typing.Annotated[ + niobot.RoomMessage, niobot.EventParser("m.room.message") + ], + ): """ Transcodes a video to H264/AAC. """ if not isinstance(message_link, niobot.RoomMessageVideo): await ctx.respond("That's not a video. %r" % message_link) return - res = await ctx.respond("\N{hourglass} Dispatching, please wait...") + res = await ctx.respond("\N{HOURGLASS} Dispatching, please wait...") ret = await self.on_message(ctx.room, message_link) - await res.edit(f"\N{white heavy check mark} Transcode complete. Result: `{ret}`") - + await res.edit( + f"\N{WHITE HEAVY CHECK MARK} Transcode complete. Result: `{ret}`" + ) diff --git a/app/modules/version.py b/app/modules/version.py index fb4744c..18127c7 100644 --- a/app/modules/version.py +++ b/app/modules/version.py @@ -1,6 +1,7 @@ """ Displays versions for the bot and its dependencies """ + import niobot import asyncio import tomllib @@ -15,11 +16,7 @@ class VersionModule(niobot.Module): pipfile = tomllib.load(pipfile_fd) process = await asyncio.create_subprocess_exec( - "git", - "rev-parse", - "--short", - "HEAD", - stdout=asyncio.subprocess.PIPE + "git", "rev-parse", "--short", "HEAD", stdout=asyncio.subprocess.PIPE ) stdout, _ = await process.communicate() git_version = stdout.decode().strip() @@ -28,9 +25,6 @@ class VersionModule(niobot.Module): for dependency in pipfile["packages"].keys(): dep_version = version(dependency) lines.append( - "* {0}: [{1}](https://pypi.org/p/{0})".format( - dependency, - dep_version - ) + "* {0}: [{1}](https://pypi.org/p/{0})".format(dependency, dep_version) ) return await ctx.respond("\n".join(lines)) diff --git a/app/modules/yd_dl.py b/app/modules/yd_dl.py index 8c9b04a..195213f 100644 --- a/app/modules/yd_dl.py +++ b/app/modules/yd_dl.py @@ -1,6 +1,7 @@ """ What it says on the tin - a command wrapper around yt-dlp. """ + import datetime import functools import math @@ -94,7 +95,9 @@ class YoutubeDLModule(niobot.Module): str(new_file), ] self.log.debug("Running command: ffmpeg %s", " ".join(args)) - process = subprocess.run(["ffmpeg", *args], stdout=subprocess.PIPE, stderr=subprocess.PIPE) + process = subprocess.run( + ["ffmpeg", *args], stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) if process.returncode != 0: raise RuntimeError(process.stderr.decode()) return new_file @@ -102,11 +105,15 @@ class YoutubeDLModule(niobot.Module): return await asyncio.to_thread(inner) @staticmethod - async def upload_to_0x0(name: str, data: typing.IO[bytes], mime_type: str | None = None) -> str: + async def upload_to_0x0( + name: str, data: typing.IO[bytes], mime_type: str | None = None + ) -> str: if not mime_type: import magic - mime_type = await asyncio.to_thread(magic.from_buffer, data.read(4096), mime=True) + mime_type = await asyncio.to_thread( + magic.from_buffer, data.read(4096), mime=True + ) data.seek(0) async with httpx.AsyncClient() as client: response = await client.post( @@ -123,7 +130,11 @@ class YoutubeDLModule(niobot.Module): @niobot.command() async def ytdl( - self, ctx: niobot.Context, url: str, snip: Optional[str] = None, download_format: Optional[str] = None + self, + ctx: niobot.Context, + url: str, + snip: Optional[str] = None, + download_format: Optional[str] = None, ): """ Downloads a video from YouTube or other source @@ -161,7 +172,9 @@ class YoutubeDLModule(niobot.Module): await response.edit(content="Fetching metadata (step 1/10)") try: # noinspection PyTypeChecker - extracted_info = await asyncio.to_thread(downloader.extract_info, url, download=False) + extracted_info = await asyncio.to_thread( + downloader.extract_info, url, download=False + ) except DownloadError as e: extracted_info = { "title": "error", @@ -185,22 +198,38 @@ class YoutubeDLModule(niobot.Module): title = textwrap.shorten(title, 100) webpage_url = extracted_info.get("webpage_url", url) - chosen_format = extracted_info.get("format") or chosen_format or str(uuid.uuid4()) - chosen_format_id = extracted_info.get("format_id") or str(uuid.uuid4()) + chosen_format = ( + extracted_info.get("format") + or chosen_format + or str(uuid.uuid4()) + ) + chosen_format_id = extracted_info.get("format_id") or str( + uuid.uuid4() + ) final_extension = extracted_info.get("ext") or "mp4" - format_note = extracted_info.get("format_note", "%s (%s)" % (chosen_format, chosen_format_id)) or "" + format_note = ( + extracted_info.get( + "format_note", "%s (%s)" % (chosen_format, chosen_format_id) + ) + or "" + ) resolution = extracted_info.get("resolution") or "1x1" fps = extracted_info.get("fps", 0.0) or 0.0 vcodec = extracted_info.get("vcodec") or "h264" acodec = extracted_info.get("acodec") or "aac" - filesize = extracted_info.get("filesize", extracted_info.get("filesize_approx", 1)) - likes = extracted_info.get("like_count", extracted_info.get("average_rating", 0)) + filesize = extracted_info.get( + "filesize", extracted_info.get("filesize_approx", 1) + ) + likes = extracted_info.get( + "like_count", extracted_info.get("average_rating", 0) + ) views = extracted_info.get("view_count", 0) lines = [] if chosen_format and chosen_format_id: lines.append( - "* Chosen format: `%s` (`%s`)" % (chosen_format, chosen_format_id), + "* Chosen format: `%s` (`%s`)" + % (chosen_format, chosen_format_id), ) if format_note: lines.append("* Format note: %r" % format_note) @@ -225,14 +254,18 @@ class YoutubeDLModule(niobot.Module): ) try: - await asyncio.to_thread(functools.partial(downloader.download, [url])) + await asyncio.to_thread( + functools.partial(downloader.download, [url]) + ) except DownloadError as e: logging.error(e, exc_info=True) return await response.edit( f"# Error!\n\nDownload failed:\n```\n{e}\n```", ) try: - file: Path = next(temp_dir.glob("*." + extracted_info.get("ext", "*"))) + file: Path = next( + temp_dir.glob("*." + extracted_info.get("ext", "*")) + ) except StopIteration: ext = extracted_info.get("ext", "*") self.log.warning( @@ -252,7 +285,9 @@ class YoutubeDLModule(niobot.Module): except ValueError: trim_start, trim_end = snip, None trim_start = trim_start or "00:00:00" - trim_end = trim_end or extracted_info.get("duration_string", "00:30:00") + trim_end = trim_end or extracted_info.get( + "duration_string", "00:30:00" + ) new_file = temp_dir / ("output" + file.suffix) args = [ "-hwaccel", @@ -308,7 +343,9 @@ class YoutubeDLModule(niobot.Module): size_megabits = (size_bytes * 8) / 1024 / 1024 eta_seconds = size_megabits / 20 - await response.edit(content=f"Uploading (ETA: {humanize.naturaldelta(eta_seconds)})...") + await response.edit( + content=f"Uploading (ETA: {humanize.naturaldelta(eta_seconds)})..." + ) views = views or 0 likes = likes or 0 try: @@ -329,7 +366,9 @@ class YoutubeDLModule(niobot.Module): httpx.HTTPStatusError, ) as e: self.log.error(e, exc_info=True) - await response.edit(content=f"# Error\n\nUpload failed:\n```\n{e}\n```") + await response.edit( + content=f"# Error\n\nUpload failed:\n```\n{e}\n```" + ) else: await response.edit( content=f"# [Downloaded {title}!]({webpage_url})\n\nViews: {views:,} | Likes: {likes:,}"