nex/nonsensebot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add msc command

Browse source
This commit is contained in:
Nexus 2024-09-15 23:03:24 +01:00
parent 0a694476be
commit b9e635f20b
14 changed files with 378 additions and 160 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Pipfile
View file

@ -12,6 +12,7 @@ httpx = "*"
ollama = "*"
nio-bot = {git = "git+https://github.com/nexy7574/nio-bot.git"}
thefuzz = "*"
redis = "*"
[dev-packages]
ruff = "*"

19
Pipfile.lock generated
View file

@ -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",

17
app/__main__.py
View file

@ -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)

42
app/main.py
View file

@ -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")

8
app/models.py
View file

@ -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")

37
app/modules/ai.py
View file

@ -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():

67
app/modules/evaluation.py
View file

@ -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 &nbsp;
code = code.replace("\u00a0", " ") # 00A0 is &nbsp;
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))

28
app/modules/fediverse_preview.py
View file

@ -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 = "<blockquote>%s</blockquote>" % rendered
body = "<a href=\"%s\">@%s:</a><br>%s" % (
body = '<a href="%s">@%s:</a><br>%s' % (
"https://%s/@%s" % (parsed.netloc, username),
username,
text_body,
)
if data.get("media_attachments"):
body += "<p>{:,} attachments</p>".format(len(data["media_attachments"]))
body += "<p>{:,} attachments</p>".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)

18
app/modules/info.py
View file

@ -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)

24
app/modules/latency.py
View file

@ -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
)
)

108
app/modules/msc_getter.py
View file

@ -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.")

86
app/modules/ts_transcode.py
View file

@ -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}`"
)

12
app/modules/version.py
View file

@ -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))

71
app/modules/yd_dl.py
View file

@ -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:,}"