diff --git a/app/__main__.py b/app/__main__.py index e5c0b30..811b96f 100644 --- a/app/__main__.py +++ b/app/__main__.py @@ -7,15 +7,20 @@ log = logging.getLogger(__name__) @click.group() -def cli(): - logging.basicConfig(level=logging.INFO) +@click.option( + "--log-level", + default="INFO", + type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False), + help="Set the log level." +) +def cli(log_level: str): + logging.basicConfig(level=log_level.upper()) @cli.command() def run(): """Runs the bot""" log.info("Starting bot.") - logging.getLogger("nio.rooms").setLevel(logging.WARNING) from .main import bot run_async(bot.start(access_token=bot.cfg["bot"]["access_token"])) diff --git a/app/main.py b/app/main.py index 4392cbc..42280c7 100644 --- a/app/main.py +++ b/app/main.py @@ -8,6 +8,10 @@ import tortoise from pathlib import Path log = logging.getLogger(__name__) +logging.getLogger("nio.rooms").setLevel(logging.WARNING) +logging.getLogger("nio.crypto.log").setLevel(logging.WARNING) +logging.getLogger("peewee").setLevel(logging.WARNING) +logging.getLogger("nio.responses").setLevel(logging.WARNING) with open("config.toml", "rb") as fd: @@ -73,3 +77,7 @@ async def on_command_error(ctx: niobot.Context, exc: Exception): @bot.on_event("command_complete") async def on_command_complete(ctx: niobot.Context, _): log.info("Command %s completed.", ctx.command) + + +if __name__ == "__main__": + tortoise.run_async(bot.start(access_token=bot.cfg["bot"]["access_token"])) diff --git a/app/modules/ts_transcode.py b/app/modules/ts_transcode.py new file mode 100644 index 0000000..63d1502 --- /dev/null +++ b/app/modules/ts_transcode.py @@ -0,0 +1,192 @@ +import logging +import shutil +import typing +import urllib.parse + +import niobot +import tempfile +import asyncio +from pathlib import Path + + +class TruthSocialTranscode(niobot.Module): + USER_IDS = { + "@trumpbot:shronk.net", + "@tatebot:shronk.net", + "@nex:nexy7574.co.uk" + } + + def __init__(self, bot: niobot.NioBot): + super().__init__(bot) + self.log = logging.getLogger(__name__) + if not shutil.which("ffprobe"): + raise RuntimeError("ffprobe is required for this module.") + if not shutil.which("ffmpeg"): + raise RuntimeError("ffmpeg is required for this module.") + + # noinspection PyTypeChecker + self.bot.add_event_callback(self.on_message, niobot.RoomMessageVideo) + + def __teardown__(self): + for callback in self.bot.event_callbacks: + if callback.func == self.on_message: + self.bot.event_callbacks.remove(callback) + super().__teardown__() + + 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. + # return + # event: niobot.RoomMessageVideo + + if event.sender not in self.USER_IDS: + # We are not interested in messages from other users. + self.log.info("Ignored user: %s", event.sender) + return "IGNORED" + + # 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 + ) + http_url = "https://" + http_url + name = event.flattened()["content.filename"] + + with tempfile.NamedTemporaryFile(suffix="-%s" % name) as temp_fd: + # async with self.bot.client_session.get(http_url) as resp: + # temp_fd.write(await resp.read()) + # temp_fd.flush() + dl = await self.bot.download(event.url, save_to=Path(temp_fd.name)) + if isinstance(dl, niobot.DownloadError): + self.log.error("Error downloading video: %r", dl) + await self.bot.send_message( + room, + "Error downloading video, unable to transcode: %r" % dl, + 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", + temp_fd.name, + stdout=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")) + await self.bot.send_message( + room, + "Error probing video, unable to determine whether a transcode is required.", + reply_to=event.event_id + ) + return "ERROR_PROBE" + + codec = stdout.decode("utf-8").strip() + if codec in ["h264", "vp8", "vp9"]: + # We don't need to transcode this. + self.log.info("Video is already in a supported format: %s", codec) + return "NO_NEED" + + # We now need to transcode this. For compatibility, we'll force it to be H264 and AAC. + # We will, however, use the High profile for H264, as almost everything supports it now, plus better + # 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 + ) + args = [ + "ffmpeg", + "-hwaccel", + "auto", + "-v", + "error", + "-i", + temp_fd.name, + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-profile:v", + "high", + "-preset", + "medium", + "-b:v", + "5M", + "-maxrate", + "5M", + "-bufsize", + "10M", + "-minrate", + "1M", + "-c:a", + "aac", + "-b:a", + "96k", + "-vf", + "scale=-2:'min(1080,ih)'", + "-movflags", + "faststart", + "-y", + "-f", + "mp4", + out_fd.name + ] + transcode = await asyncio.create_subprocess_exec( + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + self.log.info("Starting transcode...") + stdout, stderr = await transcode.communicate() + self.log.info("Finished transcode.") + if transcode.returncode != 0: + self.log.error("Error transcoding video:\n%s", stderr.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 + ) + 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 + ) + self.log.info("Uploaded transcoded video") + return "OK" + + async def on_message(self, room, event): + try: + return await self.on_message_internal(room, event) + except Exception as e: + self.log.exception("Error processing video", exc_info=e) + return "ERROR %r" % e + + @staticmethod + def is_allowed(ctx: niobot.Context) -> bool: + return ctx.message.sender.split(":", 1)[1] in {"shronk.net", "nicroxio.co.uk", "nexy7574.co.uk"} + + @niobot.command("transcode") + 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...") + ret = await self.on_message(ctx.room, message_link) + await res.edit(f"\N{white heavy check mark} Transcode complete. Result: `{ret}`") +