✨图片消息进行解密

Author: JustUndertaker

.env

@@ -34,3 +34,8 @@ cache_path = "./file_cache"
 # 文件缓存天数
 cache_days = 3
 
+# 聊天图片解密地址
+image_path = "./image_decode"
+
+# 下载pc图片超时时间(s)，超时的图片消息将不会发送
+image_timeout = 30

.gitignore

@@ -130,3 +130,4 @@ dmypy.json
 *.code-workspace
 file_cache/
 logs/
+image_decode/

ntchat_client/config.py

@@ -158,6 +158,10 @@ class Config(BaseConfig):
     """文件缓存目录"""
     cache_days: int = 3
     """文件缓存天数"""
+    image_path: str = "./image_decode"
+    """聊天图片解密地址"""
+    image_timeout: int = 30
+    """下载pc图片超时时间(s)，超时的图片消息将不会发送"""
 
     class Config:
         extra = "allow"

ntchat_client/wechat/image_decode.py

@@ -0,0 +1,71 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+
+@dataclass
+class FileTypes:
+    """文件格式"""
+
+    file_type: str
+    """格式后缀"""
+    key: int
+    """密钥"""
+
+
+class FileDecoder:
+    """文件解密类"""
+
+    file_map: dict[str, tuple[int, int]] = {
+        "jpg": (0xFF, 0xD8),
+        "png": (0x89, 0x50),
+        "gif": (0x47, 0x49),
+    }
+    out_image_dir: Path
+    """输出文件夹"""
+    out_image_thumb: Path
+    """缩略图输出文件夹"""
+
+    def __init__(self, image_path: str) -> None:
+        self.out_image_dir = Path(image_path) / "image"
+        self.out_image_thumb = Path(image_path) / "thumb"
+        self.out_image_dir.mkdir(parents=True, exist_ok=True)
+        self.out_image_thumb.mkdir(parents=True, exist_ok=True)
+
+    def get_file_type(self, byte0: int, byte1: int) -> Optional[FileTypes]:
+        """获取文件格式及密钥"""
+        for type, value in self.file_map.items():
+            result0 = value[0] ^ byte0
+            result1 = value[1] ^ byte1
+            if result0 == result1:
+                return FileTypes(type, result0)
+        return None
+
+    def decode_file(self, image_file: Path, is_thumb: bool) -> Optional[str]:
+        """
+        说明:
+            解密微信图片文件，并返回新的文件地址
+
+        参数:
+            * `image_file`：dat文件路径
+            * `is_thumb`：是否为缩略图
+
+        返回:
+            * `str`：解密文件路径
+        """
+        file_value = np.fromfile(image_file, dtype=np.uint8)
+        file_type = self.get_file_type(file_value[0], file_value[1])
+        if file_type is None:
+            return None
+        xor_array = np.full_like(file_value, fill_value=file_type.key)
+        out_value = np.bitwise_xor(file_value, xor_array)
+        if is_thumb:
+            out_file = self.out_image_thumb / f"{image_file.stem}.{file_type.file_type}"
+        else:
+            out_file = self.out_image_dir / f"{image_file.stem}.{file_type.file_type}"
+        with open(out_file, mode="wb") as f:
+            f.write(out_value)
+
+        return out_file.absolute()

ntchat_client/wechat/wechat.py

@@ -1,5 +1,7 @@
 import asyncio
+import time
 from asyncio import AbstractEventLoop
+from pathlib import Path
 from typing import Any, Callable, NoReturn, Optional
 
 import ntchat
@@ -17,6 +19,7 @@
 from ntchat_client.utils import escape_tag, notify
 
 from .cache import FileCache
+from .image_decode import FileDecoder
 from .qrcode import draw_qrcode
 
 wechat_client: "WeChatManager" = None
@@ -69,6 +72,10 @@ class WeChatManager:
     """http_post处理器"""
     file_cache: FileCache
     """文件缓存管理器"""
+    image_decoder: FileDecoder
+    """图片解密器"""
+    image_timeout: int
+    """图片下载超时时间"""
     msg_fiter = {
         ntchat.MT_USER_LOGIN_MSG,
         ntchat.MT_USER_LOGOUT_MSG,
@@ -88,6 +95,8 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, config: Config) -> None:
         self.config = config
         self.file_cache = FileCache(config.cache_path)
+        self.image_decoder = FileDecoder(config.image_path)
+        self.image_timeout = config.image_timeout
         self.msg_fiter |= config.msg_filter
         ntchat.set_wechat_exe_path(wechat_version="3.6.0.18")
 
@@ -208,6 +217,23 @@ def handle_ws_api(self, request: WsRequest) -> WsRequest:
             echo=echo, status=response.status, msg=response.msg, data=response.data
         )
 
+    def _handle_image(self, message: dict) -> Optional[dict]:
+        """处理图片消息，并替换字段，如果超时将会返回None"""
+        data: dict = message["data"]
+        image_file = Path(data["image"])
+        image_thumb = Path(data["image_thumb"])
+        time_count = 0
+        while True:
+            if image_file.exists() and image_thumb.exists():
+                break
+            time_count += 1
+            if time_count > self.image_timeout:
+                return None
+            time.sleep(1)
+        data["image"] = self.image_decoder.decode_file(image_file, False)
+        data["image_thumb"] = self.image_decoder.decode_file(image_thumb, True)
+        return message
+
     def on_message(self, _: ntchat.WeChat, message: dict) -> None:
         """接收消息"""
         # 过滤事件
@@ -218,6 +244,15 @@ def on_message(self, _: ntchat.WeChat, message: dict) -> None:
         if wx_id == self.self_id and not self.config.report_self:
             return
         logger.success(f"<m>wechat</m> - <g>收到wechat消息：</g>{escape_tag(str(message))}")
+        if msgtype == 11047:
+            # 群图片消息，预处理
+            logger.debug("正在解密图片地址...")
+            message = self._handle_image(message)
+            if message is None:
+                logger.error("下载图片超时，本次消息不会发送...")
+                return
+            logger.debug("解密图片已保存...")
+
         if self.loop is not None:
             if self.loop.is_running:
                 if self.ws_message_handler: