Merge branch 'main' of https://github.com/SengokuCola/MaiMBot

src/common/database.py

@@ -4,18 +4,24 @@ from typing import Optional
 class Database:
     _instance: Optional["Database"] = None
     
-    def __init__(self, host: str, port: int, db_name: str):
-        self.client = MongoClient(host, port)
+    def __init__(self, host: str, port: int, db_name: str, username: Optional[str] = None, password: Optional[str] = None, auth_source: Optional[str] = None):
+        if username and password:
+            # 如果有用户名和密码，使用认证连接
+            # TODO: 复杂情况直接支持URI吧
+            self.client = MongoClient(host, port, username=username, password=password, authSource=auth_source)
+        else:
+            # 否则使用无认证连接
+            self.client = MongoClient(host, port)
         self.db = self.client[db_name]
         
     @classmethod
-    def initialize(cls, host: str, port: int, db_name: str) -> "Database":
+    def initialize(cls, host: str, port: int, db_name: str, username: Optional[str] = None, password: Optional[str] = None, auth_source: Optional[str] = None) -> "Database":
         if cls._instance is None:
-            cls._instance = cls(host, port, db_name)
+            cls._instance = cls(host, port, db_name, username, password, auth_source)
         return cls._instance
         
     @classmethod
     def get_instance(cls) -> "Database":
         if cls._instance is None:
             raise RuntimeError("Database not initialized")
-        return cls._instance 
+        return cls._instance

src/plugins/chat/bot_config_toml

@@ -2,6 +2,9 @@
 host = "127.0.0.1"
 port = 27017
 name = "MegBot"
+username = ""  # 默认空值
+password = ""  # 默认空值
+auth_source = ""  # 默认空值
 
 [bot]
 qq = #填入你的机器人QQ

src/plugins/chat/config.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Set
 import os
 from nonebot.log import logger, default_format
 import logging
-import configparser  # 添加这行导入
-import tomli  # 添加这行导入
+import configparser
+import tomli
 
 # 禁用默认的日志输出
 # logger.remove()
@@ -24,6 +24,9 @@ class BotConfig:
     MONGODB_HOST: str = "127.0.0.1"
     MONGODB_PORT: int = 27017
     DATABASE_NAME: str = "MegBot"
+    MONGODB_USERNAME: Optional[str] = None  # 默认空值
+    MONGODB_PASSWORD: Optional[str] = None  # 默认空值
+    MONGODB_AUTH_SOURCE: Optional[str] = None  # 默认空值
     
     BOT_QQ: Optional[int] = None
     BOT_NICKNAME: Optional[str] = None
@@ -61,7 +64,10 @@ class BotConfig:
                 config.MONGODB_HOST = db_config.get("host", config.MONGODB_HOST)
                 config.MONGODB_PORT = db_config.get("port", config.MONGODB_PORT)
                 config.DATABASE_NAME = db_config.get("name", config.DATABASE_NAME)
-            
+                config.MONGODB_USERNAME = db_config.get("username", config.MONGODB_USERNAME) or None  # 空字符串转为 None
+                config.MONGODB_PASSWORD = db_config.get("password", config.MONGODB_PASSWORD) or None  # 空字符串转为 None
+                config.MONGODB_AUTH_SOURCE = db_config.get("auth_source", config.MONGODB_AUTH_SOURCE) or None  # 空字符串转为 None
+                
             if "emoji" in toml_dict:
                 emoji_config = toml_dict["emoji"]
                 config.EMOJI_CHECK_INTERVAL = emoji_config.get("check_interval", config.EMOJI_CHECK_INTERVAL)

src/plugins/chat/image_utils.py

@@ -31,7 +31,10 @@ def storage_compress_image(image_data: bytes, max_size: int = 200) -> bytes:
         db = Database(
             host=bot_config.MONGODB_HOST,
             port=bot_config.MONGODB_PORT,
-            db_name=bot_config.DATABASE_NAME
+            db_name=bot_config.DATABASE_NAME,
+            username=bot_config.MONGODB_USERNAME,
+            password=bot_config.MONGODB_PASSWORD,
+            auth_source=bot_config.MONGODB_AUTH_SOURCE
         )
         
         # 检查是否已存在相同哈希值的图片

src/plugins/schedule/schedule_generator.py

@@ -11,10 +11,19 @@ from ...common.database import Database  # 使用正确的导入语法
 # from src.plugins.schedule.schedule_llm_module import LLMModel
 # from src.common.database import Database  # 使用正确的导入语法
 
+# 获取当前文件的绝对路径
+#TODO: 这个好几个地方用需要封装
+current_dir = os.path.dirname(os.path.abspath(__file__))
+root_dir = os.path.abspath(os.path.join(current_dir, '..', '..', '..'))
+load_dotenv(os.path.join(root_dir, '.env'))
+
 Database.initialize(
-            os.getenv("MONGODB_HOST"),
-            int(os.getenv("MONGODB_PORT")),
-            os.getenv("DATABASE_NAME")
+            host= os.getenv("MONGODB_HOST"),
+            port= int(os.getenv("MONGODB_PORT")),
+            db_name=  os.getenv("DATABASE_NAME"),
+            username= os.getenv("MONGODB_USERNAME"),
+            password= os.getenv("MONGODB_PASSWORD"),
+            auth_source=os.getenv("MONGODB_AUTH_SOURCE")
         )
 
 class ScheduleGenerator: