Enhance client

python/wcf/client.py

@@ -5,10 +5,11 @@ import atexit
 import ctypes
 import logging
 import os
+import re
 import sys
 from threading import Thread
 from time import sleep
-from typing import Any, Callable, Optional
+from typing import Any, List, Callable, Optional
 
 import grpc
 
@@ -19,6 +20,42 @@ import wcf_pb2_grpc  # noqa
 
 
 class Wcf():
+    """WeChatFerry, a tool to play WeChat."""
+    class WxMsg():
+        """微信消息"""
+
+        def __init__(self, msg: wcf_pb2.WxMsg) -> None:
+            self._is_self = msg.is_self
+            self._is_group = msg.is_group
+            self.type = msg.type
+            self.id = msg.id
+            self.xml = msg.xml
+            self.sender = msg.sender
+            self.roomid = msg.roomid
+            self.content = msg.content
+
+        def __str__(self) -> str:
+            s = f"{self.sender}[{self.roomid}]\t{self.id}-{self.type}-{self.xml.replace(chr(10), '').replace(chr(9),'')}\n"
+            s += self.content
+            return s
+
+        def from_self(self) -> bool:
+            """是否自己发的消息"""
+            return self._is_self == 1
+
+        def from_group(self) -> bool:
+            """是否群聊消息"""
+            return self._is_group
+
+        def is_at(self, wxid) -> bool:
+            """是否被@：群消息，在@名单里，并且不是@所有人"""
+            return self.from_group() and re.findall(
+                f"<atuserlist>.*({wxid}).*</atuserlist>", self.xml) and not re.findall(r"@(?:所有人|all)", self.xml)
+
+        def is_text(self) -> bool:
+            """是否文本消息"""
+            return self.type == 1
+
     def __init__(self, host_port: str = "localhost:10086") -> None:
         self._enable_recv_msg = False
         self.LOG = logging.getLogger("WCF")
@@ -31,11 +68,15 @@ class Wcf():
         self._stub = wcf_pb2_grpc.WcfStub(self._channel)
         atexit.register(self.disable_recv_msg)  # 退出的时候停止消息接收，防止内存泄露
         self._is_running = True
+        self.contacts = []
+        self._SQL_TYPES = {1: int, 2: float, 3: lambda x: x.decode("utf-8"), 4: bytes, 5: lambda x: None}
+        self.self_wxid = self.get_self_wxid()
 
     def __del__(self) -> None:
         self.cleanup()
 
     def cleanup(self) -> None:
+        """停止 gRPC，关闭连接，回收资源"""
         if not self._is_running:
             return
 
@@ -48,6 +89,7 @@ class Wcf():
         self._is_running = False
 
     def keep_running(self):
+        """阻塞进程，让 RPC 一直维持连接"""
         try:
             while True:
                 sleep(1)
@@ -55,10 +97,12 @@ class Wcf():
             self.cleanup()
 
     def is_login(self) -> int:
+        """是否已经登录"""
         rsp = self._stub.RpcIsLogin(wcf_pb2.Empty())
         return rsp.status
 
     def get_self_wxid(self) -> str:
+        """获取登录账户的 wxid"""
         rsp = self._stub.RpcGetSelfWxid(wcf_pb2.Empty())
         return rsp.str
 
@@ -66,13 +110,14 @@ class Wcf():
         rsps = self._stub.RpcEnableRecvMsg(wcf_pb2.Empty())
         try:
             for rsp in rsps:
-                func(rsp)
+                func(self.WxMsg(rsp))
         except Exception as e:
             self.LOG.error(f"RpcEnableRecvMsg: {e}")
         finally:
             self.disable_recv_msg()
 
-    def enable_recv_msg(self, callback: Callable[..., Any] = None) -> bool:
+    def enable_recv_msg(self, callback: Callable[[WxMsg], None] = None) -> bool:
+        """设置接收消息回调"""
         if self._enable_recv_msg:
             return True
 
@@ -89,6 +134,7 @@ class Wcf():
         return True
 
     def disable_recv_msg(self) -> int:
+        """停止接收消息"""
         if not self._enable_recv_msg:
             return -1
 
@@ -99,33 +145,58 @@ class Wcf():
         return rsp.status
 
     def send_text(self, msg: str, receiver: str, aters: Optional[str] = "") -> int:
+        """发送文本消息"""
         rsp = self._stub.RpcSendTextMsg(wcf_pb2.TextMsg(msg=msg, receiver=receiver, aters=aters))
         return rsp.status
 
     def send_image(self, path: str, receiver: str) -> int:
+        """发送图片"""
         rsp = self._stub.RpcSendImageMsg(wcf_pb2.ImageMsg(path=path, receiver=receiver))
         return rsp.status
 
-    def get_msg_types(self) -> wcf_pb2.MsgTypes:
+    def get_msg_types(self) -> dict:
+        """获取所有消息类型"""
         rsp = self._stub.RpcGetMsgTypes(wcf_pb2.Empty())
-        return rsp
+        return dict(sorted(dict(rsp.types).items()))
 
-    def get_contacts(self) -> wcf_pb2.Contacts:
+    def get_contacts(self) -> List[dict]:
+        """获取完整通讯录"""
         rsp = self._stub.RpcGetContacts(wcf_pb2.Empty())
-        return rsp
+        for cnt in rsp.contacts:
+            gender = ""
+            if cnt.gender == 1:
+                gender = "男"
+            elif cnt.gender == 2:
+                gender = "女"
+            self.contacts.append({"wxid": cnt.wxid, "code": cnt.code, "name": cnt.name,
+                                 "country": cnt.country, "province": cnt.province, "city": cnt.city, "gender": gender})
+        return self.contacts
 
-    def get_dbs(self) -> wcf_pb2.DbNames:
+    def get_dbs(self) -> List[str]:
+        """获取所有数据库"""
         rsp = self._stub.RpcGetDbNames(wcf_pb2.Empty())
-        return rsp
+        return rsp.names
 
-    def get_tables(self, db: str) -> wcf_pb2.DbTables:
+    def get_tables(self, db: str) -> List[dict]:
+        """获取 db 中所有表"""
+        tables = []
         rsp = self._stub.RpcGetDbTables(wcf_pb2.String(str=db))
-        return rsp
+        for tbl in rsp.tables:
+            tables.append({"name": tbl.name, "sql": tbl.sql.replace("\t", "")})
+        return tables
 
-    def query_sql(self, db: str, sql: str) -> wcf_pb2.DbRows:
+    def query_sql(self, db: str, sql: str) -> List[dict]:
+        """执行 SQL"""
+        result = []
         rsp = self._stub.RpcExecDbQuery(wcf_pb2.DbQuery(db=db, sql=sql))
-        return rsp
+        for r in rsp.rows:
+            row = {}
+            for f in r.fields:
+                row[f.column] = self._SQL_TYPES[f.type](f.content)
+            result.append(row)
+        return result
 
     def accept_new_friend(self, v3: str, v4: str) -> int:
+        """通过好友验证"""
         rsp = self._stub.RpcAcceptNewFriend(wcf_pb2.Verification(v3=v3, v4=v4))
         return rsp.status