shanxia2.0/backend/app/services/doclib_sso.py

"""门户用户与 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