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", "-v", "warning" "-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\n\n%s", stderr.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 ) 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}`")