From 51d1d2004bd2589533d1e19f65eba79db641eb86 Mon Sep 17 00:00:00 2001
From: Ding Jiatong <dingjiatong@gmail.com>
Date: Sun, 17 Nov 2024 16:43:41 +0800
Subject: [PATCH] fix: gewechat image download (#140)

---
 channel/gewechat/gewechat_channel.py |  2 +-
 channel/gewechat/gewechat_message.py | 27 +++++++++++++++++++++++----
 lib/gewechat/client.py               |  5 ++---
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/channel/gewechat/gewechat_channel.py b/channel/gewechat/gewechat_channel.py
index 513d75f37..232fb9681 100644
--- a/channel/gewechat/gewechat_channel.py
+++ b/channel/gewechat/gewechat_channel.py
@@ -30,7 +30,7 @@ def __init__(self):
                 self.base_url, self.download_url, self.token, self.app_id
             )
         )
-        self.client = GewechatClient(self.base_url, self.download_url, self.token)
+        self.client = GewechatClient(self.base_url, self.token)
 
     def startup(self):
         urls = ("/v2/api/callback/collect", "channel.gewechat.gewechat_channel.Query")
diff --git a/channel/gewechat/gewechat_message.py b/channel/gewechat/gewechat_message.py
index 57a0bdb02..78bbe32f5 100644
--- a/channel/gewechat/gewechat_message.py
+++ b/channel/gewechat/gewechat_message.py
@@ -6,6 +6,7 @@
 from common.tmp_dir import TmpDir
 from config import conf
 from lib.gewechat import GewechatClient
+import requests
 
 class GeWeChatMessage(ChatMessage):
     def __init__(self, msg, client: GewechatClient):
@@ -79,7 +80,7 @@ def __init__(self, msg, client: GewechatClient):
 
     def download_voice(self):
         try:
-            voice_data = self.client.download_file(self.msg['Wxid'], self.msg_id)
+            voice_data = self.client.download_voice(self.msg['Wxid'], self.msg_id)
             with open(self.content, "wb") as f:
                 f.write(voice_data)
         except Exception as e:
@@ -87,9 +88,27 @@ def download_voice(self):
 
     def download_image(self):
         try:
-            image_data = self.client.download_file(self.msg['Wxid'], self.msg_id)
-            with open(self.content, "wb") as f:
-                f.write(image_data)
+            try:
+                # 尝试下载高清图片
+                image_info = self.client.download_image(app_id=self.app_id, xml=self.msg['Data']['Content']['string'], type=1)
+            except Exception as e:
+                logger.warning(f"[gewechat] Failed to download high-quality image: {e}")
+                # 尝试下载普通图片
+                image_info = self.client.download_image(app_id=self.app_id, xml=self.msg['Data']['Content']['string'], type=2)
+            if image_info['ret'] == 200 and image_info['data']:
+                file_url = image_info['data']['fileUrl']
+                logger.info(f"[gewechat] Download image file from {file_url}")
+                download_url = conf().get("gewechat_download_url").rstrip('/')
+                full_url = download_url + '/' + file_url
+                try:
+                    file_data = requests.get(full_url).content
+                except Exception as e:
+                    logger.error(f"[gewechat] Failed to download image file: {e}")
+                    return
+                with open(self.content, "wb") as f:
+                    f.write(file_data)
+            else:
+                logger.error(f"[gewechat] Failed to download image file: {image_info}")
         except Exception as e:
             logger.error(f"[gewechat] Failed to download image file: {e}")
 
diff --git a/lib/gewechat/client.py b/lib/gewechat/client.py
index f7445c428..09b54c73f 100644
--- a/lib/gewechat/client.py
+++ b/lib/gewechat/client.py
@@ -1,6 +1,5 @@
 from .api.contact_api import ContactApi
 from .api.download_api import DownloadApi
-from .api.download_api import DownloadApi
 from .api.favor_api import FavorApi
 from .api.group_api import GroupApi
 from .api.label_api import LabelApi
@@ -30,9 +29,9 @@ class GewechatClient:
 
     注意: 在使用任何方法之前，请确保你已经正确初始化了客户端，并且有有效的 base_url 和 token。
     """
-    def __init__(self, base_url, download_url, token):
+    def __init__(self, base_url, token):
         self._contact_api = ContactApi(base_url, token)
-        self._download_api = DownloadApi(download_url, token)
+        self._download_api = DownloadApi(base_url, token)
         self._favor_api = FavorApi(base_url, token)
         self._group_api = GroupApi(base_url, token)
         self._label_api = LabelApi(base_url, token)