Initial project import

2026-09-18 03:46:58 +08:00
commit b600da4b76
209 changed files with 45904 additions and 0 deletions
@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from services.storage.factory import create_storage_backend
+
+__all__ = ["create_storage_backend"]
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class StorageBackend(ABC):
+    """抽象存储后端基类"""
+
+    @abstractmethod
+    def load_accounts(self) -> list[dict[str, Any]]:
+        """加载所有账号数据"""
+        pass
+
+    @abstractmethod
+    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
+        """保存所有账号数据"""
+        pass
+
+    @abstractmethod
+    def load_auth_keys(self) -> list[dict[str, Any]]:
+        """加载所有鉴权密钥数据"""
+        pass
+
+    @abstractmethod
+    def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
+        """保存所有鉴权密钥数据"""
+        pass
+
+    @abstractmethod
+    def health_check(self) -> dict[str, Any]:
+        """健康检查，返回存储后端状态"""
+        pass
+
+    @abstractmethod
+    def get_backend_info(self) -> dict[str, Any]:
+        """获取存储后端信息"""
+        pass
@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from sqlalchemy import Column, String, Text, create_engine, Integer, text
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+from services.storage.base import StorageBackend
+
+Base = declarative_base()
+
+
+class AccountModel(Base):
+    """账号数据模型"""
+    __tablename__ = "accounts"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    access_token = Column(String(2048), unique=True, nullable=False, index=True)
+    data = Column(Text, nullable=False)  # JSON 格式存储完整账号数据
+
+
+class AuthKeyModel(Base):
+    """鉴权密钥数据模型"""
+    __tablename__ = "auth_keys"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    key_id = Column(String(255), unique=True, nullable=False, index=True)
+    data = Column(Text, nullable=False)
+
+
+class DatabaseStorageBackend(StorageBackend):
+    """数据库存储后端（支持 SQLite、PostgreSQL、MySQL 等）"""
+
+    def __init__(self, database_url: str):
+        self.database_url = database_url
+        self.engine = create_engine(
+            database_url,
+            pool_pre_ping=True,  # 自动检测连接是否有效
+            pool_recycle=3600,   # 1小时回收连接
+        )
+        Base.metadata.create_all(self.engine)
+        self.Session = sessionmaker(bind=self.engine)
+
+    def load_accounts(self) -> list[dict[str, Any]]:
+        """从数据库加载账号数据"""
+        session = self.Session()
+        try:
+            accounts = []
+            for row in session.query(AccountModel).all():
+                try:
+                    account_data = json.loads(row.data)
+                    if isinstance(account_data, dict):
+                        accounts.append(account_data)
+                except json.JSONDecodeError:
+                    continue
+            return accounts
+        finally:
+            session.close()
+
+    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
+        """保存账号数据到数据库"""
+        self._save_rows(AccountModel, accounts, "access_token")
+
+    def load_auth_keys(self) -> list[dict[str, Any]]:
+        """从数据库加载鉴权密钥数据"""
+        return self._load_rows(AuthKeyModel)
+
+    def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
+        """保存鉴权密钥数据到数据库"""
+        self._save_rows(AuthKeyModel, auth_keys, "id", "key_id")
+
+    def _load_rows(self, model: type[AccountModel] | type[AuthKeyModel]) -> list[dict[str, Any]]:
+        session = self.Session()
+        try:
+            items = []
+            for row in session.query(model).all():
+                try:
+                    item_data = json.loads(row.data)
+                    if isinstance(item_data, dict):
+                        items.append(item_data)
+                except json.JSONDecodeError:
+                    continue
+            return items
+        finally:
+            session.close()
+
+    def _save_rows(
+        self,
+        model: type[AccountModel] | type[AuthKeyModel],
+        items: list[dict[str, Any]],
+        source_key: str,
+        target_key: str | None = None,
+    ) -> None:
+        session = self.Session()
+        try:
+            key_column = target_key or source_key
+            existing_rows = {
+                str(getattr(row, key_column)): row
+                for row in session.query(model).all()
+            }
+            incoming_keys: set[str] = set()
+
+            for item in items:
+                if not isinstance(item, dict):
+                    continue
+                key_value = str(item.get(source_key) or "").strip()
+                if not key_value:
+                    continue
+                if key_value in incoming_keys:
+                    raise ValueError(f"Duplicate {source_key} in storage snapshot")
+
+                incoming_keys.add(key_value)
+                serialized_data = json.dumps(item, ensure_ascii=False)
+                existing_row = existing_rows.get(key_value)
+                if existing_row is None:
+                    session.add(
+                        model(
+                            **{key_column: key_value},
+                            data=serialized_data,
+                        )
+                    )
+                elif existing_row.data != serialized_data:
+                    existing_row.data = serialized_data
+
+            for key_value, row in existing_rows.items():
+                if key_value not in incoming_keys:
+                    session.delete(row)
+
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def health_check(self) -> dict[str, Any]:
+        """健康检查"""
+        try:
+            session = self.Session()
+            try:
+                # 尝试执行简单查询
+                session.execute(text("SELECT 1"))
+                count = session.query(AccountModel).count()
+                auth_key_count = session.query(AuthKeyModel).count()
+                return {
+                    "status": "healthy",
+                    "backend": "database",
+                    "database_url": self._mask_password(self.database_url),
+                    "account_count": count,
+                    "auth_key_count": auth_key_count,
+                }
+            finally:
+                session.close()
+        except Exception as e:
+            return {
+                "status": "unhealthy",
+                "backend": "database",
+                "error": str(e),
+            }
+
+    def get_backend_info(self) -> dict[str, Any]:
+        """获取存储后端信息"""
+        db_type = "unknown"
+        if "sqlite" in self.database_url:
+            db_type = "sqlite"
+        elif "postgresql" in self.database_url or "postgres" in self.database_url:
+            db_type = "postgresql"
+        elif "mysql" in self.database_url:
+            db_type = "mysql"
+        
+        return {
+            "type": "database",
+            "db_type": db_type,
+            "description": f"数据库存储 ({db_type})",
+            "database_url": self._mask_password(self.database_url),
+        }
+
+    @staticmethod
+    def _mask_password(url: str) -> str:
+        """隐藏数据库连接字符串中的密码"""
+        if "://" not in url:
+            return url
+        try:
+            protocol, rest = url.split("://", 1)
+            if "@" in rest:
+                credentials, host = rest.split("@", 1)
+                if ":" in credentials:
+                    username, _ = credentials.split(":", 1)
+                    return f"{protocol}://{username}:****@{host}"
+            return url
+        except Exception:
+            return url
@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from services.storage.base import StorageBackend
+from services.storage.database_storage import DatabaseStorageBackend
+from services.storage.git_storage import GitStorageBackend
+from services.storage.json_storage import JSONStorageBackend
+
+
+def create_storage_backend(data_dir: Path) -> StorageBackend:
+    """
+    根据环境变量创建存储后端
+    
+    环境变量：
+    - STORAGE_BACKEND: json|sqlite|postgres|git (默认 json)
+    - DATABASE_URL: 数据库连接字符串 (用于 sqlite/postgres)
+    - GIT_REPO_URL: Git 仓库地址 (用于 git)
+    - GIT_TOKEN: Git 访问令牌 (用于 git)
+    - GIT_BRANCH: Git 分支 (默认 main)
+    - GIT_FILE_PATH: Git 仓库中的文件路径 (默认 accounts.json)
+    """
+    backend_type = os.getenv("STORAGE_BACKEND", "json").lower().strip()
+    
+    print(f"[storage] Initializing storage backend: {backend_type}")
+    
+    if backend_type == "json":
+        # 本地 JSON 文件存储
+        file_path = data_dir / "accounts.json"
+        auth_keys_path = data_dir / "auth_keys.json"
+        print(f"[storage] Using JSON storage: {file_path}")
+        return JSONStorageBackend(file_path, auth_keys_path)
+    
+    elif backend_type in ("sqlite", "postgres", "postgresql", "mysql", "database"):
+        # 数据库存储
+        database_url = os.getenv("DATABASE_URL", "").strip()
+        
+        if not database_url:
+            # 如果没有指定 DATABASE_URL，使用本地 SQLite
+            database_url = f"sqlite:///{data_dir / 'accounts.db'}"
+            print(f"[storage] No DATABASE_URL provided, using local SQLite: {database_url}")
+        else:
+            print(f"[storage] Using database storage: {_mask_password(database_url)}")
+        
+        return DatabaseStorageBackend(database_url)
+    
+    elif backend_type == "git":
+        # Git 仓库存储
+        repo_url = os.getenv("GIT_REPO_URL", "").strip()
+        token = os.getenv("GIT_TOKEN", "").strip()
+        branch = os.getenv("GIT_BRANCH", "main").strip()
+        file_path = os.getenv("GIT_FILE_PATH", "accounts.json").strip()
+        auth_keys_file_path = os.getenv("GIT_AUTH_KEYS_FILE_PATH", "auth_keys.json").strip()
+        
+        if not repo_url:
+            raise ValueError(
+                "GIT_REPO_URL is required when using git storage backend. "
+                "Please set GIT_REPO_URL environment variable."
+            )
+        
+        print(f"[storage] Using Git storage: {_mask_token(repo_url)}, branch: {branch}, file: {file_path}")
+        
+        cache_dir = data_dir / "git_cache"
+        return GitStorageBackend(
+            repo_url=repo_url,
+            token=token,
+            branch=branch,
+            file_path=file_path,
+            auth_keys_file_path=auth_keys_file_path,
+            local_cache_dir=cache_dir,
+        )
+    
+    else:
+        raise ValueError(
+            f"Unknown storage backend: {backend_type}. "
+            f"Supported backends: json, sqlite, postgres, git"
+        )
+
+
+def _mask_password(url: str) -> str:
+    """隐藏数据库连接字符串中的密码"""
+    if "://" not in url:
+        return url
+    try:
+        protocol, rest = url.split("://", 1)
+        if "@" in rest:
+            credentials, host = rest.split("@", 1)
+            if ":" in credentials:
+                username, _ = credentials.split(":", 1)
+                return f"{protocol}://{username}:****@{host}"
+        return url
+    except Exception:
+        return url
+
+
+def _mask_token(url: str) -> str:
+    """隐藏 URL 中的 token"""
+    if "@" in url and "://" in url:
+        protocol, rest = url.split("://", 1)
+        if "@" in rest:
+            _, host = rest.split("@", 1)
+            return f"{protocol}://****@{host}"
+    return url
@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any
+
+from git import Repo
+from git.exc import GitCommandError
+
+from services.storage.base import StorageBackend
+
+
+class GitStorageBackend(StorageBackend):
+    """Git 私有仓库存储后端"""
+
+    def __init__(
+        self,
+        repo_url: str,
+        token: str,
+        branch: str = "main",
+        file_path: str = "accounts.json",
+        auth_keys_file_path: str = "auth_keys.json",
+        local_cache_dir: Path | None = None,
+    ):
+        self.repo_url = repo_url
+        self.token = token
+        self.branch = branch
+        self.file_path = file_path
+        self.auth_keys_file_path = auth_keys_file_path
+        
+        # 本地缓存目录
+        if local_cache_dir is None:
+            local_cache_dir = Path(tempfile.gettempdir()) / "chatgpt2api_git_cache"
+        self.local_cache_dir = local_cache_dir
+        self.local_cache_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 构建带认证的 Git URL
+        self.auth_repo_url = self._build_auth_url(repo_url, token)
+
+    @staticmethod
+    def _build_auth_url(repo_url: str, token: str) -> str:
+        """构建带认证的 Git URL"""
+        if not token:
+            return repo_url
+        
+        # 支持 HTTPS 格式：https://github.com/user/repo.git
+        if repo_url.startswith("https://"):
+            # 插入 token
+            return repo_url.replace("https://", f"https://{token}@")
+        
+        # 支持 git@ 格式：git@github.com:user/repo.git
+        # 转换为 HTTPS 格式
+        if repo_url.startswith("git@"):
+            repo_url = repo_url.replace("git@", "https://")
+            repo_url = repo_url.replace(".com:", ".com/")
+            return repo_url.replace("https://", f"https://{token}@")
+        
+        return repo_url
+
+    def _clone_or_pull(self) -> Repo:
+        """克隆或拉取仓库"""
+        repo_path = self.local_cache_dir / "repo"
+        
+        if repo_path.exists() and (repo_path / ".git").exists():
+            # 仓库已存在，拉取最新代码
+            try:
+                repo = Repo(repo_path)
+                origin = repo.remote("origin")
+                origin.pull(self.branch)
+                return repo
+            except GitCommandError:
+                # 拉取失败，删除重新克隆
+                shutil.rmtree(repo_path)
+        
+        # 克隆仓库
+        repo = Repo.clone_from(
+            self.auth_repo_url,
+            repo_path,
+            branch=self.branch,
+        )
+        return repo
+
+    def load_accounts(self) -> list[dict[str, Any]]:
+        """从 Git 仓库加载账号数据"""
+        try:
+            return self._load_json_file(self.file_path)
+        except Exception as e:
+            print(f"[git-storage] load failed: {e}")
+            raise
+
+    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
+        """保存账号数据到 Git 仓库"""
+        try:
+            self._save_json_file(self.file_path, accounts, "Update accounts data")
+        except Exception as e:
+            print(f"[git-storage] save failed: {e}")
+            raise e
+
+    def load_auth_keys(self) -> list[dict[str, Any]]:
+        """从 Git 仓库加载鉴权密钥数据"""
+        try:
+            data = self._load_json_value(self.auth_keys_file_path)
+            if isinstance(data, dict):
+                data = data.get("items")
+            return data if isinstance(data, list) else []
+        except Exception as e:
+            print(f"[git-storage] load failed: {e}")
+            raise
+
+    def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
+        """保存鉴权密钥数据到 Git 仓库"""
+        try:
+            self._save_json_file(self.auth_keys_file_path, {"items": auth_keys}, "Update auth keys data")
+        except Exception as e:
+            print(f"[git-storage] save failed: {e}")
+            raise e
+
+    def _load_json_file(self, file_path: str) -> list[dict[str, Any]]:
+        data = self._load_json_value(file_path)
+        return data if isinstance(data, list) else []
+
+    def _load_json_value(self, file_path: str) -> Any:
+        repo = self._clone_or_pull()
+        file_full_path = Path(repo.working_dir) / file_path
+        if not file_full_path.exists():
+            return None
+        return json.loads(file_full_path.read_text(encoding="utf-8"))
+
+    def _save_json_file(self, file_path: str, items: Any, message: str) -> None:
+        repo = self._clone_or_pull()
+        file_full_path = Path(repo.working_dir) / file_path
+        file_full_path.parent.mkdir(parents=True, exist_ok=True)
+        file_full_path.write_text(
+            json.dumps(items, ensure_ascii=False, indent=2) + "\n",
+            encoding="utf-8",
+        )
+        repo.index.add([file_path])
+        if repo.is_dirty():
+            repo.index.commit(message)
+            repo.remote("origin").push(self.branch)
+
+    def health_check(self) -> dict[str, Any]:
+        """健康检查"""
+        try:
+            repo = self._clone_or_pull()
+            return {
+                "status": "healthy",
+                "backend": "git",
+                "repo_url": self._mask_token(self.repo_url),
+                "branch": self.branch,
+                "file_path": self.file_path,
+                "auth_keys_file_path": self.auth_keys_file_path,
+                "last_commit": repo.head.commit.hexsha[:8],
+            }
+        except Exception as e:
+            return {
+                "status": "unhealthy",
+                "backend": "git",
+                "error": str(e),
+            }
+
+    def get_backend_info(self) -> dict[str, Any]:
+        """获取存储后端信息"""
+        return {
+            "type": "git",
+            "description": "Git 私有仓库存储",
+            "repo_url": self._mask_token(self.repo_url),
+            "branch": self.branch,
+            "file_path": self.file_path,
+            "auth_keys_file_path": self.auth_keys_file_path,
+        }
+
+    @staticmethod
+    def _mask_token(url: str) -> str:
+        """隐藏 URL 中的 token"""
+        if "@" in url and "://" in url:
+            protocol, rest = url.split("://", 1)
+            if "@" in rest:
+                _, host = rest.split("@", 1)
+                return f"{protocol}://****@{host}"
+        return url
@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from services.storage.base import StorageBackend
+
+
+class JSONStorageBackend(StorageBackend):
+    """本地 JSON 文件存储后端"""
+
+    def __init__(self, file_path: Path, auth_keys_path: Path | None = None):
+        self.file_path = file_path
+        self.auth_keys_path = auth_keys_path or file_path.with_name("auth_keys.json")
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        self.auth_keys_path.parent.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def _load_json_list(file_path: Path) -> list[dict[str, Any]]:
+        if not file_path.exists():
+            return []
+        try:
+            data = json.loads(file_path.read_text(encoding="utf-8"))
+            return data if isinstance(data, list) else []
+        except (json.JSONDecodeError, Exception):
+            return []
+
+    @staticmethod
+    def _save_json_list(file_path: Path, items: list[dict[str, Any]]) -> None:
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(
+            json.dumps(items, ensure_ascii=False, indent=2) + "\n",
+            encoding="utf-8",
+        )
+
+    def load_accounts(self) -> list[dict[str, Any]]:
+        """从 JSON 文件加载账号数据"""
+        return self._load_json_list(self.file_path)
+
+    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
+        """保存账号数据到 JSON 文件"""
+        self._save_json_list(self.file_path, accounts)
+
+    def load_auth_keys(self) -> list[dict[str, Any]]:
+        """从 JSON 文件加载鉴权密钥数据"""
+        if not self.auth_keys_path.exists():
+            return []
+        try:
+            data = json.loads(self.auth_keys_path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, Exception):
+            return []
+        if isinstance(data, dict):
+            data = data.get("items")
+        return data if isinstance(data, list) else []
+
+    def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
+        """保存鉴权密钥数据到 JSON 文件"""
+        self.auth_keys_path.parent.mkdir(parents=True, exist_ok=True)
+        self.auth_keys_path.write_text(
+            json.dumps({"items": auth_keys}, ensure_ascii=False, indent=2) + "\n",
+            encoding="utf-8",
+        )
+
+    def health_check(self) -> dict[str, Any]:
+        """健康检查"""
+        try:
+            # 检查文件是否可读写
+            if self.file_path.exists():
+                self.file_path.read_text(encoding="utf-8")
+            return {
+                "status": "healthy",
+                "backend": "json",
+                "file_exists": self.file_path.exists(),
+                "file_path": str(self.file_path),
+                "auth_keys_file_exists": self.auth_keys_path.exists(),
+                "auth_keys_file_path": str(self.auth_keys_path),
+            }
+        except Exception as e:
+            return {
+                "status": "unhealthy",
+                "backend": "json",
+                "error": str(e),
+            }
+
+    def get_backend_info(self) -> dict[str, Any]:
+        """获取存储后端信息"""
+        return {
+            "type": "json",
+            "description": "本地 JSON 文件存储",
+            "file_path": str(self.file_path),
+            "file_exists": self.file_path.exists(),
+            "auth_keys_file_path": str(self.auth_keys_path),
+            "auth_keys_file_exists": self.auth_keys_path.exists(),
+        }