Initial project import
This commit is contained in:
XM
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from threading import Lock
|
||||
from typing import Literal
|
||||
|
||||
from services.config import config
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
AuthRole = Literal["admin", "user"]
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _hash_key(value: str) -> str:
|
||||
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class AuthService:
|
||||
def __init__(self, storage: StorageBackend):
|
||||
self.storage = storage
|
||||
self._lock = Lock()
|
||||
self._items = self._load()
|
||||
self._last_used_flush_at: dict[str, datetime] = {}
|
||||
|
||||
@staticmethod
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
@staticmethod
|
||||
def _default_name(role: object) -> str:
|
||||
return "管理员密钥" if str(role or "").strip().lower() == "admin" else "普通用户"
|
||||
|
||||
def _normalize_item(self, raw: object) -> dict[str, object] | None:
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
role = self._clean(raw.get("role")).lower()
|
||||
if role not in {"admin", "user"}:
|
||||
return None
|
||||
key_hash = self._clean(raw.get("key_hash"))
|
||||
if not key_hash:
|
||||
return None
|
||||
item_id = self._clean(raw.get("id")) or uuid.uuid4().hex[:12]
|
||||
name = self._clean(raw.get("name")) or self._default_name(role)
|
||||
created_at = self._clean(raw.get("created_at")) or _now_iso()
|
||||
last_used_at = self._clean(raw.get("last_used_at")) or None
|
||||
return {
|
||||
"id": item_id,
|
||||
"name": name,
|
||||
"role": role,
|
||||
"key_hash": key_hash,
|
||||
"enabled": bool(raw.get("enabled", True)),
|
||||
"created_at": created_at,
|
||||
"last_used_at": last_used_at,
|
||||
}
|
||||
|
||||
def _load(self) -> list[dict[str, object]]:
|
||||
try:
|
||||
items = self.storage.load_auth_keys()
|
||||
except Exception:
|
||||
return []
|
||||
if not isinstance(items, list):
|
||||
return []
|
||||
return [normalized for item in items if (normalized := self._normalize_item(item)) is not None]
|
||||
|
||||
def _save(self) -> None:
|
||||
self.storage.save_auth_keys(self._items)
|
||||
|
||||
def _reload_locked(self) -> None:
|
||||
self._items = self._load()
|
||||
|
||||
@staticmethod
|
||||
def _public_item(item: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
"id": item.get("id"),
|
||||
"name": item.get("name"),
|
||||
"role": item.get("role"),
|
||||
"enabled": bool(item.get("enabled", True)),
|
||||
"created_at": item.get("created_at"),
|
||||
"last_used_at": item.get("last_used_at"),
|
||||
}
|
||||
|
||||
def list_keys(self, role: AuthRole | None = None) -> list[dict[str, object]]:
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
items = [item for item in self._items if role is None or item.get("role") == role]
|
||||
return [self._public_item(item) for item in items]
|
||||
|
||||
def _has_key_hash_locked(self, key_hash: str, *, exclude_id: str = "") -> bool:
|
||||
for item in self._items:
|
||||
item_id = self._clean(item.get("id"))
|
||||
if exclude_id and item_id == exclude_id:
|
||||
continue
|
||||
stored_hash = self._clean(item.get("key_hash"))
|
||||
if stored_hash and hmac.compare_digest(stored_hash, key_hash):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _build_key_hash_locked(self, raw_key: str, *, exclude_id: str = "") -> str:
|
||||
candidate = self._clean(raw_key)
|
||||
if not candidate:
|
||||
raise ValueError("请输入新的专用密钥")
|
||||
admin_key = self._clean(config.auth_key)
|
||||
if admin_key and hmac.compare_digest(candidate, admin_key):
|
||||
raise ValueError("这个密钥和管理员密钥冲突了，请换一个新的密钥")
|
||||
key_hash = _hash_key(candidate)
|
||||
if self._has_key_hash_locked(key_hash, exclude_id=exclude_id):
|
||||
raise ValueError("这个专用密钥已经存在，请换一个新的密钥")
|
||||
return key_hash
|
||||
|
||||
def _has_name_locked(self, name: str, *, role: AuthRole | None = None, exclude_id: str = "") -> bool:
|
||||
candidate = self._clean(name)
|
||||
if not candidate:
|
||||
return False
|
||||
for item in self._items:
|
||||
item_id = self._clean(item.get("id"))
|
||||
if exclude_id and item_id == exclude_id:
|
||||
continue
|
||||
if role is not None and item.get("role") != role:
|
||||
continue
|
||||
if self._clean(item.get("name")) == candidate:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _build_default_name_locked(self, role: AuthRole, *, exclude_id: str = "") -> str:
|
||||
base_name = self._default_name(role)
|
||||
if not self._has_name_locked(base_name, role=role, exclude_id=exclude_id):
|
||||
return base_name
|
||||
suffix = 2
|
||||
while True:
|
||||
candidate = f"{base_name} {suffix}"
|
||||
if not self._has_name_locked(candidate, role=role, exclude_id=exclude_id):
|
||||
return candidate
|
||||
suffix += 1
|
||||
|
||||
def _build_name_locked(self, name: str, *, role: AuthRole, exclude_id: str = "") -> str:
|
||||
candidate = self._clean(name)
|
||||
if not candidate:
|
||||
return self._build_default_name_locked(role, exclude_id=exclude_id)
|
||||
if self._has_name_locked(candidate, role=role, exclude_id=exclude_id):
|
||||
raise ValueError("这个名称已经在使用中了，换一个更容易区分的名称吧")
|
||||
return candidate
|
||||
|
||||
def create_key(self, *, role: AuthRole, name: str = "") -> tuple[dict[str, object], str]:
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
normalized_name = self._build_name_locked(name, role=role)
|
||||
while True:
|
||||
raw_key = f"sk-{secrets.token_urlsafe(24)}"
|
||||
try:
|
||||
key_hash = self._build_key_hash_locked(raw_key)
|
||||
break
|
||||
except ValueError:
|
||||
continue
|
||||
item = {
|
||||
"id": uuid.uuid4().hex[:12],
|
||||
"name": normalized_name,
|
||||
"role": role,
|
||||
"key_hash": key_hash,
|
||||
"enabled": True,
|
||||
"created_at": _now_iso(),
|
||||
"last_used_at": None,
|
||||
}
|
||||
self._items.append(item)
|
||||
self._save()
|
||||
return self._public_item(item), raw_key
|
||||
|
||||
def update_key(
|
||||
self,
|
||||
key_id: str,
|
||||
updates: dict[str, object],
|
||||
*,
|
||||
role: AuthRole | None = None,
|
||||
) -> dict[str, object] | None:
|
||||
normalized_id = self._clean(key_id)
|
||||
if not normalized_id:
|
||||
return None
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
for index, item in enumerate(self._items):
|
||||
if item.get("id") != normalized_id:
|
||||
continue
|
||||
if role is not None and item.get("role") != role:
|
||||
return None
|
||||
next_item = dict(item)
|
||||
next_role = "admin" if str(next_item.get("role") or "").strip().lower() == "admin" else "user"
|
||||
if "name" in updates and updates.get("name") is not None:
|
||||
next_item["name"] = self._build_name_locked(
|
||||
str(updates.get("name") or ""),
|
||||
role=next_role,
|
||||
exclude_id=normalized_id,
|
||||
)
|
||||
if "enabled" in updates and updates.get("enabled") is not None:
|
||||
next_item["enabled"] = bool(updates.get("enabled"))
|
||||
if "key" in updates and updates.get("key") is not None:
|
||||
next_item["key_hash"] = self._build_key_hash_locked(str(updates.get("key") or ""), exclude_id=normalized_id)
|
||||
self._items[index] = next_item
|
||||
self._save()
|
||||
return self._public_item(next_item)
|
||||
return None
|
||||
|
||||
def delete_key(self, key_id: str, *, role: AuthRole | None = None) -> bool:
|
||||
normalized_id = self._clean(key_id)
|
||||
if not normalized_id:
|
||||
return False
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
before = len(self._items)
|
||||
self._items = [
|
||||
item
|
||||
for item in self._items
|
||||
if not (item.get("id") == normalized_id and (role is None or item.get("role") == role))
|
||||
]
|
||||
if len(self._items) == before:
|
||||
return False
|
||||
self._save()
|
||||
return True
|
||||
|
||||
def authenticate(self, raw_key: str) -> dict[str, object] | None:
|
||||
candidate = self._clean(raw_key)
|
||||
if not candidate:
|
||||
return None
|
||||
candidate_hash = _hash_key(candidate)
|
||||
with self._lock:
|
||||
for index, item in enumerate(self._items):
|
||||
if not bool(item.get("enabled", True)):
|
||||
continue
|
||||
stored_hash = self._clean(item.get("key_hash"))
|
||||
if not stored_hash or not hmac.compare_digest(stored_hash, candidate_hash):
|
||||
continue
|
||||
next_item = dict(item)
|
||||
now = datetime.now(timezone.utc)
|
||||
next_item["last_used_at"] = now.isoformat()
|
||||
self._items[index] = next_item
|
||||
item_id = self._clean(next_item.get("id"))
|
||||
last_flush_at = self._last_used_flush_at.get(item_id)
|
||||
if last_flush_at is None or (now - last_flush_at).total_seconds() >= 60:
|
||||
try:
|
||||
self._save()
|
||||
self._last_used_flush_at[item_id] = now
|
||||
except Exception:
|
||||
pass
|
||||
return self._public_item(next_item)
|
||||
return None
|
||||
|
||||
|
||||
auth_service = AuthService(config.get_storage_backend())
|
||||
Reference in New Issue
Block a user