nonsensebot/app/modules/ts_transcode.py

230 lines
8.8 KiB
Python

"""
This module actively listens for messages in a specific room from specific users, and checks if they contain
a video which uses an unsupported video codec. If it does, it will be transcoded and re-uploaded to the room.
This is used because we mirror Truth Social into a matrix room and sometimes the videos are in H265, which a lot of
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
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:
if self.bot.is_old(event):
return "IGNORED"
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)',fps=25",
"-movflags",
"faststart",
"-y",
# "-fps_mode",
# "vfr",
"-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...")
self.log.debug("Running command: %r", shlex.join(args))
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")
@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...")
ret = await self.on_message(ctx.room, message_link)
await res.edit(
f"\N{WHITE HEAVY CHECK MARK} Transcode complete. Result: `{ret}`"
)