TikTok hardening: DNS off-loop, manual redirects, JSON cap, N/A and WebM guards (A7-2..A7-4, A7-6..A7-9, A7-12)

lumen_tiktok.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import ipaddress
+import json
 import logging
 import os
 import re
@@ -93,8 +94,8 @@ def _chunk_tiktok_media_items(items: list, chunk_size: int = TELEGRAM_MEDIA_GROU
 
 
 def _looks_like_video_bytes(data: bytes) -> bool:
-    """Видео или картинка — по "ftyp" на смещении 4: Content-Type от TikWM недостоверен, поле в JSON не помечено."""
-    return len(data) >= 12 and data[4:8] == b"ftyp"
+    """Видео или картинка — по "ftyp" на смещении 4 (MP4) или сигнатуре EBML в начале (WebM): Content-Type от TikWM недостоверен, поле в JSON не помечено."""
+    return len(data) >= 12 and (data[4:8] == b"ftyp" or data[:4] == b"\x1a\x45\xdf\xa3")
 
 
 def _slideshow_slide_urls(media_data: dict, images_to_fetch: list[str]) -> list[str]:
@@ -120,7 +121,10 @@ def _tiktok_video_candidates(media_data: dict) -> list[dict[str, Any]]:
         raw_url = media_data.get(url_key)
         if not raw_url:
             continue
-        if not raw_url.startswith("http"):
+        if raw_url.startswith("//"):
+            # Протокол-относительный URL чиним как https, иначе ниже склеилось бы в битый "https://www.tikwm.com//host/path".
+            raw_url = "https:" + raw_url
+        elif not raw_url.startswith("http"):
             raw_url = "https://www.tikwm.com" + raw_url
         try:
             size_bytes = int(media_data.get(size_key) or 0)
@@ -153,6 +157,28 @@ async def _tikwm_throttle() -> None:
         _tikwm_last_request_ts = now
 
 
+async def _read_capped_body(resp: aiohttp.ClientResponse, url: str) -> bytes | None:
+    # JSON TikWM обязан быть маленьким: читаем с тем же капом, что бинарные
+    # скачивания, иначе скомпрометированный прокси льёт в память без края.
+    content_length = resp.headers.get("Content-Length")
+    if content_length is not None:
+        try:
+            if int(content_length) > TIKTOK_DOWNLOAD_MAX_BYTES:
+                log.warning("[tikwm] Refusing %s: Content-Length %s exceeds the %d byte cap.", url, content_length, TIKTOK_DOWNLOAD_MAX_BYTES)
+                return None
+        except ValueError:
+            pass
+    chunks: list[bytes] = []
+    total = 0
+    async for chunk in resp.content.iter_chunked(65536):
+        total += len(chunk)
+        if total > TIKTOK_DOWNLOAD_MAX_BYTES:
+            log.warning("[tikwm] Aborting %s: exceeded the %d byte cap mid-stream.", url, TIKTOK_DOWNLOAD_MAX_BYTES)
+            return None
+        chunks.append(chunk)
+    return b"".join(chunks)
+
+
 async def _fetch_tikwm_media_data(
     session: aiohttp.ClientSession, resolved_url: str, headers: dict, *, hd: bool = True,
     proxy_base_url: str = "",
@@ -185,7 +211,14 @@ async def _fetch_tikwm_media_data(
                 async with session.get(api_url, timeout=12, headers=api_headers) as r:
                     if r.status == 200:
                         all_attempts_403_empty = False
-                        res = await r.json(content_type=None)
+                        raw_body = await _read_capped_body(r, api_url)
+                        if raw_body is None:
+                            continue
+                        try:
+                            res = json.loads(raw_body.decode("utf-8"))
+                        except Exception as e:
+                            log.warning("[tikwm] Endpoint %s returned undecodable JSON: %s", api_url, e)
+                            continue
                         if res.get("code") == 0 and isinstance(res.get("data"), dict):
                             log.info("[tikwm] Successfully fetched media data from %s", api_url)
                             return res.get("data")
@@ -254,7 +287,9 @@ def _write_mp3_tags(path: str, title: str, artist: str, cover: bytes | None) ->
 def _env_number_or_default(name: str, default: int) -> int:
     try:
         import bot
-        return int(bot._env_number(name, default, cast=int, min_value=1024))
+        # bot._env_number уже кастует и режет мусор дефолтом; здесь только
+        # отложенный импорт (bot.py импортирует этот модуль) и страховка.
+        return bot._env_number(name, default, cast=int, min_value=1024)
     except Exception:
         return default
 TIKTOK_DOWNLOAD_MAX_BYTES = _env_number_or_default("TIKTOK_DOWNLOAD_MAX_BYTES", 75 * 1024 * 1024)
@@ -319,7 +354,7 @@ async def _download_url_bin(session: aiohttp.ClientSession, url: str, headers: d
     if scheme not in ("http", "https"):
         log.warning("[download] Refusing non-HTTP(S) URL (scheme=%r).", scheme)
         return None
-    if not _host_resolves_to_public(urllib.parse.urlsplit(url).hostname):
+    if not await asyncio.to_thread(_host_resolves_to_public, urllib.parse.urlsplit(url).hostname):
         log.warning("[download] Refusing non-public host for URL %r.", url)
         return None
     if headers is None:
@@ -333,7 +368,8 @@ async def _download_url_bin(session: aiohttp.ClientSession, url: str, headers: d
         for _ in range(_DOWNLOAD_MAX_REDIRECTS + 1):
             async with session.get(current_url, headers=headers, timeout=60, allow_redirects=False) as resp:
                 if resp.status in _DOWNLOAD_REDIRECT_STATUSES:
-                    next_url = _checked_redirect_url(current_url, resp.headers.get("Location"))
+                    # Резолв хопа тоже в потоке: внутри getaddrinfo, loop не стопорим.
+                    next_url = await asyncio.to_thread(_checked_redirect_url, current_url, resp.headers.get("Location"))
                     if next_url is None:
                         return None
                     current_url = next_url
@@ -345,7 +381,7 @@ async def _download_url_bin(session: aiohttp.ClientSession, url: str, headers: d
                 # проверен выше. Проверку оставляем как последний рубеж.
                 final_url = getattr(resp, "url", None)
                 final_host = urllib.parse.urlsplit(str(final_url)).hostname if final_url is not None else None
-                if final_host is not None and not _host_resolves_to_public(final_host):
+                if final_host is not None and not await asyncio.to_thread(_host_resolves_to_public, final_host):
                     log.warning("[download] Refusing redirect to non-public host %r.", final_host)
                     return None
                 # Content-Length — быстрый отказ ДО скачивания (сервер может соврать/не прислать — ниже та же проверка потоково по факту).
@@ -436,9 +472,14 @@ async def _probe_video_dimensions(path: str) -> tuple[int, int, int]:
         for line in stdout.decode(errors="replace").splitlines():
             line = line.strip()
             if line.startswith("width="):
-                width = int(float(line.split("=", 1)[1] or 0))
+                raw_w = line.split("=", 1)[1]
+                # "N/A" раньше ронял float() и вместе с ним валидный duration всего пробинга.
+                if raw_w and raw_w != "N/A":
+                    width = int(float(raw_w))
             elif line.startswith("height="):
-                height = int(float(line.split("=", 1)[1] or 0))
+                raw_h = line.split("=", 1)[1]
+                if raw_h and raw_h != "N/A":
+                    height = int(float(raw_h))
             elif line.startswith("duration="):
                 raw = line.split("=", 1)[1]
                 if raw and raw != "N/A":
@@ -482,13 +523,18 @@ async def _probe_and_thumbnail_from_path(path: str) -> tuple[int, int, int, byte
     return duration, width, height, thumb_bytes
 
 
+def _write_file_sync(path: str, data: bytes) -> None:
+    with open(path, "wb") as f:
+        f.write(data)
+
+
 async def _probe_and_thumbnail_from_bytes(video_bytes: bytes) -> tuple[int, int, int, bytes | None]:
     """То же для байтов видео в памяти (видео-слайды): ffprobe/ffmpeg умеют только файлы — пишем во временный файл."""
     try:
         with tempfile.TemporaryDirectory() as tdir:
             raw_path = os.path.join(tdir, "slide.mp4")
-            with open(raw_path, "wb") as f:
-                f.write(video_bytes)
+            # Запись до 50МБ в потоке: в loop стопорила всех.
+            await asyncio.to_thread(_write_file_sync, raw_path, video_bytes)
             return await _probe_and_thumbnail_from_path(raw_path)
     except Exception as probe_exc:
         log.warning("[tiktok] Video-slide metadata probe failed, sending without: %s", probe_exc)
@@ -511,25 +557,47 @@ async def _resolve_tiktok_short(session: aiohttp.ClientSession, url: str) -> str
         "Accept-Language": "ru-RU,ru;q=0.9,en-US;q=0.8,en;q=0.7",
         "Upgrade-Insecure-Requests": "1"
     }
-    try:
-        async with session.head(url, allow_redirects=True, timeout=8, headers=headers) as resp:
-             if resp.status < 400:
-                  resolved = str(resp.url)
-                  if _looks_like_resolved_tiktok_url(resolved):
-                       return resolved
-    except Exception as e:
-         log.warning("[tiktok] HEAD resolution failed: %s", e)
+    # Редиректы разбираем вручную через _checked_redirect_url: aiohttp
+    # с allow_redirects=True сходил бы на следующий хоп до проверки хоста.
+    current_url = url
+    for _ in range(_DOWNLOAD_MAX_REDIRECTS + 1):
+        try:
+            async with session.head(current_url, allow_redirects=False, timeout=8, headers=headers) as resp:
+                if resp.status in _DOWNLOAD_REDIRECT_STATUSES:
+                    next_url = await asyncio.to_thread(_checked_redirect_url, current_url, resp.headers.get("Location"))
+                    if next_url is None:
+                        break
+                    current_url = next_url
+                    continue
+                if resp.status < 400:
+                    resolved = str(resp.url)
+                    if _looks_like_resolved_tiktok_url(resolved):
+                        return resolved
+                break
+        except Exception as e:
+            log.warning("[tiktok] HEAD resolution failed: %s", e)
+            break
 
+    current_url = url
     try:
-        async with session.get(url, allow_redirects=True, timeout=10, headers=headers) as resp:
-             resolved = str(resp.url)
-             if not _looks_like_resolved_tiktok_url(resolved):
-                   # Битый URL после GET тоже логируем — иначе остаётся только малопонятный 403 уже на стороне TikWM.
-                  log.warning(
-                       "[tiktok] Resolving short link %s didn't yield anything that looks like a post URL via either HEAD or GET (result: %s) — passing it through as-is, TikWM may reject it.", url, resolved,
-                  )
-             return resolved
+        for _ in range(_DOWNLOAD_MAX_REDIRECTS + 1):
+            async with session.get(current_url, allow_redirects=False, timeout=10, headers=headers) as resp:
+                if resp.status in _DOWNLOAD_REDIRECT_STATUSES:
+                    next_url = await asyncio.to_thread(_checked_redirect_url, current_url, resp.headers.get("Location"))
+                    if next_url is None:
+                        break
+                    current_url = next_url
+                    continue
+                # Тело осознанно не читаем: нужен только итоговый URL, соединение
+                # просто закроется — капа не требуется, байты не качаем.
+                resolved = str(resp.url)
+                if not _looks_like_resolved_tiktok_url(resolved):
+                    # Битый URL после GET тоже логируем — иначе остаётся только малопонятный 403 уже на стороне TikWM.
+                    log.warning(
+                        "[tiktok] Resolving short link %s didn't yield anything that looks like a post URL via either HEAD or GET (result: %s) — passing it through as-is, TikWM may reject it.", url, resolved,
+                    )
+                return resolved
     except Exception as e:
-         log.warning("[tiktok] GET resolution failed: %s", e)
+        log.warning("[tiktok] GET resolution failed: %s", e)
 
     return url

lumen_tiktok_flow.py

@@ -26,12 +26,25 @@ from lumen_tiktok import (
     _looks_like_video_bytes,
     _slideshow_slide_urls,
     _tiktok_video_candidates,
+    _write_file_sync,
     TikTokUserFacingError,
     _tiktok_music_page_id,
 )
 
 log = logging.getLogger("bot")
 
+
+def _write_and_tag_mp3_sync(path: str, music_bytes: bytes, title: str, artist: str, cover: bytes | None) -> None:
+    import bot
+    with open(path, "wb") as f:
+        f.write(music_bytes)
+    bot._write_mp3_tags(path, title, artist, cover)
+
+
+def _read_file_sync(path: str) -> bytes:
+    with open(path, "rb") as f:
+        return f.read()
+
 async def _send_tiktok_music(session, media_data: dict, message: Message, author: str, headers: dict) -> None:
     import bot
     music_url = media_data.get("music")
@@ -125,12 +138,10 @@ async def _send_tiktok_music(session, media_data: dict, message: Message, author
          try:
               with tempfile.TemporaryDirectory() as tmp_dir:
                    tmp_mp3_path = os.path.join(tmp_dir, "music.mp3")
-                   with open(tmp_mp3_path, "wb") as f:
-                        f.write(music_bytes)
-                   bot._write_mp3_tags(tmp_mp3_path, cleaned_title, performer_name, cover_bytes)
+                   # Запись и теги — синхронные и тяжёлые (до 50МБ): в поток, loop не стопорим.
+                   await asyncio.to_thread(_write_and_tag_mp3_sync, tmp_mp3_path, music_bytes, cleaned_title, performer_name, cover_bytes)
                    if os.path.exists(tmp_mp3_path) and os.path.getsize(tmp_mp3_path) > 0:
-                        with open(tmp_mp3_path, "rb") as f:
-                             tagged_music_bytes = f.read()
+                        tagged_music_bytes = await asyncio.to_thread(_read_file_sync, tmp_mp3_path)
                         # Длительность — явно: без неё Telegram показывает 0:00 (та же история, что была с видео).
                         music_duration = await bot._probe_audio_duration(tmp_mp3_path)
          except Exception as tag_err:
@@ -320,8 +331,8 @@ async def _send_tiktok_single_video(
               try:
                    with tempfile.TemporaryDirectory() as tdir:
                         raw_path = os.path.join(tdir, "raw_tiktok.mp4")
-                        with open(raw_path, "wb") as f:
-                             f.write(video_bytes)
+                        # Запись до 50МБ в потоке: в loop стопорила всех.
+                        await asyncio.to_thread(_write_file_sync, raw_path, video_bytes)
                         # Длительность/размеры — явно: TikTok-контейнер Telegram сам не всегда разбирает (иначе "файл" 0:00).
                         duration, width, height = await bot._probe_video_dimensions(raw_path)
                         thumb_bytes = await bot._generate_video_thumbnail(raw_path, duration)

tests/bot_test_helpers.py

@@ -34,6 +34,34 @@ class _FakeResolveSession:
         return self._get_response
 
 
+class _FakeChainResolveResponse:
+    def __init__(self, status: int, url: str = "", headers=None):
+        self.status = status
+        self.url = url
+        self.headers = headers or {}
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        return False
+
+
+class _FakeChainResolveSession:
+    """head()/get() по карте URL с записью kwargs — для проверки ручной цепочки редиректов."""
+    def __init__(self, mapping: dict):
+        self._mapping = mapping
+        self.calls = []
+
+    def head(self, url, *args, **kwargs):
+        self.calls.append(("head", url, dict(kwargs)))
+        return self._mapping[url][0]
+
+    def get(self, url, *args, **kwargs):
+        self.calls.append(("get", url, dict(kwargs)))
+        return self._mapping[url][1]
+
+
 class _FakeExc(Exception):
     def __init__(self, msg, status_code=None):
         super().__init__(msg)
@@ -276,8 +304,13 @@ class _FakeTikTokBot:
 
 class _FakeTikTokResponse:
     def __init__(self, status=200, json_body=None):
+        import json as _json
         self.status = status
         self._json_body = json_body or {}
+        # Честный фейк: настоящий aiohttp парсит .json() из тела — контент
+        # выводим из того же json, чтобы код шёл боевым путём чтения с капом.
+        self.headers = {}
+        self.content = _FakeDownloadContent([_json.dumps(self._json_body).encode("utf-8")])
 
     async def __aenter__(self):
         return self
@@ -304,9 +337,15 @@ class _FakeTikTokSession:
 
 class _FakeTikwmApiResponse:
     def __init__(self, status=200, json_body=None, body_bytes=b""):
+        import json as _json
         self.status = status
         self._json_body = json_body or {}
         self._body_bytes = body_bytes
+        # Честный фейк (см. _FakeTikTokResponse): контент из json, если байты не заданы.
+        self.headers = {}
+        self.content = _FakeDownloadContent(
+            [body_bytes] if body_bytes else [_json.dumps(self._json_body).encode("utf-8")]
+        )
 
     async def __aenter__(self):
         return self

tests/test_bot_tiktok.py

@@ -6,11 +6,16 @@ test_bot_tiktok.py — TikTok: резолв, слайдшоу, видео, му
 from types import SimpleNamespace
 import asyncio
 import bot
+import functools
 import lumen_tiktok
 import pytest
+import socket
+import threading
 import time
 from tests.bot_test_helpers import (
     _FakeAudioBot,
+    _FakeChainResolveResponse,
+    _FakeChainResolveSession,
     _FakeDownloadResponse,
     _FakeDownloadSession,
     _FakeIncomingMessage,
@@ -1485,3 +1490,132 @@ def test_single_video_status_deleted_after_music():
         bot._delete_message_quietly = original_delete
         lumen_tiktok_flow._send_tiktok_music = original_music
 
+
+# ─────────────── S9: DNS в потоке, ручные редиректы, капы, N/A, WebM ───────────────
+
+def _fake_public_dns(host, port, *args, **kwargs):
+    return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))]
+
+
+def test_download_url_bin_resolves_dns_off_loop(monkeypatch):
+    # A7-2: блокирующий getaddrinfo раньше стопорил loop на время резолва.
+    seen = {}
+
+    def _recording_dns(host, port, *args, **kwargs):
+        seen["thread"] = threading.current_thread()
+        return _fake_public_dns(host, port)
+
+    monkeypatch.setattr(socket, "getaddrinfo", _recording_dns)
+    result = asyncio.run(lumen_tiktok._download_url_bin(
+        _FakeDownloadSession(_FakeDownloadResponse([b"ok"])), "https://example.com/f.mp4"))
+    assert result == b"ok"
+    assert seen["thread"] is not threading.main_thread()
+
+
+def test_resolve_tiktok_short_disables_auto_redirect(monkeypatch):
+    # A7-3: allow_redirects=True ходил на следующий хоп до проверки хоста.
+    monkeypatch.setattr(socket, "getaddrinfo", _fake_public_dns)
+    canonical = "https://www.tiktok.com/@u/video/123"
+    session = _FakeChainResolveSession({
+        "https://vt.tiktok.com/x/": (
+            _FakeChainResolveResponse(301, headers={"Location": canonical}),
+            _FakeChainResolveResponse(200, url="https://vt.tiktok.com/x/"),
+        ),
+        canonical: (
+            _FakeChainResolveResponse(200, url=canonical),
+            _FakeChainResolveResponse(200, url=canonical),
+        ),
+    })
+    result = asyncio.run(bot._resolve_tiktok_short(session, "https://vt.tiktok.com/x/"))
+    assert result == canonical
+    assert session.calls
+    assert all(call[2].get("allow_redirects") is False for call in session.calls)
+
+
+def test_resolve_tiktok_short_refuses_non_public_hop():
+    # A7-3: внутренний адрес в Location останавливает цепочку до запроса к нему.
+    session = _FakeChainResolveSession({
+        "https://vt.tiktok.com/x/": (
+            _FakeChainResolveResponse(301, headers={"Location": "http://169.254.169.254/evil"}),
+            _FakeChainResolveResponse(404, url="https://vt.tiktok.com/x/"),
+        ),
+    })
+    result = asyncio.run(bot._resolve_tiktok_short(session, "https://vt.tiktok.com/x/"))
+    assert "169.254" not in result
+    assert all("169.254" not in url for _, url, _ in session.calls)
+
+
+def test_fetch_tikwm_media_data_skips_huge_content_length():
+    # A7-4: тело JSON раньше читалось без капа — заявленные гигабайты
+    # пропускаем до скачивания.
+    from tests.bot_test_helpers import _FakeTikwmApiResponse, _FakeTikwmApiSession
+    resp = _FakeTikwmApiResponse(status=200, json_body={"code": 0, "data": {}})
+    resp.headers = {"Content-Length": str(10 ** 12)}
+    result = asyncio.run(bot._fetch_tikwm_media_data(
+        _FakeTikwmApiSession([resp]), "https://www.tiktok.com/@u/video/1", {}))
+    assert result is None
+
+
+def test_probe_video_dimensions_survives_na_width(monkeypatch):
+    # A7-6: width=N/A раньше ронял float() и вместе с ним валидный duration.
+    class _FakeFFprobe:
+        returncode = 0
+
+        def __init__(self, out):
+            self._out = out
+
+        async def communicate(self):
+            return (self._out, b"")
+
+        async def wait(self):
+            return 0
+
+        def kill(self):
+            pass
+
+    async def _fake_exec(*args, **kwargs):
+        return _FakeFFprobe(b"width=N/A\nheight=720\nduration=10.5\n")
+
+    monkeypatch.setattr(asyncio, "create_subprocess_exec", _fake_exec)
+    assert asyncio.run(lumen_tiktok._probe_video_dimensions("/tmp/x.mp4")) == (10, 0, 720)
+
+
+def test_probe_thumbnail_from_bytes_writes_off_loop(monkeypatch):
+    # A7-7: запись файла раньше шла в loop.
+    seen = {}
+    real_to_thread = asyncio.to_thread
+
+    def _wrap(func):
+        @functools.wraps(func)
+        def _inner(*args, **kwargs):
+            seen.setdefault("threads", []).append(threading.current_thread())
+            return func(*args, **kwargs)
+        return _inner
+
+    async def _recorder(func, /, *args, **kwargs):
+        return await real_to_thread(_wrap(func), *args, **kwargs)
+
+    monkeypatch.setattr(asyncio, "to_thread", _recorder)
+    assert asyncio.run(lumen_tiktok._probe_and_thumbnail_from_bytes(b"\x00" * 16)) == (0, 0, 0, None)
+    assert seen["threads"] and all(t is not threading.main_thread() for t in seen["threads"])
+
+
+def test_write_and_tag_mp3_sync_roundtrip(tmp_path):
+    # A7-7: синхронный кусок MP3-пути вынесен в хелпер (мутаген битый вход терпит).
+    from lumen_tiktok_flow import _write_and_tag_mp3_sync
+    path = str(tmp_path / "m.mp3")
+    _write_and_tag_mp3_sync(path, b"not a real mp3", "t", "a", None)
+    import os
+    assert os.path.exists(path) and os.path.getsize(path) > 0
+
+
+def test_tiktok_video_candidates_protocol_relative_url():
+    # A7-8: "//host/path" раньше склеивался в битый tikwm-URL.
+    candidates = bot._tiktok_video_candidates({"play": "//cdn.tikwm.com/v.mp4", "size": 1})
+    assert candidates[0]["url"] == "https://cdn.tikwm.com/v.mp4"
+
+
+def test_looks_like_video_bytes_true_for_webm_ebml():
+    # A7-9: WebM-слайд раньше принимался за фото.
+    assert bot._looks_like_video_bytes(b"\x1a\x45\xdf\xa3" + b"\x00" * 16) is True
+