holy-python/auth_service.py

import hashlib
import secrets
from datetime import datetime, timedelta
from typing import Any
from urllib.parse import urlencode

import requests

from config import Config
from database import connection, json_dumps, json_loads, now_iso
from task_service import TaskError


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _normalize_user(payload: dict[str, Any]) -> dict[str, str]:
    user_id = (
        payload.get("unionId")
        or payload.get("unionid")
        or payload.get("userid")
        or payload.get("userId")
        or payload.get("openId")
        or payload.get("openid")
    )
    name = (
        payload.get("nick")
        or payload.get("nickName")
        or payload.get("name")
        or payload.get("mobile")
        or user_id
    )
    if not user_id:
        raise TaskError("DINGTALK_USER_ID_MISSING", "钉钉用户信息中缺少用户标识", 502, {"payload": payload})
    return {
        "ding_user_id": str(user_id),
        "union_id": str(payload.get("unionId") or payload.get("unionid") or ""),
        "open_id": str(payload.get("openId") or payload.get("openid") or ""),
        "name": str(name or ""),
        "avatar": str(payload.get("avatarUrl") or payload.get("avatar") or ""),
        "mobile": str(payload.get("mobile") or ""),
        "email": str(payload.get("email") or ""),
    }


def build_dingtalk_login_url(frontend_redirect: str = "/tasks") -> dict[str, str]:
    if not Config.DINGTALK_CLIENT_ID:
        raise TaskError("DINGTALK_NOT_CONFIGURED", "未配置钉钉 DINGTALK_CLIENT_ID/DINGTALK_APP_KEY", 500)
    state = secrets.token_urlsafe(24)
    params = {
        "redirect_uri": Config.DINGTALK_REDIRECT_URI,
        "response_type": "code",
        "client_id": Config.DINGTALK_CLIENT_ID,
        "scope": "openid",
        "state": f"{state}:{frontend_redirect or '/tasks'}",
        "prompt": "consent",
    }
    return {"url": f"{Config.DINGTALK_AUTH_URL}?{urlencode(params)}", "state": params["state"]}


def exchange_dingtalk_user(code: str) -> dict[str, Any]:
    if not Config.DINGTALK_CLIENT_ID or not Config.DINGTALK_CLIENT_SECRET:
        raise TaskError("DINGTALK_NOT_CONFIGURED", "未配置钉钉应用 client_id/client_secret", 500)
    token_resp = requests.post(
        Config.DINGTALK_TOKEN_URL,
        json={
            "clientId": Config.DINGTALK_CLIENT_ID,
            "clientSecret": Config.DINGTALK_CLIENT_SECRET,
            "code": code,
            "grantType": "authorization_code",
        },
        timeout=20,
    )
    if token_resp.status_code >= 400:
        raise TaskError("DINGTALK_TOKEN_FAILED", "钉钉 access_token 获取失败", 502, {"body": token_resp.text})
    token_body = token_resp.json()
    access_token = token_body.get("accessToken") or token_body.get("access_token")
    if not access_token:
        raise TaskError("DINGTALK_TOKEN_MISSING", "钉钉响应中缺少 accessToken", 502, {"body": token_body})
    user_resp = requests.get(
        Config.DINGTALK_USER_URL,
        headers={"x-acs-dingtalk-access-token": access_token},
        timeout=20,
    )
    if user_resp.status_code >= 400:
        raise TaskError("DINGTALK_USER_FAILED", "钉钉用户信息获取失败", 502, {"body": user_resp.text})
    return user_resp.json()


def upsert_user_and_session(dingtalk_user: dict[str, Any]) -> tuple[dict[str, Any], str, datetime]:
    normalized = _normalize_user(dingtalk_user)
    token = secrets.token_urlsafe(36)
    session_hash = _sha256(token)
    now = now_iso()
    expires_at_dt = datetime.now().astimezone() + timedelta(days=Config.SESSION_EXPIRE_DAYS)
    expires_at = expires_at_dt.isoformat(timespec="seconds")
    with connection() as conn:
        existing = conn.execute(
            "SELECT id FROM auth_users WHERE ding_user_id=?",
            (normalized["ding_user_id"],),
        ).fetchone()
        if existing:
            user_id = existing["id"]
            conn.execute(
                """UPDATE auth_users SET union_id=?,open_id=?,name=?,avatar=?,mobile=?,email=?,
                   raw_json=?,last_login_at=?,updated_at=? WHERE id=?""",
                (
                    normalized["union_id"], normalized["open_id"], normalized["name"],
                    normalized["avatar"], normalized["mobile"], normalized["email"],
                    json_dumps(dingtalk_user), now, now, user_id,
                ),
            )
        else:
            cursor = conn.execute(
                """INSERT INTO auth_users
                   (ding_user_id,union_id,open_id,name,avatar,mobile,email,raw_json,last_login_at,created_at,updated_at)
                   VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
                (
                    normalized["ding_user_id"], normalized["union_id"], normalized["open_id"],
                    normalized["name"], normalized["avatar"], normalized["mobile"],
                    normalized["email"], json_dumps(dingtalk_user), now, now, now,
                ),
            )
            user_id = cursor.lastrowid
        conn.execute(
            """INSERT INTO auth_sessions (session_hash,user_id,expires_at,created_at,last_seen_at)
               VALUES (?,?,?,?,?)""",
            (session_hash, user_id, expires_at, now, now),
        )
    return get_user_by_session_token(token), token, expires_at_dt


def get_user_by_session_token(token: str | None) -> dict[str, Any] | None:
    if not token:
        return None
    session_hash = _sha256(token)
    current = now_iso()
    with connection() as conn:
        row = conn.execute(
            """SELECT s.id AS session_id,s.expires_at,u.* FROM auth_sessions s
               JOIN auth_users u ON u.id=s.user_id
               WHERE s.session_hash=? AND s.revoked_at IS NULL AND s.expires_at>?""",
            (session_hash, current),
        ).fetchone()
        if not row:
            return None
        conn.execute(
            "UPDATE auth_sessions SET last_seen_at=? WHERE id=?",
            (current, row["session_id"]),
        )
        raw = json_loads(row["raw_json"]) if row["raw_json"] else {}
        return {
            "id": str(row["ding_user_id"]),
            "name": row["name"],
            "account": row["mobile"] or row["email"] or row["ding_user_id"],
            "avatar": row["avatar"],
            "union_id": row["union_id"],
            "open_id": row["open_id"],
            "raw": raw,
        }


def revoke_session(token: str | None) -> None:
    if not token:
        return
    with connection() as conn:
        conn.execute(
            "UPDATE auth_sessions SET revoked_at=? WHERE session_hash=? AND revoked_at IS NULL",
            (now_iso(), _sha256(token)),
        )