347 lines
12 KiB
Python
347 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
|
|
MAX_VIDEO_SECONDS = 3.0
|
|
MAX_STICKER_BYTES = 256 * 1024
|
|
MAX_INPUT_PIXELS = 8192 * 4096
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MediaConversionError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class IncomingMedia:
|
|
path: Path
|
|
source_unique_id: str
|
|
kind: str
|
|
emoji: str = "❤️"
|
|
content_hash: str | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ConvertedMedia:
|
|
path: Path
|
|
content_hash: str
|
|
|
|
|
|
class MediaConverter:
|
|
def __init__(
|
|
self,
|
|
ffmpeg: str = "ffmpeg",
|
|
ffprobe: str = "ffprobe",
|
|
*,
|
|
process_timeout: float = 60,
|
|
max_parallel: int = 2,
|
|
) -> None:
|
|
self.ffmpeg = ffmpeg
|
|
self.ffprobe = ffprobe
|
|
self.process_timeout = process_timeout
|
|
self._semaphore = asyncio.Semaphore(max_parallel)
|
|
|
|
async def convert(self, media: IncomingMedia) -> ConvertedMedia:
|
|
converted = await self.convert_path(media.path, media.kind)
|
|
content_hash = media.content_hash or await asyncio.to_thread(
|
|
self._sha256, media.path
|
|
)
|
|
return ConvertedMedia(converted.path, content_hash)
|
|
|
|
async def convert_path(self, source: Path, kind: str) -> ConvertedMedia:
|
|
async with self._semaphore:
|
|
try:
|
|
async with asyncio.timeout(self.process_timeout):
|
|
return await self._convert_path_unlocked(source, kind)
|
|
except TimeoutError as error:
|
|
raise MediaConversionError(
|
|
"Обработка медиа превысила допустимое время."
|
|
) from error
|
|
|
|
async def _convert_path_unlocked(self, source: Path, kind: str) -> ConvertedMedia:
|
|
probe = await self._probe(source)
|
|
if not any(
|
|
stream.get("codec_type") == "video" for stream in probe.get("streams", [])
|
|
):
|
|
raise MediaConversionError("В файле не найдено изображение или видео.")
|
|
self._validate_input(probe)
|
|
|
|
if kind == "auto" and self._needs_frame_count(probe):
|
|
probe = await self._probe(source, count_frames=True)
|
|
|
|
resolved_kind = self._resolve_kind(kind, probe)
|
|
if resolved_kind != "static":
|
|
duration = self._duration(probe)
|
|
if duration is None:
|
|
raise MediaConversionError("Не удалось определить длительность видео.")
|
|
if duration > MAX_VIDEO_SECONDS + 0.01:
|
|
raise MediaConversionError("Видео должно быть не длиннее 3 секунд.")
|
|
|
|
fd, output_name = tempfile.mkstemp(suffix=".webm", dir=source.parent)
|
|
os.close(fd)
|
|
Path(output_name).unlink(missing_ok=True)
|
|
try:
|
|
output = Path(output_name)
|
|
last_error = ""
|
|
for crf in (32, 38, 44, 50, 56, 63):
|
|
with source.open("rb") as input_file:
|
|
input_fd = input_file.fileno()
|
|
command = self._command(
|
|
source, output, resolved_kind, crf, input_fd=input_fd
|
|
)
|
|
process = await asyncio.create_subprocess_exec(
|
|
*command,
|
|
stdout=asyncio.subprocess.DEVNULL,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
pass_fds=(input_fd,),
|
|
)
|
|
_, stderr = await self._communicate(process)
|
|
last_error = stderr.decode(errors="replace").strip()
|
|
if (
|
|
process.returncode == 0
|
|
and output.exists()
|
|
and output.stat().st_size <= MAX_STICKER_BYTES
|
|
):
|
|
self._validate_output(await self._probe(output))
|
|
return ConvertedMedia(output, self._sha256(source))
|
|
output.unlink(missing_ok=True)
|
|
detail = (
|
|
last_error.splitlines()[-1]
|
|
if last_error
|
|
else "неизвестная ошибка ffmpeg"
|
|
)
|
|
logger.warning("ffmpeg rejected media: %s", detail)
|
|
raise MediaConversionError("Не удалось преобразовать медиа в видео-стикер.")
|
|
except BaseException:
|
|
Path(output_name).unlink(missing_ok=True)
|
|
raise
|
|
|
|
def _command(
|
|
self,
|
|
source: Path,
|
|
output: Path,
|
|
kind: str,
|
|
crf: int,
|
|
*,
|
|
input_fd: int,
|
|
) -> list[str]:
|
|
command = [
|
|
self.ffmpeg,
|
|
"-y",
|
|
"-v",
|
|
"error",
|
|
"-protocol_whitelist",
|
|
"fd",
|
|
"-fd",
|
|
str(input_fd),
|
|
]
|
|
if kind == "static":
|
|
command.extend(["-loop", "1"])
|
|
command.extend(["-i", "fd:", "-map", "0:v:0", "-an"])
|
|
if kind == "static":
|
|
command.extend(["-t", "1", "-r", "15"])
|
|
else:
|
|
command.extend(["-t", "3", "-r", "30"])
|
|
command.extend(
|
|
[
|
|
"-vf",
|
|
"scale='if(gte(iw,ih),512,-2)':'if(gte(iw,ih),-2,512)'",
|
|
"-c:v",
|
|
"libvpx-vp9",
|
|
"-pix_fmt",
|
|
"yuva420p",
|
|
"-b:v",
|
|
"0",
|
|
"-crf",
|
|
str(crf),
|
|
"-deadline",
|
|
"good",
|
|
"-cpu-used",
|
|
"4",
|
|
"-row-mt",
|
|
"1",
|
|
"-threads",
|
|
"2",
|
|
"-auto-alt-ref",
|
|
"0",
|
|
str(output),
|
|
]
|
|
)
|
|
return command
|
|
|
|
async def _probe(self, source: Path, *, count_frames: bool = False) -> dict:
|
|
try:
|
|
with source.open("rb") as input_file:
|
|
input_fd = input_file.fileno()
|
|
command = self._probe_command(
|
|
count_frames=count_frames, input_fd=input_fd
|
|
)
|
|
process = await asyncio.create_subprocess_exec(
|
|
*command,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
pass_fds=(input_fd,),
|
|
)
|
|
stdout, stderr = await self._communicate(process)
|
|
except FileNotFoundError as error:
|
|
raise MediaConversionError(
|
|
"ffmpeg/ffprobe не установлен на сервере."
|
|
) from error
|
|
if process.returncode != 0:
|
|
detail = stderr.decode(errors="replace").strip().splitlines()
|
|
logger.warning(
|
|
"ffprobe rejected media: %s",
|
|
detail[-1] if detail else "unknown format",
|
|
)
|
|
raise MediaConversionError("Формат медиа не поддерживается.")
|
|
return json.loads(stdout)
|
|
|
|
def _probe_command(self, *, count_frames: bool, input_fd: int) -> list[str]:
|
|
command = [
|
|
self.ffprobe,
|
|
"-v",
|
|
"error",
|
|
"-protocol_whitelist",
|
|
"fd",
|
|
"-fd",
|
|
str(input_fd),
|
|
]
|
|
if count_frames:
|
|
command.extend(["-read_intervals", "%+3.02", "-count_frames"])
|
|
command.extend(
|
|
[
|
|
"-show_streams",
|
|
"-show_format",
|
|
"-of",
|
|
"json",
|
|
"fd:",
|
|
]
|
|
)
|
|
return command
|
|
|
|
@staticmethod
|
|
async def _communicate(
|
|
process: asyncio.subprocess.Process,
|
|
) -> tuple[bytes, bytes]:
|
|
try:
|
|
stdout, stderr = await process.communicate()
|
|
return stdout or b"", stderr or b""
|
|
except BaseException:
|
|
if process.returncode is None:
|
|
process.kill()
|
|
await process.wait()
|
|
raise
|
|
|
|
@staticmethod
|
|
def _duration(probe: dict) -> float | None:
|
|
candidates = [probe.get("format", {}).get("duration")]
|
|
candidates.extend(stream.get("duration") for stream in probe.get("streams", []))
|
|
for value in candidates:
|
|
if value not in (None, "N/A"):
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
for stream in probe.get("streams", []):
|
|
frame_count = stream.get("nb_read_frames") or stream.get("nb_frames")
|
|
frame_rate = stream.get("avg_frame_rate") or stream.get("r_frame_rate")
|
|
if frame_count in (None, "N/A") or frame_rate in (None, "N/A", "0/0"):
|
|
continue
|
|
try:
|
|
fps = float(Fraction(str(frame_rate)))
|
|
if fps > 0:
|
|
return int(frame_count) / fps
|
|
except (TypeError, ValueError, ZeroDivisionError):
|
|
pass
|
|
return None
|
|
|
|
@classmethod
|
|
def _needs_frame_count(cls, probe: dict) -> bool:
|
|
if cls._duration(probe) is not None:
|
|
return False
|
|
animated_codecs = {"apng", "gif", "webp"}
|
|
return any(
|
|
stream.get("codec_type") == "video"
|
|
and stream.get("codec_name") in animated_codecs
|
|
for stream in probe.get("streams", [])
|
|
)
|
|
|
|
@classmethod
|
|
def _resolve_kind(cls, kind: str, probe: dict) -> str:
|
|
if kind != "auto":
|
|
return kind
|
|
duration = cls._duration(probe)
|
|
for stream in probe.get("streams", []):
|
|
value = stream.get("nb_read_frames") or stream.get("nb_frames")
|
|
if value not in (None, "N/A"):
|
|
try:
|
|
if int(value) > 1:
|
|
return "video"
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return "video" if duration is not None and duration > 0.1 else "static"
|
|
|
|
@classmethod
|
|
def _validate_output(cls, probe: dict) -> None:
|
|
videos = [
|
|
stream
|
|
for stream in probe.get("streams", [])
|
|
if stream.get("codec_type") == "video"
|
|
]
|
|
if len(videos) != 1 or videos[0].get("codec_name") != "vp9":
|
|
raise MediaConversionError(
|
|
"Результат конвертации должен быть в формате VP9."
|
|
)
|
|
if any(
|
|
stream.get("codec_type") == "audio" for stream in probe.get("streams", [])
|
|
):
|
|
raise MediaConversionError("Результат конвертации содержит аудио.")
|
|
width = int(videos[0].get("width", 0))
|
|
height = int(videos[0].get("height", 0))
|
|
if min(width, height) <= 0 or max(width, height) != 512:
|
|
raise MediaConversionError("Некорректный размер видео-стикера.")
|
|
duration = cls._duration(probe)
|
|
if duration is None or duration > MAX_VIDEO_SECONDS + 0.01:
|
|
raise MediaConversionError("Некорректная длительность видео-стикера.")
|
|
frame_rate = videos[0].get("avg_frame_rate", "0/1")
|
|
try:
|
|
fps = float(Fraction(str(frame_rate)))
|
|
except (ValueError, ZeroDivisionError):
|
|
fps = 0
|
|
if fps <= 0 or fps > 30.01:
|
|
raise MediaConversionError("Частота кадров видео-стикера превышает 30 FPS.")
|
|
|
|
@staticmethod
|
|
def _validate_input(probe: dict) -> None:
|
|
video = next(
|
|
(
|
|
stream
|
|
for stream in probe.get("streams", [])
|
|
if stream.get("codec_type") == "video"
|
|
),
|
|
None,
|
|
)
|
|
if video is None:
|
|
return
|
|
width = int(video.get("width", 0))
|
|
height = int(video.get("height", 0))
|
|
if min(width, height) <= 0 or width * height > MAX_INPUT_PIXELS:
|
|
raise MediaConversionError("Слишком большое разрешение входного медиа.")
|
|
|
|
@staticmethod
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as file:
|
|
for chunk in iter(lambda: file.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|