From d614afa952680c215e39e6ff07d9a6addd054faf Mon Sep 17 00:00:00 2001 From: john <32427341+typefield@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:44:18 +0800 Subject: [PATCH 1/3] feat: richText resource extraction and download_file_to_file - normalize: richText picture/file segments now extract into IncomingMessage.resources. Defensive narrowing via isinstance checks; unknown/dirty segments are skipped without affecting the rest; download codes deduped per message. - DingTalkChannel.download_file_to_file: streaming download to a local path (64 KiB chunks, no whole-file buffering), SSRF-guarded like download_file, same-dir temp file + atomic os.replace, parent dir must exist, no partial file left behind. Returns bytes written. - Tests: resource extraction/dedup/dirty-data; streaming download success, missing-parent-dir, and SSRF-block regressions. --- src/dingtalk_channel_sdk/channel.py | 43 ++++++++++ .../normalize/converters/richtext.py | 34 ++++++-- src/dingtalk_channel_sdk/normalize/message.py | 2 +- tests/test_download_to_file.py | 83 +++++++++++++++++++ tests/test_normalize_converters.py | 39 +++++++++ 5 files changed, 195 insertions(+), 6 deletions(-) create mode 100644 tests/test_download_to_file.py diff --git a/src/dingtalk_channel_sdk/channel.py b/src/dingtalk_channel_sdk/channel.py index 4fb8b74..09ad881 100644 --- a/src/dingtalk_channel_sdk/channel.py +++ b/src/dingtalk_channel_sdk/channel.py @@ -4,6 +4,8 @@ import asyncio import json +import os +import tempfile import time import urllib.request from datetime import timedelta @@ -220,6 +222,47 @@ def _fetch() -> bytes: return await asyncio.to_thread(_fetch) + async def download_file_to_file(self, url: str, dest_path: str, timeout: float = 60.0) -> int: + """流式下载文件到本地路径,不整块载入内存。 + + SSRF 防护同 download_file;父目录必须已存在;先写同目录临时文件再 + 原子重命名,失败不落半截文件。返回写入的字节数。 + """ + assert_public_url(url, allowlist=self.cfg.ssrf_allowlist) + + def _fetch() -> int: + dest = os.path.abspath(dest_path) + parent = os.path.dirname(dest) + if not os.path.isdir(parent): + raise FileNotFoundError(f"parent directory does not exist: {parent}") + n = 0 + fd, tmp = tempfile.mkstemp(prefix="." + os.path.basename(dest) + ".tmp-", dir=parent) + tmp_open = True + try: + with os.fdopen(fd, "wb") as out: + tmp_open = False + req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT}) + with urllib.request.urlopen(req, timeout=timeout) as resp: + while True: + chunk = resp.read(64 * 1024) + if not chunk: + break + out.write(chunk) + n += len(chunk) + os.replace(tmp, dest) + tmp = None + return n + finally: + if tmp_open: + os.close(fd) + if tmp is not None: + try: + os.remove(tmp) + except OSError: + pass + + return await asyncio.to_thread(_fetch) + async def mark_thinking(self, conversation_id: str, msg_id: str) -> None: """在用户消息上打"🤔Thinking"状态章(仅人发的消息)。""" await self.emotion.mark_thinking(conversation_id, msg_id) diff --git a/src/dingtalk_channel_sdk/normalize/converters/richtext.py b/src/dingtalk_channel_sdk/normalize/converters/richtext.py index 67f428f..355aac9 100644 --- a/src/dingtalk_channel_sdk/normalize/converters/richtext.py +++ b/src/dingtalk_channel_sdk/normalize/converters/richtext.py @@ -1,21 +1,45 @@ -"""富文本(richText)转换器:拼接正文与 @提及。""" +"""富文本(richText)转换器:拼接正文、@提及与内嵌媒体资源。""" from __future__ import annotations from typing import Any, Dict, List, Tuple -def convert_rich_text(content: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]]]: - """从 richText 数组提取拼接文本与 @提及(userId / 手机号)。""" +def convert_rich_text(content: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]], List[Dict[str, Any]]]: + """从 richText 数组提取拼接文本、@提及(userId / 手机号)与内嵌媒体资源。 + + picture/file 段提取为资源(picture 段值即下载码)。脏数据防御:段值非 + 字符串或下载码为空时跳过该段,不影响其余段落;同一下载码在单条消息内 + 去重。 + """ mentions: List[Dict[str, Any]] = [] + resources: List[Dict[str, Any]] = [] + seen: set = set() parts: List[str] = [] for item in (content or {}).get("richText", []): + if not isinstance(item, dict): + continue item_type = item.get("type", "") if item_type == "text": - parts.append(item.get("text", "")) + text = item.get("text", "") + if isinstance(text, str): + parts.append(text) elif item_type == "at": for uid in item.get("atUserIds") or []: mentions.append({"userId": uid}) for mob in item.get("atMobiles") or []: mentions.append({"userId": mob, "name": mob}) - return "".join(parts), mentions + elif item_type == "picture": + code = item.get("picture") + if isinstance(code, str) and code and code not in seen: + seen.add(code) + resources.append({"type": "image", "downloadCode": code}) + elif item_type == "file": + code = item.get("downloadCode") + name = item.get("fileName") + if isinstance(code, str) and code and code not in seen: + seen.add(code) + resources.append( + {"type": "file", "downloadCode": code, "fileName": name if isinstance(name, str) else ""} + ) + return "".join(parts), mentions, resources diff --git a/src/dingtalk_channel_sdk/normalize/message.py b/src/dingtalk_channel_sdk/normalize/message.py index 231ba2c..60db17d 100644 --- a/src/dingtalk_channel_sdk/normalize/message.py +++ b/src/dingtalk_channel_sdk/normalize/message.py @@ -67,7 +67,7 @@ def parse_content( if msg_type == "text": text = convert_text(content) elif msg_type == "richText": - text, mentions = convert_rich_text(content) + text, mentions, resources = convert_rich_text(content) elif msg_type == "picture": text, resources = convert_picture(content) elif msg_type == "file": diff --git a/tests/test_download_to_file.py b/tests/test_download_to_file.py new file mode 100644 index 0000000..1c4f88b --- /dev/null +++ b/tests/test_download_to_file.py @@ -0,0 +1,83 @@ +"""download_file_to_file 流式落盘单测。""" + +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +from dingtalk_channel_sdk import DingTalkChannel + +MEDIA = b"dingtalk-media-bytes" * 512 + + +def _fake_server(): + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/v1.0/oauth2/accessToken": + self._json(200, {"accessToken": "tok-1", "expireIn": 7200}) + elif self.path == "/media.bin": + body = MEDIA + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self._json(404, {}) + + def _json(self, code, obj): + body = json.dumps(obj).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + return ThreadingHTTPServer(("127.0.0.1", 0), Handler) + + +def _channel(base): + return DingTalkChannel( + client_id="ding-test", + client_secret="s", + api_base=base, + ssrf_allowlist=["127.0.0.1"], + ) + + +async def test_download_file_to_file(tmp_path): + server = _fake_server() + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + base = f"http://127.0.0.1:{server.server_address[1]}" + ch = _channel(base) + + dest = tmp_path / "media.bin" + n = await ch.download_file_to_file(f"{base}/media.bin", str(dest)) + assert n == len(MEDIA) + assert dest.read_bytes() == MEDIA + + # 原有内存下载语义保持不变 + assert await ch.download_file(f"{base}/media.bin") == MEDIA + finally: + server.shutdown() + + +async def test_download_file_to_file_missing_parent(tmp_path): + ch = DingTalkChannel(client_id="a", client_secret="b", ssrf_allowlist=["127.0.0.1"]) + dest = tmp_path / "no-such-dir" / "media.bin" + with pytest.raises(FileNotFoundError): + await ch.download_file_to_file("http://127.0.0.1:1/x", str(dest)) + assert not (tmp_path / "no-such-dir").exists() + + +async def test_download_file_to_file_ssrf_blocked(tmp_path): + ch = DingTalkChannel(client_id="a", client_secret="b") + with pytest.raises(Exception): + await ch.download_file_to_file("http://127.0.0.1:1/x", str(tmp_path / "x.bin")) diff --git a/tests/test_normalize_converters.py b/tests/test_normalize_converters.py index 084a10a..4133511 100644 --- a/tests/test_normalize_converters.py +++ b/tests/test_normalize_converters.py @@ -136,3 +136,42 @@ def test_convert_reply_quote_matrix(): def test_convert_reply_body_only(): assert parse_content("reply", {"text": "只有正文"}, [])[0] == "只有正文" assert parse_content("reply", {}, [])[0] == "[引用消息]" + + +# ── richText 附件资源 ── + + +def test_convert_rich_text_resources(): + content = { + "richText": [ + {"type": "text", "text": "图1 "}, + {"type": "picture", "picture": "dc-1"}, + {"type": "picture", "picture": "dc-1"}, + {"type": "picture", "picture": "dc-2"}, + {"type": "file", "downloadCode": "dc-3", "fileName": "report.pdf"}, + {"type": "text", "text": " 图2"}, + ] + } + text, resources, _ = parse_content("richText", content, []) + assert text == "图1 图2" + assert resources == [ + {"type": "image", "downloadCode": "dc-1"}, + {"type": "image", "downloadCode": "dc-2"}, + {"type": "file", "downloadCode": "dc-3", "fileName": "report.pdf"}, + ] + + +def test_convert_rich_text_resources_dirty_data(): + content = { + "richText": [ + {"type": "picture", "picture": 123}, + {"type": "picture", "picture": ""}, + {"type": "picture"}, + {"type": "file", "downloadCode": 42}, + {"type": "text", "text": "ok"}, + "junk-segment", + ] + } + text, resources, _ = parse_content("richText", content, []) + assert text == "ok" + assert resources == [] From 0508c577e0e8a2a727e39f0689b7f34d6b81d9cb Mon Sep 17 00:00:00 2001 From: typefield Date: Fri, 4 Sep 2026 00:49:02 +0800 Subject: [PATCH 2/3] feat: port richText resource extraction and download_file_to_file from lark channel-sdk - normalize: richText picture/file segments now extract into IncomingMessage.resources (analog of lark's post attachment zone). Defensive narrowing via isinstance checks; unknown/dirty segments are skipped without affecting the rest; download codes deduped per message. - DingTalkChannel.download_file_to_file: streaming download to a local path (64 KiB chunks, no whole-file buffering), SSRF-guarded like download_file, same-dir temp file + atomic os.replace, parent dir must exist, no partial file left behind. Returns bytes written. - Tests: resource extraction/dedup/dirty-data; streaming download success, missing-parent-dir, and SSRF-block regressions. --- src/dingtalk_channel_sdk/channel.py | 2 +- src/dingtalk_channel_sdk/normalize/converters/richtext.py | 6 +++--- tests/test_download_to_file.py | 2 +- tests/test_normalize_converters.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/dingtalk_channel_sdk/channel.py b/src/dingtalk_channel_sdk/channel.py index 09ad881..8e9578f 100644 --- a/src/dingtalk_channel_sdk/channel.py +++ b/src/dingtalk_channel_sdk/channel.py @@ -223,7 +223,7 @@ def _fetch() -> bytes: return await asyncio.to_thread(_fetch) async def download_file_to_file(self, url: str, dest_path: str, timeout: float = 60.0) -> int: - """流式下载文件到本地路径,不整块载入内存。 + """流式下载文件到本地路径,不整块载入内存(对齐 lark channel-sdk 的 downloadResourceToFile)。 SSRF 防护同 download_file;父目录必须已存在;先写同目录临时文件再 原子重命名,失败不落半截文件。返回写入的字节数。 diff --git a/src/dingtalk_channel_sdk/normalize/converters/richtext.py b/src/dingtalk_channel_sdk/normalize/converters/richtext.py index 355aac9..f2534c9 100644 --- a/src/dingtalk_channel_sdk/normalize/converters/richtext.py +++ b/src/dingtalk_channel_sdk/normalize/converters/richtext.py @@ -8,9 +8,9 @@ def convert_rich_text(content: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]], List[Dict[str, Any]]]: """从 richText 数组提取拼接文本、@提及(userId / 手机号)与内嵌媒体资源。 - picture/file 段提取为资源(picture 段值即下载码)。脏数据防御:段值非 - 字符串或下载码为空时跳过该段,不影响其余段落;同一下载码在单条消息内 - 去重。 + 对齐 lark channel-sdk 的富文本附件区能力:picture/file 段提取为资源 + (picture 段值即下载码)。脏数据防御:段值非字符串或下载码为空时跳过 + 该段,不影响其余段落;同一下载码在单条消息内去重。 """ mentions: List[Dict[str, Any]] = [] resources: List[Dict[str, Any]] = [] diff --git a/tests/test_download_to_file.py b/tests/test_download_to_file.py index 1c4f88b..148edf8 100644 --- a/tests/test_download_to_file.py +++ b/tests/test_download_to_file.py @@ -1,4 +1,4 @@ -"""download_file_to_file 流式落盘单测。""" +"""download_file_to_file 流式落盘单测(对齐 lark channel-sdk downloadResourceToFile)。""" from __future__ import annotations diff --git a/tests/test_normalize_converters.py b/tests/test_normalize_converters.py index 4133511..3c880cb 100644 --- a/tests/test_normalize_converters.py +++ b/tests/test_normalize_converters.py @@ -138,7 +138,7 @@ def test_convert_reply_body_only(): assert parse_content("reply", {}, [])[0] == "[引用消息]" -# ── richText 附件资源 ── +# ── richText 附件资源(对齐 lark channel-sdk 富文本附件区)── def test_convert_rich_text_resources(): From 8e138532d77ada8e88926998ce0145b61790a6cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=89=E6=BE=9C?= Date: Sun, 20 Sep 2026 13:53:51 +0800 Subject: [PATCH 3/3] fix: resolve code review findings for richText images and download_file_to_file - Support downloadCode and pictureDownloadCode with fallback to picture for richText images - Use POST with JSON body for messageFiles/download in reply - Prevent SSRF redirect bypass with custom HTTPRedirectHandler - Check Content-Length to prevent overwriting with truncated files on early EOF - Add cancellation check to prevent background thread overwriting target - Provide to_thread compatibility helper for Python 3.8 - Add websockets to dependencies - Add GitHub Actions CI workflow Co-Authored-By: Claude Opus 4.7 --- .github/workflows/ci.yml | 28 +++++++ pyproject.toml | 4 +- src/dingtalk_channel_sdk/channel.py | 56 +++++++++++-- src/dingtalk_channel_sdk/compat.py | 17 ++++ src/dingtalk_channel_sdk/httpx.py | 4 +- src/dingtalk_channel_sdk/media.py | 5 +- .../normalize/converters/richtext.py | 10 ++- src/dingtalk_channel_sdk/reply.py | 8 +- tests/test_download_to_file.py | 82 ++++++++++++++++++- tests/test_normalize_converters.py | 13 +-- 10 files changed, 200 insertions(+), 27 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 src/dingtalk_channel_sdk/compat.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..c888e25 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,28 @@ +name: CI + +on: + push: + branches: ["**"] + pull_request: + branches: ["**"] + +jobs: + test: + name: Python Test + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: "pip" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + - name: Test + run: pytest -v diff --git a/pyproject.toml b/pyproject.toml index b5e44df..c3da2ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,9 @@ readme = "README.md" authors = [{name = "DingTalk Team"}] license = {text = "MIT"} requires-python = ">=3.8" -dependencies = [] +dependencies = [ + "websockets>=10.0", +] [project.optional-dependencies] redis = ["redis>=4.5.0"] diff --git a/src/dingtalk_channel_sdk/channel.py b/src/dingtalk_channel_sdk/channel.py index 8e9578f..52abaa1 100644 --- a/src/dingtalk_channel_sdk/channel.py +++ b/src/dingtalk_channel_sdk/channel.py @@ -6,11 +6,14 @@ import json import os import tempfile +import threading import time import urllib.request from datetime import timedelta from typing import Any, Awaitable, Callable, List, Optional +from .compat import to_thread + from .safety.batching import BatchConfig, BatchedMessage, MessageBatcher from .bot_identity import BotIdentity, BotIdentityProvider from .card import CardClient @@ -40,6 +43,16 @@ RejectHandler = Callable[[RejectEvent], Awaitable[None]] +class _SSRFSafeRedirectHandler(urllib.request.HTTPRedirectHandler): + def __init__(self, allowlist: Optional[List[str]] = None): + super().__init__() + self.allowlist = allowlist + + def redirect_request(self, req, fp, code, msg, headers, newurl): + assert_public_url(newurl, allowlist=self.allowlist) + return super().redirect_request(req, fp, code, msg, headers, newurl) + + class DingTalkChannel: """用法: @@ -217,20 +230,30 @@ async def download_file(self, url: str, timeout: float = 60.0) -> bytes: def _fetch() -> bytes: req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT}) - with urllib.request.urlopen(req, timeout=timeout) as resp: - return resp.read() + opener = urllib.request.build_opener(_SSRFSafeRedirectHandler(allowlist=self.cfg.ssrf_allowlist)) + with opener.open(req, timeout=timeout) as resp: + data = resp.read() + cl = resp.headers.get("Content-Length") + if cl is not None: + try: + expected = int(cl) + if len(data) != expected: + raise OSError(f"download truncated: expected {expected} bytes, got {len(data)}") + except ValueError: + pass + return data - return await asyncio.to_thread(_fetch) + return await to_thread(_fetch) async def download_file_to_file(self, url: str, dest_path: str, timeout: float = 60.0) -> int: - """流式下载文件到本地路径,不整块载入内存(对齐 lark channel-sdk 的 downloadResourceToFile)。 + """流式下载文件到本地路径,不整块载入内存。 SSRF 防护同 download_file;父目录必须已存在;先写同目录临时文件再 原子重命名,失败不落半截文件。返回写入的字节数。 """ assert_public_url(url, allowlist=self.cfg.ssrf_allowlist) - def _fetch() -> int: + def _fetch(cancel_event: threading.Event) -> int: dest = os.path.abspath(dest_path) parent = os.path.dirname(dest) if not os.path.isdir(parent): @@ -242,13 +265,27 @@ def _fetch() -> int: with os.fdopen(fd, "wb") as out: tmp_open = False req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT}) - with urllib.request.urlopen(req, timeout=timeout) as resp: + opener = urllib.request.build_opener(_SSRFSafeRedirectHandler(allowlist=self.cfg.ssrf_allowlist)) + with opener.open(req, timeout=timeout) as resp: + cl = resp.headers.get("Content-Length") + expected: Optional[int] = None + if cl is not None: + try: + expected = int(cl) + except ValueError: + expected = None while True: + if cancel_event.is_set(): + raise RuntimeError("download cancelled") chunk = resp.read(64 * 1024) if not chunk: break out.write(chunk) n += len(chunk) + if expected is not None and n != expected: + raise OSError(f"download truncated: expected {expected} bytes, got {n}") + if cancel_event.is_set(): + raise RuntimeError("download cancelled") os.replace(tmp, dest) tmp = None return n @@ -261,7 +298,12 @@ def _fetch() -> int: except OSError: pass - return await asyncio.to_thread(_fetch) + cancel_ev = threading.Event() + try: + return await to_thread(_fetch, cancel_ev) + except asyncio.CancelledError: + cancel_ev.set() + raise async def mark_thinking(self, conversation_id: str, msg_id: str) -> None: """在用户消息上打"🤔Thinking"状态章(仅人发的消息)。""" diff --git a/src/dingtalk_channel_sdk/compat.py b/src/dingtalk_channel_sdk/compat.py new file mode 100644 index 0000000..99481af --- /dev/null +++ b/src/dingtalk_channel_sdk/compat.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +import asyncio +import functools +import sys +from typing import Any, Callable, TypeVar + +T = TypeVar("T") + + +async def to_thread(func: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Python 3.8+ compatible asyncio.to_thread helper.""" + if sys.version_info >= (3, 9): + return await asyncio.to_thread(func, *args, **kwargs) + loop = asyncio.get_running_loop() + pfunc = functools.partial(func, *args, **kwargs) + return await loop.run_in_executor(None, pfunc) diff --git a/src/dingtalk_channel_sdk/httpx.py b/src/dingtalk_channel_sdk/httpx.py index 20f48a9..225cadb 100644 --- a/src/dingtalk_channel_sdk/httpx.py +++ b/src/dingtalk_channel_sdk/httpx.py @@ -7,6 +7,8 @@ import urllib.request from typing import Any, Dict, Optional +from .compat import to_thread + class ApiError(Exception): """钉钉 API 错误;is_qps_limit 判定 403 + code 含 QpsLimit(SPEC §6)。""" @@ -49,4 +51,4 @@ def _request_sync(method: str, url: str, headers: Dict[str, str], body: Optional async def http_json(method: str, url: str, headers: Optional[Dict[str, str]] = None, body: Optional[dict] = None) -> Dict[str, Any]: - return await asyncio.to_thread(_request_sync, method, url, headers or {}, body) + return await to_thread(_request_sync, method, url, headers or {}, body) diff --git a/src/dingtalk_channel_sdk/media.py b/src/dingtalk_channel_sdk/media.py index 1eee630..5f72ce8 100644 --- a/src/dingtalk_channel_sdk/media.py +++ b/src/dingtalk_channel_sdk/media.py @@ -8,6 +8,7 @@ import urllib.parse import urllib.request +from .compat import to_thread from .config import Config DEFAULT_OAPI_BASE = "https://oapi.dingtalk.com" @@ -36,7 +37,7 @@ async def upload_media(self, media_type: str, filename: str, data: bytes, conten media_type: image | file | video | voice """ - token = await asyncio.to_thread(self._get_token) + token = await to_thread(self._get_token) if not content_type: content_type = "image/jpeg" if media_type == "image" else "application/octet-stream" @@ -64,7 +65,7 @@ def _do_upload() -> dict: except urllib.error.HTTPError as e: raise RuntimeError(f"media/upload: http {e.code} {e.read().decode('utf-8', 'replace')}") from e - out = await asyncio.to_thread(_do_upload) + out = await to_thread(_do_upload) if out.get("errcode") not in (0, None): raise RuntimeError(f"media/upload: errcode={out.get('errcode')} {out.get('errmsg', '')}") media_id = out.get("media_id") or "" diff --git a/src/dingtalk_channel_sdk/normalize/converters/richtext.py b/src/dingtalk_channel_sdk/normalize/converters/richtext.py index f2534c9..d5bebd5 100644 --- a/src/dingtalk_channel_sdk/normalize/converters/richtext.py +++ b/src/dingtalk_channel_sdk/normalize/converters/richtext.py @@ -8,8 +8,8 @@ def convert_rich_text(content: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]], List[Dict[str, Any]]]: """从 richText 数组提取拼接文本、@提及(userId / 手机号)与内嵌媒体资源。 - 对齐 lark channel-sdk 的富文本附件区能力:picture/file 段提取为资源 - (picture 段值即下载码)。脏数据防御:段值非字符串或下载码为空时跳过 + 富文本附件区能力:picture/file 段提取为资源 + (兼容 downloadCode / pictureDownloadCode / picture)。脏数据防御:段值非字符串或下载码为空时跳过 该段,不影响其余段落;同一下载码在单条消息内去重。 """ mentions: List[Dict[str, Any]] = [] @@ -30,7 +30,11 @@ def convert_rich_text(content: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any] for mob in item.get("atMobiles") or []: mentions.append({"userId": mob, "name": mob}) elif item_type == "picture": - code = item.get("picture") + code = item.get("downloadCode") + if not isinstance(code, str) or not code: + code = item.get("pictureDownloadCode") + if not isinstance(code, str) or not code: + code = item.get("picture") if isinstance(code, str) and code and code not in seen: seen.add(code) resources.append({"type": "image", "downloadCode": code}) diff --git a/src/dingtalk_channel_sdk/reply.py b/src/dingtalk_channel_sdk/reply.py index c514dd7..f675301 100644 --- a/src/dingtalk_channel_sdk/reply.py +++ b/src/dingtalk_channel_sdk/reply.py @@ -6,7 +6,6 @@ import time from typing import Optional -from urllib.parse import quote from .card import CardClient, CardStreamer @@ -166,11 +165,10 @@ async def download_url(self, download_code: str, msg_id: str) -> str: """换取消息附件下载地址(E9)。""" token = await self.tokens.get() out = await http_json( - "GET", - f"{self.cfg.api_base}/v1.0/robot/messageFiles/download" - f"?downloadCode={quote(download_code)}&messageId={quote(msg_id)}" - f"&robotCode={quote(self.cfg.client_id)}", + "POST", + f"{self.cfg.api_base}/v1.0/robot/messageFiles/download", {"x-acs-dingtalk-access-token": token}, + {"downloadCode": download_code, "robotCode": self.cfg.client_id}, ) return out.get("downloadUrl", "") diff --git a/tests/test_download_to_file.py b/tests/test_download_to_file.py index 148edf8..0021694 100644 --- a/tests/test_download_to_file.py +++ b/tests/test_download_to_file.py @@ -1,9 +1,11 @@ -"""download_file_to_file 流式落盘单测(对齐 lark channel-sdk downloadResourceToFile)。""" +"""download_file_to_file 流式落盘单测。""" from __future__ import annotations +import asyncio import json import threading +import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest @@ -25,6 +27,19 @@ def do_GET(self): self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) + elif self.path == "/truncated.bin": + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", "1000") + self.end_headers() + self.wfile.write(b"only 10 bytes") + elif self.path == "/slow.bin": + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(MEDIA))) + self.end_headers() + time.sleep(0.3) + self.wfile.write(MEDIA) else: self._json(404, {}) @@ -42,12 +57,12 @@ def log_message(self, *args): return ThreadingHTTPServer(("127.0.0.1", 0), Handler) -def _channel(base): +def _channel(base, allowlist=None): return DingTalkChannel( client_id="ding-test", client_secret="s", api_base=base, - ssrf_allowlist=["127.0.0.1"], + ssrf_allowlist=allowlist if allowlist is not None else ["127.0.0.1"], ) @@ -81,3 +96,64 @@ async def test_download_file_to_file_ssrf_blocked(tmp_path): ch = DingTalkChannel(client_id="a", client_secret="b") with pytest.raises(Exception): await ch.download_file_to_file("http://127.0.0.1:1/x", str(tmp_path / "x.bin")) + + +async def test_download_file_truncated_fails_cleanly(tmp_path): + server = _fake_server() + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + base = f"http://127.0.0.1:{server.server_address[1]}" + ch = _channel(base) + dest = tmp_path / "truncated.bin" + with pytest.raises(OSError, match="download truncated"): + await ch.download_file_to_file(f"{base}/truncated.bin", str(dest)) + assert not dest.exists(), "截断文件不应落盘" + finally: + server.shutdown() + + +async def test_download_file_ssrf_redirect_bypass(tmp_path): + target = _fake_server() + threading.Thread(target=target.serve_forever, daemon=True).start() + + class RedirectHandler(BaseHTTPRequestHandler): + def do_GET(self): + target_port = target.server_address[1] + self.send_response(302) + self.send_header("Location", f"http://127.0.0.1:{target_port}/media.bin") + self.end_headers() + + def log_message(self, *args): + pass + + redirect_server = ThreadingHTTPServer(("127.0.0.1", 0), RedirectHandler) + threading.Thread(target=redirect_server.serve_forever, daemon=True).start() + + try: + redirect_base = f"http://127.0.0.1:{redirect_server.server_address[1]}" + ch = _channel(redirect_base, allowlist=[f"127.0.0.1:{redirect_server.server_address[1]}"]) + dest = tmp_path / "redirect.bin" + with pytest.raises(Exception): + await ch.download_file_to_file(f"{redirect_base}/redirect", str(dest)) + assert not dest.exists() + finally: + redirect_server.shutdown() + target.shutdown() + + +async def test_download_file_cancellation_does_not_overwrite(tmp_path): + server = _fake_server() + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + base = f"http://127.0.0.1:{server.server_address[1]}" + ch = _channel(base) + dest = tmp_path / "slow.bin" + task = asyncio.create_task(ch.download_file_to_file(f"{base}/slow.bin", str(dest))) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0.4) + assert not dest.exists(), "已取消的任务不应覆盖目标文件" + finally: + server.shutdown() diff --git a/tests/test_normalize_converters.py b/tests/test_normalize_converters.py index 3c880cb..62da11d 100644 --- a/tests/test_normalize_converters.py +++ b/tests/test_normalize_converters.py @@ -138,17 +138,19 @@ def test_convert_reply_body_only(): assert parse_content("reply", {}, [])[0] == "[引用消息]" -# ── richText 附件资源(对齐 lark channel-sdk 富文本附件区)── +# ── richText 附件资源 ── def test_convert_rich_text_resources(): content = { "richText": [ {"type": "text", "text": "图1 "}, + {"type": "picture", "downloadCode": "dc-1"}, + {"type": "picture", "pictureDownloadCode": "dc-1"}, {"type": "picture", "picture": "dc-1"}, - {"type": "picture", "picture": "dc-1"}, - {"type": "picture", "picture": "dc-2"}, - {"type": "file", "downloadCode": "dc-3", "fileName": "report.pdf"}, + {"type": "picture", "pictureDownloadCode": "dc-2"}, + {"type": "picture", "picture": "dc-3"}, + {"type": "file", "downloadCode": "dc-4", "fileName": "report.pdf"}, {"type": "text", "text": " 图2"}, ] } @@ -157,7 +159,8 @@ def test_convert_rich_text_resources(): assert resources == [ {"type": "image", "downloadCode": "dc-1"}, {"type": "image", "downloadCode": "dc-2"}, - {"type": "file", "downloadCode": "dc-3", "fileName": "report.pdf"}, + {"type": "image", "downloadCode": "dc-3"}, + {"type": "file", "downloadCode": "dc-4", "fileName": "report.pdf"}, ]