92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
# coding: utf-8
|
|
import argparse
|
|
import getpass
|
|
import logging
|
|
import os
|
|
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("-s", "--serve", action="store_true", help="开启HTTP服务")
|
|
parser.add_argument("--debug", action="store_true", help="调试模式")
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if sys.stdout.encoding == "cp936":
|
|
sys.stdout = UnicodeStreamFilter(sys.stdout)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
if not sys.platform.startswith("win"):
|
|
import coloredlogs
|
|
|
|
coloredlogs.install(level="INFO")
|
|
|
|
args = parse_args()
|
|
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(args.data_dir, os.X_OK):
|
|
print("数据目录不存在或权限不足")
|
|
exit()
|
|
conn = sqlite3.connect(args.db_file)
|
|
while True:
|
|
try:
|
|
conn.execute(
|
|
"CREATE TABLE IF NOT EXISTS msg (\
|
|
msg_id TEXT NOT NULL,\
|
|
create_time INTEGER NOT NULL,\
|
|
msg_type INTEGER NOT NULL,\
|
|
content TEXT,\
|
|
from_name TEXT,\
|
|
from_nickname TEXT,\
|
|
to_name TEXT,\
|
|
to_nickname TEXT,\
|
|
group_name TEXT\
|
|
);"
|
|
)
|
|
break
|
|
except sqlite3.DatabaseError as e:
|
|
if "file is not a database" in e.args:
|
|
print("密码错误，请重试")
|
|
else:
|
|
print("打开数据库失败", e)
|
|
except KeyboardInterrupt:
|
|
exit(1)
|
|
except Exception as e:
|
|
print("打开数据库失败", e)
|
|
exit(1)
|
|
msg_queue = queue.Queue(1024)
|
|
wxbot = WXBot(conn)
|
|
wxbot_instance['current'] = wxbot
|
|
wxbot.DEBUG = args.debug
|
|
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()
|