ai did his job.

2026-08-08 03:53:10 +03:00
parent 2ffc773113
commit a830623fcd
23 changed files with 3309 additions and 2 deletions
@@ -0,0 +1,346 @@
+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()