Merge pull request chigwell#206 from vsem-azamat/fix/bounded-media-downloads

fix: bound and stage media downloads

Author: chigwell

telegram_mcp/runtime.py

@@ -703,6 +703,7 @@ async def ensure_connected(cl: TelegramClient = None):
     "edit_chat_photo": {".jpg", ".jpeg", ".png", ".webp"},
 }
 MAX_FILE_BYTES: dict[str, int] = {
+    "download_media": 200 * 1024 * 1024,  # 200 MB
     "send_file": 200 * 1024 * 1024,  # 200 MB
     "upload_file": 200 * 1024 * 1024,
     "send_voice": 100 * 1024 * 1024,

telegram_mcp/tools/media.py

@@ -1,5 +1,10 @@
 """Media MCP tools."""
 
+import os
+import shutil
+import tempfile
+from uuid import uuid4
+
 from telegram_mcp.runtime import *
 
 from telegram_mcp.contact_sheet import ContactSheetUnavailable, build_contact_sheet
@@ -16,6 +21,19 @@
 PHOTO_SHEET_MAXIMUM_TILES = 12
 
 
+def _known_media_size(message) -> Optional[int]:
+    """Return Telegram's declared media size when available."""
+    size = getattr(getattr(message, "file", None), "size", None)
+    if isinstance(size, int) and size >= 0:
+        return size
+    size = getattr(getattr(getattr(message, "media", None), "document", None), "size", None)
+    return size if isinstance(size, int) and size >= 0 else None
+
+
+class _DownloadLimitExceeded(Exception):
+    """Stop an in-progress media download once it exceeds the configured limit."""
+
+
 @mcp.tool(annotations=ToolAnnotations(title="Send File", openWorldHint=True, destructiveHint=True))
 @with_account(readonly=False)
 @validate_id("chat_id")
@@ -202,7 +220,12 @@ async def download_media(
         if not msg or not msg.media:
             return "No media found in the specified message."
 
-        default_name = f"telegram_{chat_id}_{message_id}_{int(time.time())}"
+        limit = MAX_FILE_BYTES["download_media"]
+        declared_size = _known_media_size(msg)
+        if declared_size is not None and declared_size > limit:
+            return f"Media is too large for download_media (limit: {limit} bytes)."
+
+        default_name = f"telegram_{chat_id}_{message_id}_{int(time.time())}_{uuid4().hex}"
         out_path, path_error = await _resolve_writable_file_path(
             raw_path=file_path,
             default_filename=default_name,
@@ -212,23 +235,50 @@ async def download_media(
         if path_error:
             return path_error
 
-        # Strip user-supplied extension so Telethon auto-detects the real media type.
-        # If a path with extension is passed (e.g. ticket.jpg), Telethon writes to that
-        # exact path even if the file is actually a PDF. Stripping the suffix lets
-        # Telethon append the correct extension based on the actual file content.
-        out_path_for_dl = out_path.with_suffix("")
-        downloaded = await cl.download_media(msg, file=str(out_path_for_dl))
-        if not downloaded:
-            return f"Download failed for message {message_id}."
-
-        final_path = Path(downloaded).resolve(strict=True)
-        roots, roots_error = await _ensure_allowed_roots(ctx, "download_media")
-        if roots_error:
-            return roots_error
-        if not _path_is_within_any_root(final_path, roots):
-            return "Download failed: resulting path is outside allowed roots."
-
-        return f"Media downloaded to {final_path}."
+        temp_dir = Path(
+            tempfile.mkdtemp(prefix=".telegram-mcp-download-", dir=str(out_path.parent))
+        )
+        try:
+            staged_name = out_path.with_suffix("").name if out_path.suffix else out_path.name
+            temp_requested_path = temp_dir / staged_name
+
+            def enforce_download_limit(received: int, total: int) -> None:
+                if received > limit or (total and total > limit):
+                    raise _DownloadLimitExceeded
+
+            try:
+                downloaded = await cl.download_media(
+                    msg,
+                    file=str(temp_requested_path),
+                    progress_callback=enforce_download_limit,
+                )
+            except _DownloadLimitExceeded:
+                return f"Media is too large for download_media (limit: {limit} bytes)."
+
+            if not downloaded:
+                return f"Download failed for message {message_id}."
+
+            temp_final_path = Path(downloaded).resolve(strict=True)
+            if temp_final_path.parent != temp_dir.resolve():
+                return "Download failed: resulting temporary path is invalid."
+            if temp_final_path.stat().st_size > limit:
+                return f"Media is too large for download_media (limit: {limit} bytes)."
+
+            final_path = out_path
+            if temp_final_path.suffix:
+                final_path, path_error = await _resolve_writable_file_path(
+                    raw_path=str(out_path.with_suffix(temp_final_path.suffix)),
+                    default_filename=default_name,
+                    ctx=ctx,
+                    tool_name="download_media",
+                )
+                if path_error:
+                    return path_error
+
+            os.replace(temp_final_path, final_path)
+            return f"Media downloaded to {final_path}."
+        finally:
+            shutil.rmtree(temp_dir, ignore_errors=True)
     except Exception as e:
         return log_and_format_error(
             "download_media",

tests/test_media_download.py

@@ -0,0 +1,222 @@
+import inspect
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from telegram_mcp.tools import media
+
+_download_media = inspect.unwrap(media.download_media)
+
+
+async def _call_download(monkeypatch, target, client, message):
+    monkeypatch.setattr(media, "get_client", lambda account=None: client)
+
+    async def _resolve_entity(chat_id, cl):
+        assert cl is client
+        return "entity"
+
+    async def _resolve_path(**kwargs):
+        raw_path = kwargs["raw_path"]
+        return (Path(raw_path).resolve() if raw_path else target), None
+
+    async def _allowed_roots(ctx, tool_name):
+        return [target.parent], None
+
+    monkeypatch.setattr(media, "resolve_entity", _resolve_entity)
+    monkeypatch.setattr(media, "_resolve_writable_file_path", _resolve_path)
+    monkeypatch.setattr(media, "_ensure_allowed_roots", _allowed_roots)
+    return await _download_media(123, message.id, file_path=str(target))
+
+
+async def _call_default_download(monkeypatch, root, client, message):
+    monkeypatch.setattr(media, "get_client", lambda account=None: client)
+
+    async def _resolve_entity(chat_id, cl):
+        assert cl is client
+        return "entity"
+
+    async def _resolve_path(**kwargs):
+        raw_path = kwargs["raw_path"]
+        if raw_path:
+            return Path(raw_path).resolve(), None
+        downloads = root / "downloads"
+        downloads.mkdir(exist_ok=True)
+        return (downloads / kwargs["default_filename"]).resolve(), None
+
+    monkeypatch.setattr(media, "resolve_entity", _resolve_entity)
+    monkeypatch.setattr(media, "_resolve_writable_file_path", _resolve_path)
+    return await _download_media(123, message.id)
+
+
+def _message(size):
+    return SimpleNamespace(id=7, media=object(), file=SimpleNamespace(size=size))
+
+
+class _Client:
+    def __init__(self, message, download):
+        self.message = message
+        self._download = download
+        self.download_calls = 0
+
+    async def get_messages(self, entity, ids):
+        assert entity == "entity"
+        assert ids == self.message.id
+        return self.message
+
+    async def download_media(self, message, file, **kwargs):
+        self.download_calls += 1
+        return await self._download(Path(file), kwargs)
+
+
+@pytest.mark.asyncio
+async def test_default_download_paths_do_not_collide_when_time_is_frozen(tmp_path, monkeypatch):
+    message = _message(size=3)
+    payloads = iter((b"one", b"two"))
+
+    async def successful_download(file, kwargs):
+        downloaded = Path(f"{file}.webp")
+        downloaded.write_bytes(next(payloads))
+        return str(downloaded)
+
+    client = _Client(message, successful_download)
+    monkeypatch.setattr(media.time, "time", lambda: 1_700_000_000)
+
+    first = await _call_default_download(monkeypatch, tmp_path, client, message)
+    second = await _call_default_download(monkeypatch, tmp_path, client, message)
+
+    prefix = "Media downloaded to "
+    first_path = Path(first.removeprefix(prefix).removesuffix("."))
+    second_path = Path(second.removeprefix(prefix).removesuffix("."))
+    assert first_path != second_path
+    assert first_path.read_bytes() == b"one"
+    assert second_path.read_bytes() == b"two"
+
+
+@pytest.mark.asyncio
+async def test_download_media_rejects_declared_oversize_before_transfer(tmp_path, monkeypatch):
+    target = tmp_path / "large.bin"
+    message = _message(size=5)
+
+    async def fail_download(file, kwargs):
+        raise AssertionError("download must not start")
+
+    client = _Client(message, fail_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, target, client, message)
+
+    assert "too large" in result.lower()
+    assert client.download_calls == 0
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_download_media_aborts_in_progress_oversize_and_cleans_staging(
+    tmp_path, monkeypatch
+):
+    target = tmp_path / "bounded.bin"
+    message = _message(size=None)
+
+    async def oversized_download(file, kwargs):
+        partial = Path(f"{file}.bin")
+        partial.write_bytes(b"12345")
+        kwargs["progress_callback"](5, 10)
+        return str(partial)
+
+    client = _Client(message, oversized_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, target, client, message)
+
+    assert "too large" in result.lower()
+    assert client.download_calls == 1
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_download_media_verifies_final_size_and_cleans_staging(tmp_path, monkeypatch):
+    target = tmp_path / "stale-size.bin"
+    message = _message(size=4)
+
+    async def oversized_download(file, kwargs):
+        downloaded = Path(f"{file}.bin")
+        downloaded.write_bytes(b"12345")
+        return str(downloaded)
+
+    client = _Client(message, oversized_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, target, client, message)
+
+    assert "too large" in result.lower()
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_download_media_cleans_extension_appended_file_after_false_return(
+    tmp_path, monkeypatch
+):
+    target = tmp_path / "missing.bin"
+    message = _message(size=3)
+
+    async def failed_download(file, kwargs):
+        Path(f"{file}.part").write_bytes(b"123")
+        return None
+
+    client = _Client(message, failed_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, target, client, message)
+
+    assert "download failed" in result.lower()
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_download_media_cleans_extension_appended_file_after_exception(
+    tmp_path, monkeypatch
+):
+    target = tmp_path / "broken.bin"
+    message = _message(size=3)
+
+    async def failed_download(file, kwargs):
+        Path(f"{file}.part").write_bytes(b"123")
+        raise RuntimeError("network failed")
+
+    client = _Client(message, failed_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, target, client, message)
+
+    assert "error occurred" in result.lower()
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.mark.asyncio
+async def test_download_media_stages_atomically_and_preserves_detected_extension(
+    tmp_path, monkeypatch
+):
+    requested = tmp_path / "photo.jpg"
+    final_path = tmp_path / "photo.webp"
+    final_path.write_bytes(b"old")
+    message = _message(size=3)
+
+    async def successful_download(file, kwargs):
+        assert file != requested.with_suffix("")
+        assert file.is_relative_to(tmp_path)
+        assert final_path.read_bytes() == b"old"
+        kwargs["progress_callback"](3, 3)
+        downloaded = Path(f"{file}.webp")
+        downloaded.write_bytes(b"new")
+        assert final_path.read_bytes() == b"old"
+        return str(downloaded)
+
+    client = _Client(message, successful_download)
+    monkeypatch.setitem(media.MAX_FILE_BYTES, "download_media", 4)
+
+    result = await _call_download(monkeypatch, requested, client, message)
+
+    assert result == f"Media downloaded to {final_path.resolve()}."
+    assert final_path.read_bytes() == b"new"
+    assert list(tmp_path.iterdir()) == [final_path]