From 55d90b40bd4619049506b5c88b9759fc909c3b1a Mon Sep 17 00:00:00 2001 From: mattuy Date: Sat, 7 Aug 2021 14:56:07 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=8E=BB=E9=99=A4sqlite=E5=8A=A0?= =?UTF-8?q?=E5=AF=86=E8=83=BD=E5=8A=9B=E4=BB=A5=E9=80=82=E9=85=8D=E6=A0=91?= =?UTF-8?q?=E8=8E=93=E6=B4=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 257 ++++++++-------------------------------- message_queue.py | 44 +++++++ requirements.txt | 1 + script/install-deps.sh | 9 ++ wxbot.py | 258 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 358 insertions(+), 211 deletions(-) create mode 100644 message_queue.py create mode 100755 script/install-deps.sh create mode 100644 wxbot.py diff --git a/main.py b/main.py index d5052c7..f31c4ea 100644 --- a/main.py +++ b/main.py @@ -3,232 +3,60 @@ import argparse import getpass import logging -import sys -import time import os -from weixin import WebWeixin -from sqlcipher3 import dbapi2 as sqlite +import queue +import websockets +import sys +import threading +import asyncio +import sqlite3 +from wxbot import WXBot +from message_queue import wxbot_instance, message_consumer + + +class WXBotThreadStarter(threading.Thread): + def __init__(self, wxbot): + threading.Thread.__init__(self) + self.wxbot = wxbot + + def run(self): + self.wxbot.start() def parse_args(): parser = argparse.ArgumentParser() - parser.add_argument('-f', '--db-file', required=True, help='数据库文件') - parser.add_argument('-d', '--data-dir',required=True, help='文件存储目录') - parser.add_argument('--debug', action='store_true', help='调试模式') + parser.add_argument("-f", "--db-file", required=True, help="数据库文件") + parser.add_argument("-d", "--data-dir", required=True, help="文件存储目录") + parser.add_argument("-s", "--serve", action="store_true", help="开启HTTP服务") + parser.add_argument("--debug", action="store_true", help="调试模式") args = parser.parse_args() return args -class WXBot(WebWeixin): - def getUserNickName(self, id): - if id == self.User['UserName']: - return self.User['NickName'] - for member in self.ContactList: - if member['UserName'] == id: - return member['NickName'] - return None - # 获取用户备注（没有则返回昵称）或群名称，若本地无相关数据返回_unknown - def getUserOrGroupName(self, id): - name = '_unknown' - if id == self.User['UserName']: - return self.User['NickName'] # 自己 - - if id[:2] == '@@': - # 群 - name = self.getGroupName(id) - else: - # 特殊账号 - for member in self.SpecialUsersList: - if member['UserName'] == id: - name = member['RemarkName'] if member[ - 'RemarkName'] else member['NickName'] - # 公众号或服务号 - for member in self.PublicUsersList: - if member['UserName'] == id: - name = member['RemarkName'] if member[ - 'RemarkName'] else member['NickName'] - # 直接联系人 - for member in self.ContactList: - if member['UserName'] == id: - name = member['RemarkName'] if member[ - 'RemarkName'] else member['NickName'] - # 群友 - for member in self.GroupMemeberList: - if member['UserName'] == id: - name = member['DisplayName'] if member[ - 'DisplayName'] else member['NickName'] - return name - - def composeMessageID(self, msgID, timestamp): - return time.strftime('%Y%m%d%H%M%S', time.localtime(timestamp)) + msgID - - def saveMessageImage(self, msgid, composedID): - url = self.base_uri + \ - '/webwxgetmsgimg?MsgID=%s&skey=%s' % (msgid, self.skey) - data = self._get(url, 'webwxgetmsgimg') - return self.saveFile(composedID[:4], composedID + '.jpg', data) - - def saveVoice(self, msgid, composedID): - url = self.base_uri + \ - '/webwxgetvoice?msgid=%s&skey=%s' % (msgid, self.skey) - data = self._get(url, api='webwxgetvoice') - return self.saveFile(composedID[:4], composedID + '.mp3', data) - - def saveVideo(self, msgid, composedID): - url = self.base_uri + \ - '/webwxgetvideo?msgid=%s&skey=%s' % (msgid, self.skey) - data = self._get(url, api='webwxgetvideo') - return self.saveFile(composedID[:4], composedID + '.mp4', data) - - def saveFile(self, subdir, filename, data): - if data == '': - return False - dirName = os.path.join(dataDir, subdir) - if not os.path.exists(dirName): - os.makedirs(dirName) - fn = os.path.join(dirName, filename) - with open(fn, 'wb') as f: - f.write(data) - f.close() - return True - - def handleMsg(self, r): - if self.DEBUG: - print("处理消息, NewMessageCount =", r['AddMsgCount']) - for msg in r['AddMsgList']: - msgID = msg['MsgId'] - composedID = self.composeMessageID(msgID, msg['CreateTime']) - msgType = msg['MsgType'] - if msg['FromUserName'] in self.SpecialUsers or msg['ToUserName'] in self.SpecialUsers: - # 过滤特殊帐号消息 - continue - elif msgType in (50, 51, 52, 53, 9999): - #无用消息 - continue - sql = "INSERT INTO msg (\ - msg_id,\ - create_time,\ - msg_type,\ - content,\ - from_name,\ - from_nickname,\ - to_name,\ - to_nickname,\ - group_name\ - ) VALUES (\ - :msg_id,\ - :create_time,\ - :msg_type,\ - :content,\ - :from_name,\ - :from_nickname,\ - :to_name,\ - :to_nickname,\ - :group_name\ - )" - params = { - "msg_id": composedID, - "create_time": msg['CreateTime'], - "msg_type": msgType, - "content": None, - "from_name": self.getUserOrGroupName(msg['FromUserName']), - "from_nickname": self.getUserNickName(msg['FromUserName']), - "to_name": self.getUserOrGroupName(msg['ToUserName']), - "to_nickname": self.getUserNickName(msg['ToUserName']), - "group_name": None - } - if msg['FromUserName'] == self.User['UserName']: - params['from_name'] = '_self' - elif msg['ToUserName'] == self.User['UserName']: - params['to_name'] = '_self' - if msg['FromUserName'][:2] == '@@': - # 群消息 - if ":
" in msg['Content']: - [people, content] = msg['Content'].split(':
', 1) - params['group_name'] = self.getUserOrGroupName(msg['FromUserName']) - params['from_name'] = self.getUserOrGroupName(people) - params['to_name'] = params['to_nickname'] = None - elif msg['ToUserName'][:2] == '@@': - # 发送的群消息 - params['group_name'] = self.getUserOrGroupName(msg['ToUserName']) - params['from_name'] = self.User['NickName'] - params['to_name'] = params['to_nickname'] = None - if msgType == 1: - # 文字消息 - if params['from_name'] == '_unknown' or params['to_name'] == '_unknown': - # 有陌生人的文字消息，可能是新添加的好友，重新获取联系人列表 - self.webwxgetcontact() - # 重新获取用户名称 - params['from_name'] = self.getUserOrGroupName( - msg['FromUserName']) - params['to_name'] = self.getUserOrGroupName( - msg['ToUserName']) - params['content'] = msg['Content'].replace('<', '<').replace('>', '>') - elif msgType == 3: - params['content'] = "[图片](%s)" % composedID - self.saveMessageImage(msgID, composedID) - elif msgType == 34: - params['content'] = "[语音](%s)" % composedID - self.saveVoice(msgID, composedID) - elif msgType == 43: - params['content'] = "[视频](%s)" % composedID - self.saveVideo(msgID, composedID) - elif msgType == 62: - params['msg_type'] = 43 - params['content'] = "[小视频](%s)" % composedID - self.saveVideo(msgID, composedID) - elif msgType == 42: - params['msg_type'] = 0 - params['content'] = "[名片](%s)" % msg['RecommendInfo']['NickName'] - elif msgType == 47: - if msg['HasProductId'] == 1: - # 有版权的表情包，无法在Web端查看 - params['msg_type'] = 0 - params['content'] = "[表情包图片]" - else: - params['msg_type'] = 3 - params['content'] = "[动画表情](%s)" % composedID - self.saveMessageImage(msgID, composedID) - elif msgType == 49: - params['msg_type'] = 0 - params['content'] = "[链接：%s](%s)" % (msg['FileName'], msg['Url']) - elif msgType == 10000: - params['msg_type'] = 0 - params['content'] = "[系统消息](%s)" % (msg['Content']) - elif msgType == 10002: - params['msg_type'] = 0 - params['content'] = "[系统消息](%s撤回了一条消息)" % params['from_name'] - else: - params['msg_type'] = 0 - params['content'] = "[未识别消息]" + msg['Content'] - conn.execute(sql, params) - conn.commit() - - -if __name__ == '__main__': - if sys.stdout.encoding == 'cp936': +if __name__ == "__main__": + if sys.stdout.encoding == "cp936": sys.stdout = UnicodeStreamFilter(sys.stdout) logger = logging.getLogger(__name__) - if not sys.platform.startswith('win'): + if not sys.platform.startswith("win"): import coloredlogs - coloredlogs.install(level='INFO') + + coloredlogs.install(level="INFO") args = parse_args() - dataDir = args.data_dir - if not os.access(os.path.dirname(args.db_file), os.X_OK) and not os.access(args.db_file, os.W_OK): + if not os.access(os.path.dirname(args.db_file), os.X_OK) and not os.access( + args.db_file, os.W_OK + ): print("无法创建或打开数据库文件") exit() - elif not os.access(dataDir, os.X_OK): + elif not os.access(args.data_dir, os.X_OK): print("数据目录不存在或权限不足") exit() - conn = sqlite.connect(args.db_file) + conn = sqlite3.connect(args.db_file) while True: try: - passwd = getpass.getpass("输入数据库密码：") - if passwd != '': - conn.execute("PRAGMA KEY = %s" % passwd) - conn.execute("CREATE TABLE IF NOT EXISTS msg (\ + conn.execute( + "CREATE TABLE IF NOT EXISTS msg (\ msg_id TEXT NOT NULL,\ create_time INTEGER NOT NULL,\ msg_type INTEGER NOT NULL,\ @@ -238,10 +66,11 @@ if __name__ == '__main__': to_name TEXT,\ to_nickname TEXT,\ group_name TEXT\ - );") + );" + ) break - except sqlite.DatabaseError as e: - if 'file is not a database' in e.args: + except sqlite3.DatabaseError as e: + if "file is not a database" in e.args: print("密码错误，请重试") else: print("打开数据库失败", e) @@ -250,7 +79,13 @@ if __name__ == '__main__': except Exception as e: print("打开数据库失败", e) exit(1) - wxbot = WXBot() + msg_queue = queue.Queue(1024) + wxbot = WXBot(conn) + wxbot_instance['current'] = wxbot wxbot.DEBUG = args.debug - wxbot.saveFolder = dataDir - wxbot.start() + wxbot.saveFolder = args.data_dir + WXBotThreadStarter(wxbot).start() + if args.serve: + start_server = websockets.serve(message_consumer, "localhost", 8083) + asyncio.get_event_loop().run_until_complete(start_server) + asyncio.get_event_loop().run_forever() diff --git a/message_queue.py b/message_queue.py new file mode 100644 index 0000000..52afc4b --- /dev/null +++ b/message_queue.py @@ -0,0 +1,44 @@ +import queue +import time +import asyncio +import websockets +import json + + +async def message_consumer(websocket, path): + subscription = wxbot_instance.subscribe() + while not websocket.closed: + try: + message = subscription.queue.get_nowait() + subscription.lastRead = time.time() + if message == None: + await ws.close(code=1000) + break + await websocket.send(json.dumps(message)) + except queue.Empty: + await asyncio.sleep(1) + continue + except WebSocketDisconnected: + wxbot_instance.unsubscribe(subscription) + subscription = None + return + except: + wxbot_instance.unsubscribe(subscription) + subscription = None + await ws.close(code=1001) + if subscription != None: + wxbot_instance.unsubscribe() + + +wxbot_instance = {"current": None} + + +class MessageTransfer: + def __init__(self, msg_queue): + message_queue = msg_queue + self.msg_queue = msg_queue + + def consume_message(): + msg = msg_queue.get() + for session in sessions: + pass diff --git a/requirements.txt b/requirements.txt index ed0735b..56bd657 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,4 @@ pyqrcode certifi argparse sqlcipher3-binary +websockets \ No newline at end of file diff --git a/script/install-deps.sh b/script/install-deps.sh new file mode 100755 index 0000000..a59ee5b --- /dev/null +++ b/script/install-deps.sh @@ -0,0 +1,9 @@ +#!/bin/bash +if ! (cd "$(dirname "$0")" && cd ..); then + echo 进入项目根目录失败 + exit 1 +fi + +python3 -m venv ./venv +source ./venv/bin/activate +pip install -r requirements.txt diff --git a/wxbot.py b/wxbot.py new file mode 100644 index 0000000..181306a --- /dev/null +++ b/wxbot.py @@ -0,0 +1,258 @@ +from weixin import WebWeixin +import sys +import time +import queue +import logging +import os + + +class Subscription: + def __init__(self, queue): + self.lastRead = time.time() + self.queue = queue + + +# 扩展一下将数据存到数据库 +class WXBot(WebWeixin): + def __init__(self, db_conn): + super(WXBot, self).__init__() + self.db_conn = db_conn + self.subscriptions = [] # type: List[Subscription] + + def getUserNickName(self, id): + if id == self.User["UserName"]: + return self.User["NickName"] + for member in self.ContactList: + if member["UserName"] == id: + return member["NickName"] + return None + + # 获取用户备注（没有则返回昵称）或群名称，若本地无相关数据返回_unknown + def getUserOrGroupName(self, id): + name = "_unknown" + if id == self.User["UserName"]: + return self.User["NickName"] # 自己 + + if id[:2] == "@@": + # 群 + name = self.getGroupName(id) + else: + # 特殊账号 + for member in self.SpecialUsersList: + if member["UserName"] == id: + name = ( + member["RemarkName"] + if member["RemarkName"] + else member["NickName"] + ) + # 公众号或服务号 + for member in self.PublicUsersList: + if member["UserName"] == id: + name = ( + member["RemarkName"] + if member["RemarkName"] + else member["NickName"] + ) + # 直接联系人 + for member in self.ContactList: + if member["UserName"] == id: + name = ( + member["RemarkName"] + if member["RemarkName"] + else member["NickName"] + ) + # 群友 + for member in self.GroupMemeberList: + if member["UserName"] == id: + name = ( + member["DisplayName"] + if member["DisplayName"] + else member["NickName"] + ) + return name + + # 订阅新消息 + def subscribe(self): + print("new subscription") + sub = Subscription(queue.Queue(1024)) + self.subscriptions.append(sub) + return sub + + def unsubscribe(self, subscription): + self.subscriptions.remove(subscription) + + def composeMessageID(self, msgID, timestamp): + return time.strftime("%Y%m%d%H%M%S", time.localtime(timestamp)) + msgID + + def saveMessageImage(self, msgid, composedID): + url = self.base_uri + "/webwxgetmsgimg?MsgID=%s&skey=%s" % (msgid, self.skey) + data = self._get(url, "webwxgetmsgimg") + return self.saveFile(composedID[:4], composedID + ".jpg", data) + + def saveVoice(self, msgid, composedID): + url = self.base_uri + "/webwxgetvoice?msgid=%s&skey=%s" % (msgid, self.skey) + data = self._get(url, api="webwxgetvoice") + return self.saveFile(composedID[:4], composedID + ".mp3", data) + + def saveVideo(self, msgid, composedID): + url = self.base_uri + "/webwxgetvideo?msgid=%s&skey=%s" % (msgid, self.skey) + data = self._get(url, api="webwxgetvideo") + return self.saveFile(composedID[:4], composedID + ".mp4", data) + + def saveFile(self, subdir, filename, data): + if data == "": + return False + dirName = os.path.join(self.saveFolder, subdir) + if not os.path.exists(dirName): + os.makedirs(dirName) + fn = os.path.join(dirName, filename) + with open(fn, "wb") as f: + f.write(data) + f.close() + return True + + def handleMsg(self, r): + if self.DEBUG: + print("处理消息, NewMessageCount =", r["AddMsgCount"]) + for msg in r["AddMsgList"]: + msgID = msg["MsgId"] + composedID = self.composeMessageID(msgID, msg["CreateTime"]) + msgType = msg["MsgType"] + if ( + msg["FromUserName"] in self.SpecialUsers + or msg["ToUserName"] in self.SpecialUsers + ): + # 过滤特殊帐号消息 + continue + elif msgType in (50, 51, 52, 53, 9999): + # 无用消息 + continue + sql = "INSERT INTO msg (\ + msg_id,\ + create_time,\ + msg_type,\ + content,\ + from_name,\ + from_nickname,\ + to_name,\ + to_nickname,\ + group_name\ + ) VALUES (\ + :msg_id,\ + :create_time,\ + :msg_type,\ + :content,\ + :from_name,\ + :from_nickname,\ + :to_name,\ + :to_nickname,\ + :group_name\ + )" + + params = { + "msg_id": composedID, + "create_time": msg["CreateTime"], + "msg_type": msgType, + "content": None, + "from_id": msg["ToUserName"], + "from_name": self.getUserOrGroupName(msg["FromUserName"]), + "from_nickname": self.getUserNickName(msg["FromUserName"]), + "to_id": msg["ToUserName"], + "to_name": self.getUserOrGroupName(msg["ToUserName"]), + "to_nickname": self.getUserNickName(msg["ToUserName"]), + "group_id": None, + "group_name": None, + } + if msg["FromUserName"] == self.User["UserName"]: + params["from_name"] = "_self" + elif msg["ToUserName"] == self.User["UserName"]: + params["to_name"] = "_self" + if msg["FromUserName"][:2] == "@@": + # 群消息 + params["from_id"] = params["from_name"] = params["from_nickname"] = None + params["to_id"] = params["to_name"] = params["to_nickname"] = None + params["group_id"] = msg["FromUserName"] + params["group_name"] = self.getUserOrGroupName(msg["FromUserName"]) + if ":
" in msg["Content"]: + [people, content] = msg["Content"].split(":
", 1) + params["from_id"] = people + params["from_name"] = self.getUserOrGroupName(people) + elif msg["ToUserName"][:2] == "@@": + # 发送的群消息 + params["group_id"] = msg["ToUserName"] + params["group_name"] = self.getUserOrGroupName(msg["ToUserName"]) + params["from_id"] = self.User["UserName"] + params["from_name"] = self.User["NickName"] + params["to_id"] = params["to_name"] = params["to_nickname"] = None + if msgType == 1: + # 文字消息 + if params["from_name"] == "_unknown" or params["to_name"] == "_unknown": + # 有陌生人的文字消息，可能是新添加的好友，重新获取联系人列表 + self.webwxgetcontact() + # 重新获取用户名称 + params["from_name"] = self.getUserOrGroupName(msg["FromUserName"]) + params["to_name"] = self.getUserOrGroupName(msg["ToUserName"]) + params["content"] = ( + msg["Content"].replace("<", "<").replace(">", ">") + ) + elif msgType == 3: + params["content"] = "[图片](%s)" % composedID + self.saveMessageImage(msgID, composedID) + elif msgType == 34: + params["content"] = "[语音](%s)" % composedID + self.saveVoice(msgID, composedID) + elif msgType == 43: + params["content"] = "[视频](%s)" % composedID + self.saveVideo(msgID, composedID) + elif msgType == 62: + params["msg_type"] = 43 + params["content"] = "[小视频](%s)" % composedID + self.saveVideo(msgID, composedID) + elif msgType == 42: + params["msg_type"] = 0 + params["content"] = "[名片](%s)" % msg["RecommendInfo"]["NickName"] + elif msgType == 47: + if msg["HasProductId"] == 1: + # 有版权的表情包，无法在Web端查看 + params["msg_type"] = 0 + params["content"] = "[表情包图片]" + else: + params["msg_type"] = 3 + params["content"] = "[动画表情](%s)" % composedID + self.saveMessageImage(msgID, composedID) + elif msgType == 49: + params["msg_type"] = 0 + params["content"] = "[链接：%s](%s)" % (msg["FileName"], msg["Url"]) + elif msgType == 10000: + params["msg_type"] = 0 + params["content"] = "[系统消息](%s)" % (msg["Content"]) + elif msgType == 10002: + params["msg_type"] = 0 + params["content"] = "[系统消息](%s撤回了一条消息)" % params["from_name"] + else: + params["msg_type"] = 0 + params["content"] = "[未识别消息]" + msg["Content"] + self.db_conn.execute(sql, params) + self.db_conn.commit() + for subscription in self.subscriptions: + if ( + not subscription.queue.empty() + and time.time() - subscription.lastRead > 120 + ): + # 超时关闭 + subscription.queue.put_nowait(None) + self.subscriptions.remove(subscription) + # 将消息放入队列，推送到客户端 + sent_by_self = msg["FromUserName"] == self.User["UserName"] + message = { + "uid": params["to_id"] if sent_by_self else params["from_id"], + "content": params["content"], + "name": params["to_name"] if sent_by_self else params["from_name"], + "self": sent_by_self, + "type": params["msg_type"], + "gid": params["group_id"], + "group": params["group_name"], + } + if subscription.queue.full(): + subscription.queue.get_nowait() + subscription.queue.put_nowait(message)