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: self.base_url = ( base_url or os.environ.get("JMS_AUTH_BASE_URL") or "https://jms-mes-api.jzplat.com" ).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, )