refactor!: Rewrite cookie logic; feat: Add CNM protocol verify; Add log system; fix: Variable length text cannot be primary key;

BiDuang · BiDuang · commit 017b06a40f52 · 2025-04-01T00:27:01.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -183,5 +183,9 @@ Services/Config/config.json
 # plugins
 Plugins/*/
 
+# logs
+Logs/*.lock
+Logs/*.log.*
+
 # version lock
 uv.lock
diff --git a/Logs/.gitkeep b/Logs/.gitkeep
diff --git a/Models/database.py b/Models/database.py
@@ -10,7 +10,7 @@
 class UserDb(Base):
     __tablename__ = "user"
     user_id: Mapped[str] = mapped_column(VARCHAR(32), primary_key=True)
-    username: Mapped[str] = mapped_column(TEXT, unique=True, nullable=False)
+    username: Mapped[str] = mapped_column(VARCHAR(255), nullable=False)
     email: Mapped[str] = mapped_column(TEXT, nullable=False)
     password: Mapped[str] = mapped_column(TEXT, nullable=False)
     created_at: Mapped[datetime] = mapped_column(
@@ -22,5 +22,5 @@ class PwdDb(Base):
     __tablename__ = "src_pwd"
     __table_args__ = (PrimaryKeyConstraint("user_id", "source"),)
     user_id: Mapped[str] = mapped_column(VARCHAR(32), nullable=False)
-    source: Mapped[str] = mapped_column(TEXT, nullable=False)
+    source: Mapped[str] = mapped_column(VARCHAR(10), nullable=False)
     data: Mapped[str] = mapped_column(TEXT, nullable=False)
diff --git a/Models/plugins.py b/Models/plugins.py
@@ -26,7 +26,7 @@ def search(self, keyword: str) -> list[BaseComicInfo]:
         pass
 
     @abstractmethod
-    def album(self, album_id: str) -> list[ComicInfo]:
+    def album(self, album_id: str) -> ComicInfo:
         pass
 
 
diff --git a/Models/user.py b/Models/user.py
@@ -1,6 +1,5 @@
 import json
 from datetime import datetime
-from http.cookies import BaseCookie
 
 from pydantic import BaseModel
 
@@ -14,24 +13,28 @@ class User(BaseModel):
 
 class UserData:
     uid: str
-    plugin_cookies: dict[str, BaseCookie[str]]
+    plugin_cookies: dict[str, dict[str, str]] | None
 
-    def __init__(self, uid: str, plugin_cookies: dict[str, BaseCookie[str]]):
+    def __init__(
+        self, uid: str, plugin_cookies: dict[str, dict[str, str]] | None = None
+    ) -> None:
         self.uid = uid
         self.plugin_cookies = plugin_cookies
 
-    def set_src_cookies(self, src: str, cookies: BaseCookie[str]) -> None:
-        self.plugin_cookies[src] = cookies
+    def set_src_cookies(self, src: str, cookies: dict[str, str | None]) -> None:
+        if self.plugin_cookies is None:
+            self.plugin_cookies = dict()
+        _cookies = {k: v for k, v in cookies.items() if v is not None}
+        self.plugin_cookies[src] = _cookies
 
-    def get_src_cookies(self, src: str) -> BaseCookie[str]:
-        if src in self.plugin_cookies:
-            return self.plugin_cookies[src]
-        else:
-            return BaseCookie[str]()
+    def get_src_cookies(self, src: str) -> dict[str, str]:
+        return self.plugin_cookies.get(src, dict()) if self.plugin_cookies else dict()
 
     def __str__(self):
-        return json.dumps(
-            {k: v.output(header="").strip() for k, v in self.plugin_cookies.items()}
+        return (
+            json.dumps(self.plugin_cookies, ensure_ascii=False)
+            if self.plugin_cookies
+            else "{}"
         )
 
 
diff --git a/README.md b/README.md
@@ -29,10 +29,10 @@ For Red Hat / CentOS:
 sudo yum install python3-devel mysql-devel pkgconfig
 ```
 
-Restore the dependencies:
+Restore the dependencies(Include the plugins' dependencies):
 
 ```bash
-uv sync
+uv sync --all-packages
 ```
 
 Modify the configuration file located at `Services/Config/config.toml.sample` and rename it to `config.toml`.
diff --git a/Routers/core.py b/Routers/core.py
@@ -1,4 +1,5 @@
 from fastapi import APIRouter
+
 from Models.response import BaseResponse, StandardResponse
 from Services.Modulator.manager import plugin_manager
 
@@ -10,6 +11,12 @@ def get_status() -> StandardResponse[None]:
     return StandardResponse(message="pong")
 
 
+@core_router.get("/sources", response_model=BaseResponse[set[str]])
+def get_sources() -> StandardResponse[set[str]]:
+    sources = plugin_manager.registered_source
+    return StandardResponse[set[str]](data=sources)
+
+
 @core_router.get("/protocol", response_model=BaseResponse[str])
 def get_cnm_version() -> StandardResponse[str]:
-    return StandardResponse[str](data=plugin_manager.cnm_version)
+    return StandardResponse[str](data=plugin_manager.cnm_version.__str__())
diff --git a/Services/Config/config.py b/Services/Config/config.py
@@ -29,11 +29,16 @@ class PluginConfig(BaseModel):
     strict_load: bool
 
 
+class LogConfig(BaseModel):
+    log_level: str
+
+
 class Config(BaseModel):
     security: SecurityConfig
     database: DatabaseConfig
     email: EmailConfig
     plugin: PluginConfig
+    log: LogConfig = LogConfig(log_level="INFO")
 
     @classmethod
     def load(cls):
diff --git a/Services/Config/config.toml.sample b/Services/Config/config.toml.sample
@@ -15,4 +15,7 @@ address =
 password =
 
 [plugin]
-strict_load = false
+strict_load = false
+
+# [log]
+# log_level =  # Set to debug, info, warning, error, or critical
diff --git a/Services/Mail/mail.py b/Services/Mail/mail.py
@@ -39,7 +39,7 @@ def _send_email(addr: str, subject: str, body: str):
 
 
 def send_captcha(addr: str, purpose: Purpose, ip: str) -> str:
-    captcha = str(secure_rng.randrange(10000, 99999))
+    captcha = str(secure_rng.randrange(100001, 999999))
     body = captcha_template.format(captcha=captcha, purpose=purpose.__str__(), ip=ip)
     _send_email(addr, purpose.__str__(), body)
     return captcha
diff --git a/Services/Modulator/manager.py b/Services/Modulator/manager.py
@@ -1,21 +1,21 @@
 import importlib
 import json
-import logging
 import os
-from http.cookies import BaseCookie
 from pathlib import Path
 from typing import Set
+import logging
 
 import toml
+from packaging.version import Version, parse
 
 from Models.plugins import BasePlugin, Plugin
 from Services.Config.config import config
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("[CNM]")
 
 
 class PluginManager:
-    cnm_version = "0.3.1"
+    cnm_version = Version("0.3.1")
 
     def __init__(self):
         self.strict = config.plugin.strict_load
@@ -30,11 +30,15 @@ def load_plugins(self) -> None:
                 if not self.load_plugin(
                     Path(os.path.join("Plugins", plugin)).resolve()
                 ):
-                    logger.error("An error occurred while loading plugins")
+                    logger.error("An error occurred while loading plugin")
                     if self.strict:
+                        logger.error(
+                            f"According to **STRICT** mode, server will abort if plugin {plugin} fails to load"
+                        )
                         raise RuntimeError("An error occurred while loading plugins")
 
     def load_plugin(self, plugin_dir: Path) -> bool:
+        logger.info(f"Loading plugin {plugin_dir.name}")
         try:
             with open(
                 plugin_dir.joinpath("pyproject.toml"), "r", encoding="utf-8"
@@ -45,15 +49,30 @@ def load_plugin(self, plugin_dir: Path) -> bool:
                 logger.warning(f"Plugin {plugin_name} already loaded")
                 return False
 
+            version = parse(plugin_info["tool"]["cnm"]["version"])
+            if (
+                self.cnm_version.major != version.major
+                or self.cnm_version.minor != version.minor
+            ):
+                logger.error(
+                    f"Plugin {plugin_name}'s CNM version {version} is not compatible with server's CNM version {self.cnm_version}"
+                )
+                return False
+
             src_list: list[str] = []
             for src in plugin_info["tool"]["cnm"]["source"]:
+                if len(src) > 10:
+                    logger.error(
+                        f"Failed to load {plugin_name}, source id {src} is too long"
+                    )
+                    return False
                 if src in self.registered_source:
                     logger.error(
                         f"Failed to load {plugin_name}, source {plugin_info['plugin']['source']} has already been registered"
                     )
                     return False
                 else:
-                    logger.info(f"Registering source {src}")
+                    logger.info(f"Registering source {src}...")
                     src_list.append(src)
 
             module = importlib.import_module(f"Plugins.{plugin_dir.name}.main")
@@ -65,9 +84,9 @@ def load_plugin(self, plugin_dir: Path) -> bool:
                         Plugin(
                             name=plugin_name,
                             version=plugin_info["project"]["version"],
-                            cnm_version=plugin_info["plugin"]["cnm"]["version"],
-                            source=plugin_info["plugin"]["cnm"]["source"],
-                            service=plugin_info["cnm"]["service"],
+                            cnm_version=plugin_info["tool"]["cnm"]["version"],
+                            source=plugin_info["tool"]["cnm"]["source"],
+                            service=plugin_info["tool"]["cnm"]["service"],
                             instance=instance,
                         )
                     )
@@ -81,14 +100,16 @@ def load_plugin(self, plugin_dir: Path) -> bool:
                 return False
         except ModuleNotFoundError as module_err:
             logger.error(
-                f"Failed to load {plugin_dir.name}, plugin requires some dependencies: {module_err.msg}"
+                f" Failed to load {plugin_dir.name}, plugin requires some dependencies: {module_err.msg}"
             )
             return False
         except FileNotFoundError:
             logger.error(f"Failed to load plugin {plugin_dir.name}'s information")
             return False
         except ImportError:
-            logger.error(f"Failed to import plugin {plugin_dir.name}'s main module")
+            logger.error(
+                f"Failed to import plugin {plugin_dir.name}'s main module"
+            )
             return False
         except KeyError:
             logger.error(f"Failed to load plugin {plugin_dir.name}'s information")
@@ -100,7 +121,7 @@ def load_plugin(self, plugin_dir: Path) -> bool:
             logger.exception(e)
             return False
 
-    def unload_plugin(self) -> None:
+    def unload_plugins(self) -> None:
         while len(self.plugins) > 0:
             plugin = self.plugins.pop()
             plugin.instance.on_unload()
@@ -114,23 +135,27 @@ def get_source(self, source: str) -> Plugin | None:
 
 class PluginUtils:
     @staticmethod
-    def load_cookies(cookies_str: str | None) -> dict[str, BaseCookie[str]]:
-        cookies = dict[str, BaseCookie[str]]()
-        if cookies_str is None or cookies_str == "":
-            return cookies
+    def load_cookies(cookies_str: str | None) -> dict[str, dict[str, str]]:
+        if not cookies_str:
+            return {}
 
         try:
-            plugin_cookies: dict[str, str] = json.loads(cookies_str)
-
-            for src in plugin_manager.registered_source:
-                if src in plugin_cookies:
-                    cookies[src] = BaseCookie[str](plugin_cookies[src])
-                else:
-                    cookies[src] = BaseCookie[str]()
-        except Exception:
-            logger.warning("Failed to load cookies")
+            plugin_cookies = json.loads(cookies_str)
+            if not isinstance(plugin_cookies, dict):
+                raise ValueError("Invalid cookies format")
+
+            return {
+                src: plugin_cookies.get(src, {})
+                for src in plugin_manager.registered_source
+            }
+        except (json.JSONDecodeError, ValueError) as e:
+            logger.warning(f"Failed to load cookies: {e}")
+        except Exception as e:
+            logger.exception(
+                "Unexpected error occurred while loading cookies", exc_info=e
+            )
 
-        return cookies
+        return {}
 
 
 plugin_manager = PluginManager()
diff --git a/main.py b/main.py
@@ -1,13 +1,19 @@
+import logging
 from contextlib import asynccontextmanager
 
+import fastapi
+import uvicorn
+from concurrent_log_handler import ConcurrentTimedRotatingFileHandler
 from fastapi import FastAPI, HTTPException
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
+from rich.logging import RichHandler
 from slowapi.errors import RateLimitExceeded
 
 from Models.response import http_exception_handler, validation_exception_handler
 from Routers.comic import comic_router
 from Routers.user import user_router
+from Services.Config.config import config
 from Services.Database.database import Base, engine
 from Services.Limiter.limiter import (
     LimitUploadSize,
@@ -16,14 +22,31 @@
 )
 from Services.Modulator.manager import plugin_manager
 
+logging.basicConfig(
+    level=config.log.log_level,
+    format="%(asctime)s - %(name)s [%(levelname)s] : %(message)s",
+    datefmt="[%X]",
+    handlers=[
+        RichHandler(rich_tracebacks=True, tracebacks_suppress=[fastapi]),
+        ConcurrentTimedRotatingFileHandler(
+            "Logs/latest.log", when="midnight", interval=1
+        ),
+    ],
+)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.getLogger("httpcore").setLevel(logging.WARNING)
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     Base.metadata.create_all(engine, checkfirst=True)
+    plugin_manager.load_plugins()
     yield
+    plugin_manager.unload_plugins()
+
 
+app = FastAPI(lifespan=lifespan)
 
-app = FastAPI()
 
 app.state.limiter = freq_limiter  # type: ignore
 app.add_exception_handler(RateLimitExceeded, RateLimitExceeded_handler)  # type: ignore
@@ -39,7 +62,8 @@ async def lifespan(app: FastAPI):
 )
 app.add_middleware(LimitUploadSize, max_upload_size=1024 * 1024 * 25)  # ~25MB
 
-plugin_manager.load_plugins()
-
 app.include_router(user_router)
 app.include_router(comic_router)
+
+if __name__ == "__main__":
+    uvicorn.run(app, log_level="trace")
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,6 +8,7 @@ requires-python = ">=3.12"
 dependencies = [
     "aiocache>=0.12.3",
     "bcrypt>=4.3.0",
+    "concurrent-log-handler>=0.9.25",
     "fastapi[standard]>=0.115.12",
     "mysqlclient>=2.2.7",
     "nest-asyncio>=1.6.0",
