zkxw/backend/app/core/login_crypto.py

"""登录凭据 RSA 解密（与前端 JSEncrypt PKCS#1 v1.5 兼容）。"""
from __future__ import annotations

import base64
import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding

from app.core.config import settings

logger = logging.getLogger(__name__)

_private_key = None  # type: ignore
_public_pem: Optional[str] = None

# 未配置 LOGIN_RSA_PRIVATE_KEY_PEM 时，各 Worker/实例共用此内置密钥（勿按进程随机生成）
_DEFAULT_PEM_PATH = Path(__file__).resolve().parent / "login_rsa_default.pem"


def _normalize_pem_text(pem: str) -> str:
    text = (pem or "").strip()
    if "\\n" in text and "-----BEGIN" in text:
        text = text.replace("\\n", "\n")
    return text


def _load_private_key_from_pem(pem: str):
    return serialization.load_pem_private_key(
        _normalize_pem_text(pem).encode("utf-8"),
        password=None,
        backend=default_backend(),
    )


def ensure_initialized() -> None:
    """加载 RSA 密钥对：优先环境变量，否则使用内置固定默认密钥。"""
    global _private_key, _public_pem
    if _private_key is not None:
        return

    pem = _normalize_pem_text(settings.LOGIN_RSA_PRIVATE_KEY_PEM or "")
    if pem:
        _private_key = _load_private_key_from_pem(pem)
        logger.info("登录 RSA：已自配置加载私钥")
    elif _DEFAULT_PEM_PATH.is_file():
        _private_key = _load_private_key_from_pem(_DEFAULT_PEM_PATH.read_text(encoding="utf-8"))
        logger.warning(
            "登录 RSA：未配置 LOGIN_RSA_PRIVATE_KEY_PEM，已使用内置默认密钥；"
            "生产环境建议配置独立 PEM。"
        )
    else:
        raise RuntimeError(
            f"登录 RSA 私钥不可用：未配置 LOGIN_RSA_PRIVATE_KEY_PEM，"
            f"且缺少内置文件 {_DEFAULT_PEM_PATH}"
        )

    _public_pem = (
        _private_key.public_key()
        .public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        .decode("utf-8")
    )


def get_public_pem() -> str:
    ensure_initialized()
    assert _public_pem is not None
    return _public_pem


def decrypt_login_payload(b64_cipher: str) -> Dict[str, Any]:
    """Base64 密文 -> JSON 对象（须含 username、password、ts 毫秒时间戳）。"""
    ensure_initialized()
    raw = base64.b64decode(b64_cipher)
    decrypted = _private_key.decrypt(raw, padding.PKCS1v15())
    text = decrypted.decode("utf-8")
    return json.loads(text)


def _validate_ts(ts_val: Any) -> bool:
    if ts_val is None:
        return False
    try:
        client_ms = int(ts_val)
    except (TypeError, ValueError):
        return False
    now_ms = int(time.time() * 1000)
    skew_ms = max(30, int(settings.LOGIN_ENCRYPT_TS_SKEW_SEC)) * 1000
    return abs(now_ms - client_ms) <= skew_ms


def extract_login_credentials(body: dict) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    """
    从请求体解析用户名、密码。
    返回 (username, password, error_message)；error_message 非空时表示失败。
    """
    enc = body.get("encryptedLogin")
    if isinstance(enc, str) and enc.strip():
        try:
            plain = decrypt_login_payload(enc.strip())
        except Exception:
            return None, None, "登录凭证无效或已过期"

        if not _validate_ts(plain.get("ts")):
            return None, None, "登录凭证已过期，请重新登录"

        username = (plain.get("username") or "").strip()
        password = plain.get("password")
        if password is None:
            password = ""
        elif not isinstance(password, str):
            password = str(password)

        if not username or not password:
            return None, None, "用户名和密码不能为空"
        return username, password, None

    if settings.LOGIN_ALLOW_PLAINTEXT_LOGIN:
        username = (body.get("username") or "").strip()
        password = body.get("password") or ""
        if not username or not password:
            return None, None, "用户名和密码不能为空"
        return username, password, None

    return None, None, "请使用加密方式登录"