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)