From cc1724508ff4be6dabbf442b10138d46770c1477 Mon Sep 17 00:00:00 2001 From: typefield Date: Fri, 4 Sep 2026 00:49:02 +0800 Subject: [PATCH] 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 | 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..8e9578f 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: + """流式下载文件到本地路径，不整块载入内存（对齐 lark channel-sdk 的 downloadResourceToFile）。 + + 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..f2534c9 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 / 手机号）与内嵌媒体资源。 + + 对齐 lark channel-sdk 的富文本附件区能力：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..148edf8 --- /dev/null +++ b/tests/test_download_to_file.py @@ -0,0 +1,83 @@ +"""download_file_to_file 流式落盘单测（对齐 lark channel-sdk downloadResourceToFile）。""" + +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..3c880cb 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 附件资源（对齐 lark channel-sdk 富文本附件区）── + + +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 == []