From 29fbf699456105dc7701f1dae2e7f319d5a9f734 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E8=B6=85?= <lichao@ahxz.vip>
Date: Sun, 9 Apr 2023 11:59:31 +0800
Subject: [PATCH] feat: Add configuration items to support custom data
 directories and facilitate the storage of itchat.pkl

---
 channel/wechat/wechat_channel.py |  7 ++++---
 config.py                        | 14 +++++++++++---
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/channel/wechat/wechat_channel.py b/channel/wechat/wechat_channel.py
index b11611d..f2825f3 100644
--- a/channel/wechat/wechat_channel.py
+++ b/channel/wechat/wechat_channel.py
@@ -18,7 +18,7 @@ from lib import itchat
 from lib.itchat.content import *
 from bridge.reply import *
 from bridge.context import *
-from config import conf
+from config import conf, get_data_path
 from common.time_check import time_checker
 from common.expired_dict import ExpiredDict
 from plugins import *
@@ -97,13 +97,14 @@ class WechatChannel(ChatChannel):
         itchat.instance.receivingRetryCount = 600  # 修改断线超时时间
         # login by scan QRCode
         hotReload = conf().get('hot_reload', False)
+        status_path = os.path.join(get_data_path(), "itchat.pkl")
         try:
-            itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback)
+            itchat.auto_login(enableCmdQR=2, hotReload=hotReload, statusStorageDir=status_path, qrCallback=qrCallback)
         except Exception as e:
             if hotReload:
                 logger.error("Hot reload failed, try to login without hot reload")
                 itchat.logout()
-                os.remove("itchat.pkl")
+                os.remove(status_path)
                 itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback)
             else:
                 raise e
diff --git a/config.py b/config.py
index 8745bf9..7f4270d 100644
--- a/config.py
+++ b/config.py
@@ -93,6 +93,8 @@ available_setting = {
 
     "debug": False,  # 是否开启debug模式，开启后会打印更多日志
 
+    "config_data_path": "",    # 数据目录
+
     # 插件配置
     "plugin_trigger_prefix": "$",  # 规范插件提供聊天相关指令的前缀，建议不要和管理员指令前缀"#"冲突
 }
@@ -130,7 +132,7 @@ class Config(dict):
 
     def load_user_datas(self):
         try:
-            with open('user_datas.pkl', 'rb') as f:
+            with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'rb') as f:
                 self.user_datas = pickle.load(f)
                 logger.info("[Config] User datas loaded.")
         except FileNotFoundError as e:
@@ -141,7 +143,7 @@ class Config(dict):
 
     def save_user_datas(self):
         try:
-            with open('user_datas.pkl', 'wb') as f:
+            with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'wb') as f:
                 pickle.dump(self.user_datas, f)
                 logger.info("[Config] User datas saved.")
         except Exception as e:
@@ -196,6 +198,12 @@ def read_file(path):
     with open(path, mode='r', encoding='utf-8') as f:
         return f.read()
 
-
 def conf():
     return config
+
+def get_data_path():
+    data_path = os.path.join(get_root(), conf().get('config_data_path', ""))
+    if not os.path.exists(data_path):
+        logger.info("[INIT] data path not exists, create it: {}".format(data_path))
+        os.makedirs(data_path)
+    return data_path
\ No newline at end of file