aps-agent/server/integrations/jms_auth_client.py

from __future__ import annotations

import os
from typing import Any
from urllib.parse import urlparse

import httpx


class JmsAuthClientError(Exception):
    def __init__(self, code: str, message: str, status_code: int) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.status_code = status_code


class JmsAuthClient:
    """Small typed boundary around the JMS authentication endpoints."""

    def __init__(
        self,
        *,
        base_url: str | None = None,
        timeout_seconds: float | None = None,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        configured_base_url = (base_url or os.environ.get("JMS_AUTH_BASE_URL") or "").strip()
        if not configured_base_url:
            raise JmsAuthClientError(
                "AUTH_CONFIG_INVALID", "JMS 认证服务地址未配置（JMS_AUTH_BASE_URL）", 503,
            )
        self.base_url = configured_base_url.rstrip("/")
        try:
            self.timeout_seconds = timeout_seconds or float(
                os.environ.get("JMS_AUTH_TIMEOUT_SECONDS") or "10"
            )
        except ValueError as exc:
            raise JmsAuthClientError(
                "AUTH_CONFIG_INVALID", "JMS 认证超时时间配置无效", 503,
            ) from exc
        if self.timeout_seconds <= 0:
            raise JmsAuthClientError(
                "AUTH_CONFIG_INVALID", "JMS 认证超时时间必须大于 0", 503,
            )
        self.transport = transport
        parsed = urlparse(self.base_url)
        if parsed.scheme != "https" or not parsed.netloc:
            raise JmsAuthClientError(
                "AUTH_CONFIG_INVALID",
                "JMS 认证服务地址必须是有效的 HTTPS 地址",
                503,
            )

    async def _request(
        self,
        method: str,
        path: str,
        *,
        params: dict[str, str] | None = None,
        json_body: dict[str, Any] | None = None,
        form_body: dict[str, str] | None = None,
        token_name: str | None = None,
        token_value: str | None = None,
    ) -> Any:
        headers: dict[str, str] = {"Accept": "application/json"}
        if token_name and token_value:
            headers[token_name] = token_value
        try:
            async with httpx.AsyncClient(
                base_url=self.base_url,
                timeout=self.timeout_seconds,
                follow_redirects=False,
                transport=self.transport,
            ) as client:
                response = await client.request(
                    method,
                    path,
                    params=params,
                    json=json_body,
                    data=form_body,
                    headers=headers,
                )
        except httpx.TimeoutException as exc:
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_TIMEOUT", "JMS 认证服务响应超时", 504,
            ) from exc
        except httpx.RequestError as exc:
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_UNAVAILABLE", "无法连接 JMS 认证服务", 502,
            ) from exc

        try:
            envelope = response.json()
        except ValueError as exc:
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 认证服务返回了无效响应", 502,
            ) from exc
        if not isinstance(envelope, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 认证服务返回了无效响应", 502,
            )

        remote_code = envelope.get("code", 0)
        success = response.is_success and remote_code in (0, "0", None)
        if not success:
            message = str(envelope.get("message") or envelope.get("msg") or "JMS 认证请求失败")
            numeric_code = int(remote_code) if str(remote_code).isdigit() else response.status_code
            if path == "/admin/index/login":
                diagnostic = f"{message} {envelope.get('error') or ''}".lower()
                if "验证码" in diagnostic or "captcha" in diagnostic:
                    raise JmsAuthClientError(
                        "AUTH_CAPTCHA_INVALID", "验证码错误，请重新输入", 400,
                    )
                if any(word in diagnostic for word in ("停用", "禁用", "锁定", "disabled", "locked")):
                    raise JmsAuthClientError(
                        "AUTH_USER_DISABLED", "账号已停用或锁定，请联系管理员", 403,
                    )
                if any(word in diagnostic for word in ("租户", "企业", "tenant")):
                    raise JmsAuthClientError(
                        "AUTH_TENANT_INVALID", "企业名称或账号不匹配", 400,
                    )
                credential_rejected = (
                    response.status_code in {400, 401, 403}
                    or numeric_code in {400, 401, 403}
                    or any(word in diagnostic for word in (
                        "validation failure", "bad_request", "bad request",
                        "用户名", "用户不存在", "账号", "密码", "credential",
                    ))
                )
                if credential_rejected:
                    raise JmsAuthClientError(
                        "AUTH_INVALID_CREDENTIALS", "账号或密码错误", 401,
                    )
                raise JmsAuthClientError(
                    "AUTH_UPSTREAM_ERROR", "认证服务暂时不可用，请稍后重试", 502,
                )
            if response.status_code >= 500 or numeric_code >= 500:
                raise JmsAuthClientError("AUTH_UPSTREAM_ERROR", message, 502)
            status_code = response.status_code if response.status_code in {400, 401, 403} else 400
            raise JmsAuthClientError("AUTH_UPSTREAM_REJECTED", message, status_code)
        return envelope.get("data")

    async def precheck(self, *, tenant_code: str, username: str = "") -> dict[str, Any]:
        data = await self._request(
            "GET",
            "/admin/index/login/precheck",
            params={"tenantCode": tenant_code, "username": username},
        )
        if not isinstance(data, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 登录预检响应格式无效", 502,
            )
        return data

    async def tenants(self) -> list[dict[str, Any]]:
        data = await self._request("GET", "/admin/index/tenant")
        if not isinstance(data, list):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 租户列表响应格式无效", 502,
            )
        return [row for row in data if isinstance(row, dict)]

    async def captcha(self) -> dict[str, str]:
        data = await self._request("GET", "/admin/index/captcha")
        if not isinstance(data, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 验证码响应格式无效", 502,
            )
        return {str(key): str(value) for key, value in data.items() if value is not None}

    async def login(
        self,
        *,
        tenant_code: str,
        tenant_name: str,
        username: str,
        password: str,
        captcha_id: str = "",
        captcha_code: str = "",
    ) -> dict[str, Any]:
        form_body = {
            "tenantCode": tenant_code,
            "tenantName": tenant_name,
            "username": username,
            "password": password,
            "device": "PC",
        }
        if captcha_id:
            form_body["captchaId"] = captcha_id
        if captcha_code:
            form_body["captchaCode"] = captcha_code
        data = await self._request(
            "POST",
            "/admin/index/login",
            form_body=form_body,
        )
        if not isinstance(data, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 登录响应格式无效", 502,
            )
        return data

    async def profile(self, *, token_name: str, token_value: str) -> dict[str, Any]:
        data = await self._request(
            "GET", "/admin/auth/profile", token_name=token_name, token_value=token_value,
        )
        if not isinstance(data, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 用户信息响应格式无效", 502,
            )
        return data

    async def tenant(self, *, token_name: str, token_value: str) -> dict[str, Any]:
        data = await self._request(
            "GET", "/admin/auth/tenant", token_name=token_name, token_value=token_value,
        )
        if not isinstance(data, dict):
            raise JmsAuthClientError(
                "AUTH_UPSTREAM_INVALID", "JMS 租户信息响应格式无效", 502,
            )
        return data

    async def logout(self, *, token_name: str, token_value: str) -> None:
        await self._request(
            "POST", "/admin/auth/logout", token_name=token_name, token_value=token_value,
        )