307 lines
10 KiB
Python
307 lines
10 KiB
Python
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import uuid
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from enum import Enum
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.repositories.api_key_repository import APIKeyRepository
|
||
from app.services.key_encryption import KeyEncryption, get_key_encryption
|
||
from app.services.key_verifier import KeyStatus, KeyVerifierFactory
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class KeySource(str, Enum):
|
||
SYSTEM = "system"
|
||
USER = "user"
|
||
ENV = "env"
|
||
|
||
|
||
@dataclass
|
||
class KeyCredentials:
|
||
"""统一凭证格式，支持单Key和双Key"""
|
||
api_key: str
|
||
secret_key: str | None = None
|
||
|
||
def to_dict(self) -> dict:
|
||
if self.secret_key:
|
||
return {"api_key": self.api_key, "secret_key": self.secret_key}
|
||
return {"api_key": self.api_key}
|
||
|
||
def to_json(self) -> str:
|
||
return json.dumps(self.to_dict())
|
||
|
||
|
||
@dataclass
|
||
class APIKeyConfig:
|
||
engine_type: str
|
||
key_source: KeySource
|
||
encrypted_key: str
|
||
key_hint: str
|
||
status: KeyStatus = KeyStatus.UNKNOWN
|
||
priority: int = 0
|
||
last_verified_at: str | None = None
|
||
created_at: str | None = None
|
||
user_id: str | None = None
|
||
|
||
|
||
class APIKeyManager:
|
||
_ENCRYPTION_KEY = os.getenv("API_KEY_ENCRYPTION_KEY", "geo-platform-default-key-change-in-production")
|
||
_USABLE_STATUSES = frozenset({KeyStatus.ACTIVE, KeyStatus.UNKNOWN})
|
||
_FALLBACK_SOURCES = frozenset({KeySource.SYSTEM, KeySource.ENV})
|
||
_ENV_MAPPING = {
|
||
"chatgpt": "OPENAI_API_KEY",
|
||
"perplexity": "PERPLEXITY_API_KEY",
|
||
"kimi": "MOONSHOT_API_KEY",
|
||
"wenxin": "BAIDU_QIANFAN_API_KEY",
|
||
"doubao": "DOUBAO_API_KEY",
|
||
"deepseek": "DEEPSEEK_API_KEY",
|
||
"qwen": "DASHSCOPE_API_KEY",
|
||
"gemini": "GOOGLE_API_KEY",
|
||
"yuanbao": "HUNYUAN_API_KEY",
|
||
}
|
||
|
||
def __init__(self):
|
||
self._keys: dict[str, list[APIKeyConfig]] = {}
|
||
|
||
def add_key(
|
||
self,
|
||
engine_type: str,
|
||
credentials: str | dict,
|
||
source: KeySource = KeySource.SYSTEM,
|
||
user_id: str | None = None,
|
||
priority: int = 0,
|
||
) -> APIKeyConfig:
|
||
if isinstance(credentials, dict):
|
||
key_hint = self._create_dual_key_hint(credentials)
|
||
credentials_json = json.dumps(credentials)
|
||
encrypted_key = self._encrypt(credentials_json)
|
||
else:
|
||
key_hint = self._mask_key(credentials)
|
||
encrypted_key = self._encrypt(credentials)
|
||
config = APIKeyConfig(
|
||
engine_type=engine_type,
|
||
key_source=source,
|
||
encrypted_key=encrypted_key,
|
||
key_hint=key_hint,
|
||
status=KeyStatus.UNKNOWN,
|
||
priority=priority,
|
||
user_id=user_id,
|
||
)
|
||
if engine_type not in self._keys:
|
||
self._keys[engine_type] = []
|
||
self._keys[engine_type].append(config)
|
||
self._keys[engine_type].sort(key=lambda k: k.priority, reverse=True)
|
||
return config
|
||
|
||
def _create_dual_key_hint(self, credentials: dict) -> str:
|
||
api_key = credentials.get("api_key", "")
|
||
secret_key = credentials.get("secret_key", "")
|
||
api_hint = self._mask_key(api_key) if api_key else "***"
|
||
secret_hint = self._mask_key(secret_key) if secret_key else "***"
|
||
return f"{api_hint}|{secret_hint}"
|
||
|
||
def get_key(self, engine_type: str, user_id: str | None = None) -> str | dict | None:
|
||
configs = self._keys.get(engine_type, [])
|
||
config = self._find_best_key(configs, user_id)
|
||
if not config:
|
||
return None
|
||
decrypted = self._decrypt(config.encrypted_key)
|
||
try:
|
||
return json.loads(decrypted)
|
||
except (json.JSONDecodeError, TypeError):
|
||
return decrypted
|
||
|
||
def get_credentials(self, engine_type: str, user_id: str | None = None) -> KeyCredentials | None:
|
||
key_data = self.get_key(engine_type, user_id)
|
||
if not key_data:
|
||
return None
|
||
if isinstance(key_data, dict):
|
||
return KeyCredentials(
|
||
api_key=key_data.get("api_key", ""),
|
||
secret_key=key_data.get("secret_key")
|
||
)
|
||
return KeyCredentials(api_key=key_data)
|
||
|
||
def get_all_keys(self, engine_type: str, user_id: str | None = None) -> dict | str | None:
|
||
return self.get_key(engine_type, user_id)
|
||
|
||
def _find_best_key(self, configs: list[APIKeyConfig], user_id: str | None = None) -> APIKeyConfig | None:
|
||
if user_id:
|
||
for c in configs:
|
||
if (
|
||
c.key_source == KeySource.USER
|
||
and c.user_id == user_id
|
||
and c.status in self._USABLE_STATUSES
|
||
):
|
||
return c
|
||
for c in configs:
|
||
if c.key_source in self._FALLBACK_SOURCES and c.status in self._USABLE_STATUSES:
|
||
return c
|
||
return None
|
||
|
||
def get_any_available_key(self, engine_type: str) -> str | None:
|
||
for c in self._keys.get(engine_type, []):
|
||
if c.status in self._USABLE_STATUSES:
|
||
return self._decrypt(c.encrypted_key)
|
||
return None
|
||
|
||
def remove_key(self, engine_type: str, key_hint: str) -> bool:
|
||
configs = self._keys.get(engine_type, [])
|
||
for i, c in enumerate(configs):
|
||
if c.key_hint == key_hint:
|
||
configs.pop(i)
|
||
return True
|
||
return False
|
||
|
||
def list_keys(self, engine_type: str | None = None) -> list[APIKeyConfig]:
|
||
if engine_type:
|
||
return self._keys.get(engine_type, [])
|
||
result = []
|
||
for configs in self._keys.values():
|
||
result.extend(configs)
|
||
return result
|
||
|
||
async def verify_key(self, engine_type: str, api_key: str) -> KeyStatus:
|
||
if not api_key or len(api_key) < 10:
|
||
return KeyStatus.INVALID
|
||
try:
|
||
return await KeyVerifierFactory.verify(engine_type, api_key)
|
||
except Exception as e:
|
||
logger.warning(f"[api_key_manager] Key verification failed: {e}")
|
||
return KeyStatus.UNKNOWN
|
||
|
||
def _encrypt(self, plaintext: str) -> str:
|
||
return get_key_encryption().encrypt(plaintext)
|
||
|
||
def _decrypt(self, ciphertext: str) -> str:
|
||
return get_key_encryption().decrypt(ciphertext)
|
||
|
||
def _mask_key(self, key: str) -> str:
|
||
if len(key) <= 8:
|
||
return "***"
|
||
return key[:3] + "..." + key[-3:]
|
||
|
||
def load_env_keys(self):
|
||
for engine, env_var in self._ENV_MAPPING.items():
|
||
key = os.getenv(env_var, "")
|
||
if key:
|
||
self.add_key(engine, key, source=KeySource.ENV, priority=0)
|
||
|
||
async def add_key_async(
|
||
self,
|
||
session: AsyncSession,
|
||
engine_type: str,
|
||
api_key: str,
|
||
source: KeySource = KeySource.USER,
|
||
user_id: uuid.UUID | None = None,
|
||
priority: int = 0,
|
||
) -> APIKeyConfig:
|
||
config = APIKeyConfig(
|
||
engine_type=engine_type,
|
||
key_source=source,
|
||
encrypted_key=self._encrypt(api_key),
|
||
key_hint=self._mask_key(api_key),
|
||
status=KeyStatus.UNKNOWN,
|
||
priority=priority,
|
||
user_id=str(user_id) if user_id else None,
|
||
)
|
||
repository = APIKeyRepository(session)
|
||
await repository.create(
|
||
user_id=user_id,
|
||
engine_type=engine_type,
|
||
encrypted_key=config.encrypted_key,
|
||
key_hint=config.key_hint,
|
||
key_source=source.value,
|
||
status=config.status.value,
|
||
priority=priority,
|
||
)
|
||
return config
|
||
|
||
async def get_key_async(
|
||
self,
|
||
session: AsyncSession,
|
||
engine_type: str,
|
||
user_id: uuid.UUID | None = None,
|
||
) -> str | None:
|
||
repository = APIKeyRepository(session)
|
||
keys = await repository.get_by_user(user_id, engine_type)
|
||
for key in keys:
|
||
source = KeySource(key.key_source)
|
||
status = KeyStatus(key.status)
|
||
if user_id:
|
||
if source == KeySource.USER and status in self._USABLE_STATUSES:
|
||
return self._decrypt(key.encrypted_key)
|
||
if source in self._FALLBACK_SOURCES and status in self._USABLE_STATUSES:
|
||
return self._decrypt(key.encrypted_key)
|
||
return None
|
||
|
||
async def get_any_available_key_async(
|
||
self,
|
||
session: AsyncSession,
|
||
engine_type: str,
|
||
) -> str | None:
|
||
repository = APIKeyRepository(session)
|
||
keys = await repository.get_by_user(None, engine_type)
|
||
for key in keys:
|
||
status = KeyStatus(key.status)
|
||
if status in self._USABLE_STATUSES:
|
||
return self._decrypt(key.encrypted_key)
|
||
return None
|
||
|
||
async def remove_key_async(
|
||
self,
|
||
session: AsyncSession,
|
||
user_id: uuid.UUID,
|
||
engine_type: str,
|
||
key_hint: str,
|
||
) -> bool:
|
||
repository = APIKeyRepository(session)
|
||
keys = await repository.get_by_user(user_id, engine_type)
|
||
for key in keys:
|
||
if key.key_hint == key_hint:
|
||
deleted = await repository.delete(key.id)
|
||
return deleted
|
||
return False
|
||
|
||
async def list_keys_async(
|
||
self,
|
||
session: AsyncSession,
|
||
user_id: uuid.UUID | None = None,
|
||
engine_type: str | None = None,
|
||
) -> list[APIKeyConfig]:
|
||
repository = APIKeyRepository(session)
|
||
if user_id:
|
||
keys = await repository.get_by_user(user_id, engine_type)
|
||
else:
|
||
keys = await repository.list_all(engine_type=engine_type)
|
||
return [
|
||
APIKeyConfig(
|
||
engine_type=k.engine_type,
|
||
key_source=KeySource(k.key_source),
|
||
encrypted_key=k.encrypted_key,
|
||
key_hint=k.key_hint,
|
||
status=KeyStatus(k.status),
|
||
priority=k.priority,
|
||
last_verified_at=str(k.last_verified_at) if k.last_verified_at else None,
|
||
created_at=str(k.created_at) if k.created_at else None,
|
||
user_id=str(k.user_id) if k.user_id else None,
|
||
)
|
||
for k in keys
|
||
]
|
||
|
||
async def update_key_status_async(
|
||
self,
|
||
session: AsyncSession,
|
||
key_id: uuid.UUID,
|
||
status: KeyStatus,
|
||
) -> bool:
|
||
repository = APIKeyRepository(session)
|
||
updated = await repository.update(key_id, status=status.value)
|
||
return updated is not None
|