258 lines
7.9 KiB
Python
258 lines
7.9 KiB
Python
"""门户用户与 Cloudreve 账号一一对应，并签发一次性进入票据。"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import secrets
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from typing import Optional
|
||
|
||
from app.core.config import settings
|
||
|
||
_lock = threading.Lock()
|
||
_tickets: dict[str, tuple[float, dict]] = {}
|
||
_cached_admin_token: str = ""
|
||
_cached_admin_token_until: float = 0
|
||
|
||
|
||
def cloudreve_email(user_name: str) -> str:
|
||
"""文档库邮箱：p.{门户用户名}@sso.shanxia.local。用户名需唯一。"""
|
||
raw = (user_name or "").strip()
|
||
safe = "".join(ch for ch in raw if ch.isalnum() or ch in ".-_")
|
||
if not safe:
|
||
safe = hashlib.sha256(raw.encode()).hexdigest()[:16]
|
||
return f"p.{safe}@sso.shanxia.local"
|
||
|
||
|
||
def _legacy_id_email(user_id: str) -> str:
|
||
"""改用用户名之前的邮箱，只用来把已有账号改过去，避免另建一个空账号。"""
|
||
raw = (user_id or "").strip().lower()
|
||
safe = "".join(ch for ch in raw if ch.isalnum() or ch == "-")
|
||
if not safe:
|
||
safe = hashlib.sha256(raw.encode()).hexdigest()[:32]
|
||
return f"p.{safe}@sso.shanxia.local"
|
||
|
||
|
||
def _password(user_id: str) -> str:
|
||
digest = hmac.new(
|
||
settings.SECRET_KEY.encode(),
|
||
f"doclib:{user_id}".encode(),
|
||
hashlib.sha256,
|
||
).hexdigest()
|
||
return digest[:24]
|
||
|
||
|
||
def issue_ticket(login: dict) -> str:
|
||
code = secrets.token_urlsafe(24)
|
||
with _lock:
|
||
now = time.time()
|
||
expired = [key for key, (exp, _) in _tickets.items() if exp < now]
|
||
for key in expired:
|
||
_tickets.pop(key, None)
|
||
_tickets[code] = (now + 60, login)
|
||
return code
|
||
|
||
|
||
def consume_ticket(code: str) -> Optional[dict]:
|
||
with _lock:
|
||
item = _tickets.pop((code or "").strip(), None)
|
||
if not item:
|
||
return None
|
||
exp, login = item
|
||
if exp < time.time():
|
||
return None
|
||
return login
|
||
|
||
|
||
def _request(method: str, path: str, payload: Optional[dict] = None, token: str = "") -> dict:
|
||
url = settings.DOCLIB_BASE_URL.rstrip("/") + "/api/v4" + path
|
||
body = None if payload is None else json.dumps(payload).encode()
|
||
req = urllib.request.Request(url, data=body, method=method)
|
||
req.add_header("Content-Type", "application/json")
|
||
if token:
|
||
req.add_header("Authorization", "Bearer " + token)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
data = json.load(resp)
|
||
except urllib.error.HTTPError as exc:
|
||
exc.read()
|
||
raise RuntimeError(f"文档库接口失败（{exc.code}）") from exc
|
||
if data.get("code") not in (0, None) and "data" not in data:
|
||
raise RuntimeError(data.get("msg") or "文档库接口失败")
|
||
if isinstance(data, dict) and data.get("code") not in (0, None):
|
||
raise RuntimeError(data.get("msg") or "文档库接口失败")
|
||
return data
|
||
|
||
|
||
def _admin_access_token() -> str:
|
||
global _cached_admin_token, _cached_admin_token_until
|
||
now = time.time()
|
||
if _cached_admin_token and now < _cached_admin_token_until:
|
||
return _cached_admin_token
|
||
data = _request(
|
||
"POST",
|
||
"/session/token",
|
||
{"email": settings.DOCLIB_ADMIN_EMAIL, "password": settings.DOCLIB_ADMIN_PASSWORD},
|
||
)
|
||
token = data["data"]["token"]["access_token"]
|
||
_cached_admin_token = token
|
||
_cached_admin_token_until = now + 20 * 60
|
||
return token
|
||
|
||
|
||
def _find_user(email: str, token: str) -> Optional[dict]:
|
||
data = _request(
|
||
"POST",
|
||
"/admin/user",
|
||
{
|
||
"page": 1,
|
||
"page_size": 20,
|
||
"order_by": "id",
|
||
"order_direction": "asc",
|
||
"conditions": {"user_email": email},
|
||
},
|
||
token,
|
||
)
|
||
users = (data.get("data") or {}).get("users") or []
|
||
for user in users:
|
||
if str(user.get("email") or "").lower() == email.lower():
|
||
return user
|
||
return None
|
||
|
||
|
||
def _ensure_user(email: str, nick: str, token: str) -> dict:
|
||
existing = _find_user(email, token)
|
||
display = (nick or email.split("@", 1)[0])[:50]
|
||
if existing is None:
|
||
created = _request(
|
||
"PUT",
|
||
"/admin/user",
|
||
{
|
||
"password": _password(email),
|
||
"user": {
|
||
"id": 0,
|
||
"email": email,
|
||
"nick": display,
|
||
"status": "active",
|
||
"group_users": 2,
|
||
},
|
||
},
|
||
token,
|
||
)
|
||
return created.get("data") or {}
|
||
if (existing.get("nick") or "") != display:
|
||
_request(
|
||
"PUT",
|
||
f"/admin/user/{existing['id']}",
|
||
{
|
||
"user": {
|
||
"id": existing["id"],
|
||
"email": email,
|
||
"nick": display,
|
||
"status": existing.get("status") or "active",
|
||
"group_users": existing.get("group_users") or 2,
|
||
}
|
||
},
|
||
token,
|
||
)
|
||
return existing
|
||
|
||
|
||
def _move_user(user: dict, email: str, nick: str, token: str) -> dict:
|
||
display = (nick or email.split("@", 1)[0])[:50]
|
||
_request(
|
||
"PUT",
|
||
f"/admin/user/{user['id']}",
|
||
{
|
||
"password": _password(email),
|
||
"user": {
|
||
"id": user["id"],
|
||
"email": email,
|
||
"nick": display,
|
||
"status": user.get("status") or "active",
|
||
"group_users": user.get("group_users") or 2,
|
||
},
|
||
},
|
||
token,
|
||
)
|
||
moved = dict(user)
|
||
moved["email"] = email
|
||
moved["nick"] = display
|
||
return moved
|
||
|
||
|
||
def ensure_portal_user(user_id: str, nick: str, token: str) -> dict:
|
||
"""按门户用户名找到或改名已有文档库账号，没有则新建。"""
|
||
name = (nick or "").strip() or (user_id or "").strip()
|
||
email = cloudreve_email(name)
|
||
current = _find_user(email, token)
|
||
if current is not None:
|
||
return _move_user(current, email, name, token)
|
||
legacy = _find_user(_legacy_id_email(user_id), token) if user_id else None
|
||
if legacy is not None:
|
||
return _move_user(legacy, email, name, token)
|
||
return _ensure_user(email, name, token)
|
||
|
||
|
||
def _login(email: str) -> dict:
|
||
# 密码按邮箱派生，和门户用户名一一对应。
|
||
data = _request(
|
||
"POST",
|
||
"/session/token",
|
||
{"email": email, "password": _password(email)},
|
||
)
|
||
return data["data"]
|
||
|
||
|
||
def _reset_password(user: dict, email: str, token: str) -> None:
|
||
_request(
|
||
"PUT",
|
||
f"/admin/user/{user['id']}",
|
||
{
|
||
"password": _password(email),
|
||
"user": {
|
||
"id": user["id"],
|
||
"email": email,
|
||
"nick": user.get("nick") or email.split("@", 1)[0],
|
||
"status": user.get("status") or "active",
|
||
"group_users": user.get("group_users") or 2,
|
||
},
|
||
},
|
||
token,
|
||
)
|
||
|
||
|
||
def open_session(user_id: str, nick: str) -> dict:
|
||
"""确保 Cloudreve 用户存在并登录，返回一次性进入地址。"""
|
||
name = (nick or "").strip() or (user_id or "").strip()
|
||
email = cloudreve_email(name)
|
||
token = _admin_access_token()
|
||
user = ensure_portal_user(user_id, name, token)
|
||
try:
|
||
login = _login(email)
|
||
except RuntimeError:
|
||
_admin_token_clear()
|
||
token = _admin_access_token()
|
||
user = _find_user(email, token) or user
|
||
_reset_password(user, email, token)
|
||
login = _login(email)
|
||
code = issue_ticket(login)
|
||
portal = settings.DOCLIB_PORTAL_ORIGIN.rstrip("/")
|
||
doclib = settings.DOCLIB_PUBLIC_URL.rstrip("/")
|
||
open_url = f"{doclib}/sso?code={code}&portal={portal}"
|
||
return {
|
||
"openUrl": open_url,
|
||
"cloudreveEmail": email,
|
||
"cloudreveUserId": (login.get("user") or {}).get("id"),
|
||
}
|
||
|
||
|
||
def _admin_token_clear() -> None:
|
||
global _cached_admin_token, _cached_admin_token_until
|
||
_cached_admin_token = ""
|
||
_cached_admin_token_until = 0
|