Merge pull request #1427 from AstrBotDevs/fix-gewechat

Improve: 优化 Gewechat 下文件回调逻辑

astrbot/core/platform/sources/gewechat/client.py

@@ -3,6 +3,7 @@ import base64
 import datetime
 import os
 import re
+import uuid
 import threading
 
 import aiohttp
@@ -63,7 +64,7 @@ class SimpleGewechatClient:
             "/astrbot-gewechat/callback", view_func=self._callback, methods=["POST"]
         )
         self.server.add_url_rule(
-            "/astrbot-gewechat/file/<file_id>",
+            "/astrbot-gewechat/file/<file_token>",
             view_func=self._handle_file,
             methods=["GET"],
         )
@@ -81,6 +82,11 @@ class SimpleGewechatClient:
 
         self.shutdown_event = asyncio.Event()
 
+        self.staged_files = {}
+        """存储了允许外部访问的文件列表。auth_token: file_path。通过 register_file 方法注册。"""
+
+        self.lock = asyncio.Lock()
+
     async def get_token_id(self):
         """获取 Gewechat Token。"""
         async with aiohttp.ClientSession() as session:
@@ -310,9 +316,33 @@ class SimpleGewechatClient:
 
         return quart.jsonify({"r": "AstrBot ACK"})
 
-    async def _handle_file(self, file_id):
-        file_path = f"data/temp/{file_id}"
-        return await quart.send_file(file_path)
+    async def _register_file(self, file_path: str) -> str:
+        """向 AstrBot 回调服务器 注册一个允许外部访问的文件。
+
+        Args:
+            file_path (str): 文件路径。
+        Returns:
+            str: 返回一个 auth_token，文件路径为 file_path。通过 /astrbot-gewechat/file/auth_token 得到文件。
+        """
+        async with self.lock:
+            if not os.path.exists(file_path):
+                raise Exception(f"文件不存在: {file_path}")
+
+            file_token = str(uuid.uuid4())
+            self.staged_files[file_token] = file_path
+            return file_token
+
+    async def _handle_file(self, file_token):
+        async with self.lock:
+            if file_token not in self.staged_files:
+                logger.warning(f"请求的文件 {file_token} 不存在。")
+                return quart.abort(404)
+            if not os.path.exists(self.staged_files[file_token]):
+                logger.warning(f"请求的文件 {self.staged_files[file_token]} 不存在。")
+                return quart.abort(404)
+            file_path = self.staged_files[file_token]
+            self.staged_files.pop(file_token, None)
+            return await quart.send_file(file_path)
 
     async def _set_callback_url(self):
         logger.info("设置回调，请等待...")
@@ -462,17 +492,18 @@ class SimpleGewechatClient:
                             "此次登录需要安全验证码，请在管理面板聊天页输入 /gewe_code 验证码 来验证，如 /gewe_code 123456"
                         )
                     else:
-                        status = json_blob["data"]["status"]
-                        nickname = json_blob["data"].get("nickName", "")
-                        if status == 1:
-                            logger.info(f"等待确认...{nickname}")
-                        elif status == 2:
-                            logger.info(f"绿泡泡平台登录成功: {nickname}")
-                            break
-                        elif status == 0:
-                            logger.info("等待扫码...")
-                        else:
-                            logger.warning(f"未知状态: {status}")
+                        if "status" in json_blob["data"]:
+                            status = json_blob["data"]["status"]
+                            nickname = json_blob["data"].get("nickName", "")
+                            if status == 1:
+                                logger.info(f"等待确认...{nickname}")
+                            elif status == 2:
+                                logger.info(f"绿泡泡平台登录成功: {nickname}")
+                                break
+                            elif status == 0:
+                                logger.info("等待扫码...")
+                            else:
+                                logger.warning(f"未知状态: {status}")
             await asyncio.sleep(5)
 
         if appid:

astrbot/core/platform/sources/gewechat/gewechat_event.py

@@ -83,15 +83,9 @@ class GewechatPlatformEvent(AstrMessageEvent):
 
             elif isinstance(comp, Image):
                 img_path = await comp.convert_to_file_path()
-
-                # 检查 record_path 是否在 data/temp 目录中
-                temp_directory = os.path.abspath("data/temp")
-                if os.path.commonpath([temp_directory, img_path]) != temp_directory:
-                    with open(img_path, "rb") as f:
-                        img_path = save_temp_img(f.read())
-
-                file_id = os.path.basename(img_path)
-                img_url = f"{client.file_server_url}/{file_id}"
+                # 为了安全，向 AstrBot 回调服务注册可被 gewechat 访问的文件，并获得文件 token
+                token = await client._register_file(img_path)
+                img_url = f"{client.file_server_url}/{token}"
                 logger.debug(f"gewe callback img url: {img_url}")
                 await client.post_image(to_wxid, img_url)
             elif isinstance(comp, Video):
@@ -110,20 +104,29 @@ class GewechatPlatformEvent(AstrMessageEvent):
 
                     video_url = comp.file
                     # 根据 url 下载视频
-                    video_filename = f"{uuid.uuid4()}.mp4"
-                    video_path = f"data/temp/{video_filename}"
-                    await download_file(video_url, video_path)
+                    if video_url.startswith("http"):
+                        video_filename = f"{uuid.uuid4()}.mp4"
+                        video_path = f"data/temp/{video_filename}"
+                        await download_file(video_url, video_path)
+                    else:
+                        video_path = video_url
+
+                    video_token = await client._register_file(video_path)
+                    video_callback_url = f"{client.file_server_url}/{video_token}"
 
                     # 获取视频第一帧
-                    thumb_path = f"data/temp/{uuid.uuid4()}.jpg"
+                    thumb_path = f"data/temp/gewechat_video_thumb_{uuid.uuid4()}.jpg"
+
+                    video_path = video_path.replace(" ", "\\ ")
                     try:
                         ff = FFmpeg()
-                        command = f'-i "{video_path}" -ss 0 -vframes 1 "{thumb_path}"'
+                        command = f"-i {video_path} -ss 0 -vframes 1 {thumb_path}"
                         ff.options(command)
-                        thumb_file_id = os.path.basename(thumb_path)
-                        thumb_url = f"{client.file_server_url}/{thumb_file_id}"
+                        thumb_token = await client._register_file(thumb_path)
+                        thumb_url = f"{client.file_server_url}/{thumb_token}"
                     except Exception as e:
                         logger.error(f"获取视频第一帧失败: {e}")
+
                     # 获取视频时长
                     try:
                         from pyffmpeg import FFprobe
@@ -138,15 +141,12 @@ class GewechatPlatformEvent(AstrMessageEvent):
                         logger.error(f"获取时长失败: {e}")
                         video_duration = 10
 
-                    file_id = os.path.basename(video_path)
-                    video_url = f"{client.file_server_url}/{file_id}"
+                    # 发送视频
                     await client.post_video(
-                        to_wxid, video_url, thumb_url, video_duration
+                        to_wxid, video_callback_url, thumb_url, video_duration
                     )
 
-                    # 删除临时视频和缩略图文件
-                    if os.path.exists(video_path):
-                        os.remove(video_path)
+                    # 删除临时缩略图文件
                     if os.path.exists(thumb_path):
                         os.remove(thumb_path)
             elif isinstance(comp, Record):
@@ -163,8 +163,8 @@ class GewechatPlatformEvent(AstrMessageEvent):
                 logger.info("Silk 语音文件格式转换至: " + record_path)
                 if duration == 0:
                     duration = get_wav_duration(record_path)
-                file_id = os.path.basename(silk_path)
-                record_url = f"{client.file_server_url}/{file_id}"
+                token = await client._register_file(silk_path)
+                record_url = f"{client.file_server_url}/{token}"
                 logger.debug(f"gewe callback record url: {record_url}")
                 await client.post_voice(to_wxid, record_url, duration * 1000)
             elif isinstance(comp, File):
@@ -177,10 +177,10 @@ class GewechatPlatformEvent(AstrMessageEvent):
                 else:
                     file_path = file_path
 
-                file_id = os.path.basename(file_path)
-                file_url = f"{client.file_server_url}/{file_id}"
+                token = await client._register_file(file_path)
+                file_url = f"{client.file_server_url}/{token}"
                 logger.debug(f"gewe callback file url: {file_url}")
-                await client.post_file(to_wxid, file_url, file_id)
+                await client.post_file(to_wxid, file_url, file_name)
             elif isinstance(comp, Emoji):
                 await client.post_emoji(to_wxid, comp.md5, comp.md5_len, comp.cdnurl)
             elif isinstance(comp, At):