Security: validate pick before pop (foreign tap no longer burns token), SSRF guard on download hosts

lumen_commands.py

@@ -578,16 +578,20 @@ async def handle_pick_callback(query: CallbackQuery) -> None:
         with contextlib.suppress(Exception):
             await query.answer()
         return
-    # Pop сразу: повторный тап не плодит второй ответ (чужая кнопка "сгорает" — владелец переспросит текстом).
-    rec = bot._pending_picks.pop(token, None)
+    # Сначала все проверки по записи, удаление — только перед делом: чужой тап
+    # или кривой индекс больше не сжигают кнопку владельца. Атомарность та же,
+    # что у прежнего pop-first: всё синхронно до первого await ниже.
+    rec = bot._pending_picks.get(token)
     # Язык для служебных реплик: из записи (если есть), иначе из чата кнопки.
     _qchat = query.message.chat.id if query.message and query.message.chat else None
     rec_lang = (rec or {}).get("lang") or bot._chat_lang(_qchat)
     if rec is None or rec["expires"] < time.monotonic():
         if rec is not None and query.message is not None:
             # Протухший, но известный выбор — молча выдаём свежие кнопки вместо
-            # стены "протухло, пиши текстом". Токен уже popped: второй такой тап
-            # упрётся в rec None ниже и честно покажет pick_expired. Один ресенд.
+            # стены "протухло, пиши текстом". Старый токен забираем явно (выше
+            # теперь get, а не pop). Второй такой тап упрётся в rec None ниже
+            # и честно покажет pick_expired. Один ресенд.
+            bot._pending_picks.pop(token, None)
             _purge_expired_picks()
             _enforce_pending_picks_cap()
             fresh = secrets.token_hex(4)
@@ -625,6 +629,8 @@ async def handle_pick_callback(query: CallbackQuery) -> None:
         with contextlib.suppress(Exception):
             await query.answer(_lang_t(rec_lang, "pick_not_yours"), show_alert=True)
         return
+    # Все проверки пройдены — только теперь забираем токен (см. комментарий у get выше).
+    bot._pending_picks.pop(token, None)
     choice = options[idx]
     with contextlib.suppress(Exception):
         await query.answer()

lumen_tiktok.py

@@ -8,9 +8,11 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import ipaddress
 import logging
 import os
 import re
+import socket
 import tempfile
 import time
 import urllib.parse
@@ -250,12 +252,46 @@ def _write_mp3_tags(path: str, title: str, artist: str, cover: bytes | None) ->
 TIKTOK_DOWNLOAD_MAX_BYTES = int(os.getenv("TIKTOK_DOWNLOAD_MAX_BYTES", str(75 * 1024 * 1024)))
 
 
+# ── SSRF-гард для _download_url_bin ──
+# Скомпрометированная выдача TikWM могла бы подсунуть внутренний адрес
+# (localhost, metadata-IP облака 169.254.169.254 и т.п.) — сервер сам сходил бы
+# внутрь своей сети. Литералы проверяем без DNS, хосты — через резолв
+# (fail-closed: не резолвится — не качаем). Известное ограничение: TOCTOU между
+# резолвом и коннектом (DNS-rebinding) лечится только на уровне коннектора;
+# для нашей угрозы (прямой внутренний адрес в JSON) предпроверки достаточно.
+def _host_resolves_to_public(host: str | None) -> bool:
+    if not host:
+        return False
+    host = host.strip().strip("[]").lower()
+    if host == "localhost":
+        return False
+    try:
+        return ipaddress.ip_address(host).is_global
+    except ValueError:
+        pass
+    try:
+        infos = socket.getaddrinfo(host, None, type=socket.SOCK_STREAM)
+    except (socket.gaierror, UnicodeError):
+        return False
+    addrs = {info[4][0] for info in infos}
+    if not addrs:
+        return False
+    try:
+        return all(ipaddress.ip_address(a).is_global for a in addrs)
+    except ValueError:
+        return False
+
+
 async def _download_url_bin(session: aiohttp.ClientSession, url: str, headers: dict | None = None) -> bytes | None:
-    # URL — из JSON чужого сервиса (TikWM): качаем только http(s) (AUD-D-003).
+    # URL — из JSON чужого сервиса (TikWM): качаем только http(s) (AUD-D-003)
+    # и только с публичных адресов (SSRF-гард ниже).
     scheme = urllib.parse.urlsplit(url).scheme.lower()
     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):
+        log.warning("[download] Refusing non-public host for URL %r.", url)
+        return None
     if headers is None:
         headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
@@ -266,6 +302,13 @@ async def _download_url_bin(session: aiohttp.ClientSession, url: str, headers: d
         async with session.get(url, headers=headers, timeout=60) as resp:
             if resp.status != 200:
                 return None
+            # Редирект мог увести на внутренний адрес уже после предпроверки —
+            # сверяем конечный хост тоже. У тестовых заглушек .url нет — их пропускаем.
+            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):
+                log.warning("[download] Refusing redirect to non-public host %r.", final_host)
+                return None
             # Content-Length — быстрый отказ ДО скачивания (сервер может соврать/не прислать — ниже та же проверка потоково по факту).
             content_length = resp.headers.get("Content-Length")
             if content_length is not None:

tests/test_bot_commands.py

@@ -694,6 +694,31 @@ def test_pick_callback_wrong_user_rejected(monkeypatch):
         bot._pending_picks.clear()
 
 
+def test_pick_callback_foreign_tap_does_not_burn_token(monkeypatch):
+    # Регрессия: pop до проверки владельца сжигал кнопку — чужак одним тапом
+    # лишал владельца выбора. Теперь чужой тап отклоняется, свой после — работает.
+    bot._pending_picks.clear()
+    bot._pending_picks["dd44ee55"] = {
+        "chat_id": 777, "user_id": 111, "scenario": "film",
+        "original": "посоветуй фильм", "expires": time.monotonic() + 300,
+        "lang": "ru",
+    }
+    fake_core = AsyncMock()
+    monkeypatch.setattr(bot, "_handle_message_core", fake_core)
+    try:
+        q_foreign = _make_pick_query("pick:dd44ee55:0", user_id=999)
+        asyncio.run(bot.handle_pick_callback(q_foreign))
+        assert any(alert is True for _, alert in q_foreign.answered)
+        assert "dd44ee55" in bot._pending_picks
+        q_owner = _make_pick_query("pick:dd44ee55:2")
+        asyncio.run(bot.handle_pick_callback(q_owner))
+        assert "Триллер" in q_owner.message.edits[0][0]
+        fake_core.assert_awaited_once()
+        assert "dd44ee55" not in bot._pending_picks
+    finally:
+        bot._pending_picks.clear()
+
+
 def test_pick_callback_bad_data_answered_quietly():
     for data in ("pick:", "pick:tok:notanint"):
         q = _make_pick_query(data)

tests/test_bot_tiktok.py

@@ -1055,6 +1055,35 @@ def test_download_url_bin_refuses_non_http_scheme_without_fetching():
     assert asyncio.run(lumen_tiktok._download_url_bin(session, "ftp://evil.example/x.mp4")) is None
 
 
+def test_host_resolves_to_public_blocks_internal_addresses():
+    # SSRF-гард: литералы проверяются без DNS, всё непубличное режем.
+    assert lumen_tiktok._host_resolves_to_public("8.8.8.8") is True
+    assert lumen_tiktok._host_resolves_to_public("127.0.0.1") is False
+    assert lumen_tiktok._host_resolves_to_public("10.0.0.5") is False
+    assert lumen_tiktok._host_resolves_to_public("192.168.1.20") is False
+    assert lumen_tiktok._host_resolves_to_public("169.254.169.254") is False
+    assert lumen_tiktok._host_resolves_to_public("::1") is False
+    assert lumen_tiktok._host_resolves_to_public("localhost") is False
+    assert lumen_tiktok._host_resolves_to_public("") is False
+    assert lumen_tiktok._host_resolves_to_public(None) is False
+
+
+def test_download_url_bin_refuses_internal_ips_without_fetching():
+    # Тот же гард сквозно: до session.get дело не доходит вообще.
+    class _ExplodingSession:
+        def get(self, *args, **kwargs):
+            raise AssertionError("no fetch must happen for internal IPs")
+
+    session = _ExplodingSession()
+    for url in (
+        "http://127.0.0.1/evil.mp4",
+        "http://10.1.2.3/evil.mp4",
+        "http://169.254.169.254/latest/meta-data/",
+        "http://[::1]/evil.mp4",
+    ):
+        assert asyncio.run(lumen_tiktok._download_url_bin(session, url)) is None
+
+
 def test_download_url_bin_returns_none_on_non_200_status():
     resp = _FakeDownloadResponse([b"error page"], status=404)
     session = _FakeDownloadSession(resp)