200 lines
8 KiB
Python
200 lines
8 KiB
Python
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:
|
|
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",
|
|
"-fps_mode",
|
|
"vfr",
|
|
"-r",
|
|
"25",
|
|
"-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")
|
|
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}`")
|
|
|