"""多业务共享连接池，各请求显式携带平台身份。""" from dataclasses import dataclass import hashlib import json import mimetypes from pathlib import Path from urllib.parse import unquote import httpx from extensions.shaanxi.settings import GatewaySettings from runtime.business_models import RemoteFile class GatewayError(RuntimeError): def __init__(self, operation, status=0, *, retryable=False, retry_after=None): self.status, self.retryable = status, retryable self.retry_after = retry_after super().__init__(f"{operation}失败（状态 {status}）") class GatewayStopped(GatewayError): pass @dataclass(frozen=True) class GatewayEvent: topic: str message_id: str payload: dict def relative_path(path): decoded = path for _ in range(3): decoded = unquote(decoded) if not path.startswith("/") or decoded.startswith("//") or "://" in decoded or ".." in decoded or "\\" in decoded or "?" in decoded or "#" in decoded or "\x00" in decoded: raise ValueError("必须使用安全的平台站内相对路径，参数通过 params 传递") return path class GatewayClient: def __init__(self, settings: GatewaySettings, *, transport=None, file_transport=None): self.settings = settings self.http = httpx.AsyncClient(timeout=settings.timeout, trust_env=False, follow_redirects=False, transport=transport) self.files = httpx.AsyncClient(timeout=settings.timeout, trust_env=False, follow_redirects=False, transport=file_transport or transport) async def close(self): await self.http.aclose() await self.files.aclose() def headers(self, agent_code): if not agent_code: raise ValueError("请求必须指定平台 agentCode") return {"Authorization": f"Bearer {self.settings.bearer}", "X-Agent-Code": agent_code} @staticmethod def _body(resp, operation): try: body = resp.json() except ValueError: body = {} if resp.status_code in (403, 500) and isinstance(body, dict) and "已停止" in str(body.get("detail", "")): raise GatewayStopped(operation, resp.status_code) if resp.status_code >= 400: try: retry_after = max(0, float(resp.headers["Retry-After"])) if not retry_after < float("inf"): retry_after = None except (KeyError, ValueError): retry_after = None raise GatewayError(operation, resp.status_code, retryable=resp.status_code in {408, 429, 500, 502, 503, 504}, retry_after=retry_after) if not isinstance(body, dict) or str(body.get("code")) != "200" or body.get("success") is False: raise GatewayError(operation, body.get("code", 0) if isinstance(body, dict) else 0) return body async def call(self, method, path, agent_code, **kwargs): try: response = await self.http.request(method, self.settings.base_url + path, headers=self.headers(agent_code), **kwargs) except httpx.RequestError: raise GatewayError(path, retryable=True) from None return self._body(response, path) async def register(self, business): binding = business.gateway if not binding.check_type_code: raise ValueError(f"{business.name} 尚未配置平台 checkTypeCode") return await self.call("POST", "/internal/agents/register", binding.agent_code, json={ "agentCode": binding.agent_code, "name": business.display_name, "sensorTypeList": binding.sensor_type_list, "topics": list(binding.topics), "checkTypeCode": binding.check_type_code, "enabled": True}) async def query(self, agent_code, source, resource, params=None): return await self.call("POST", "/internal/data/query", agent_code, json={"source": source, "resource": resource, "params": params or {}}) async def proxy(self, agent_code, path, *, method="GET", params=None, json_body=None): method = method.upper() if method not in {"GET", "POST", "PUT", "DELETE"}: raise ValueError("网关代理不支持该请求方法") body = await self.call("POST", "/internal/monitor/proxy", agent_code, json={"method": method, "path": relative_path(path), "params": params or {}, "json_body": json_body}) data = body.get("data") or {} status = data.get("status_code", 0) if not isinstance(status, int) or not 200 <= status < 300: raise GatewayError("监管平台代理", status, retryable=status in {429, 500, 502, 503, 504}) upstream = data.get("body") if isinstance(upstream, dict) and (upstream.get("success") is False or ("code" in upstream and str(upstream["code"]) not in {"0", "200"})): raise GatewayError("监管平台业务请求", upstream.get("code", 0)) return upstream async def token(self, agent_code): data = (await self.call("GET", "/internal/monitor/token", agent_code)).get("data") or {} if not data.get("access_token") or not data.get("client_id"): raise GatewayError("监管平台凭证下发") return data async def push(self, agent_code, payload, *, channel="ai_platform"): if channel not in {"ai_platform", "monitor_platform"}: raise ValueError("未配置的推送通道") result = await self.call("POST", f"/internal/push/{channel}", agent_code, json=payload) if channel == "ai_platform": data = result.get("data") if not isinstance(data, dict) or not data.get("checkId") or not isinstance(data.get("clueIds"), list): raise GatewayError("平台未返回完整研判回执") return result async def upload(self, agent_code, path: Path): if path.stat().st_size > self.settings.max_file_bytes: raise GatewayError("证据文件超过大小上限") with path.open("rb") as stream: response = await self.call("POST", "/internal/monitor/upload", agent_code, params={"path": self.settings.upload_path}, files={"file": (path.name, stream, mimetypes.guess_type(path.name)[0] or "application/octet-stream")}) platform = response.get("data") or {} data = platform.get("data", platform) if not isinstance(data, dict) or data.get("status") != "successed" or not data.get("path"): raise GatewayError("证据文件上传") return data["path"] async def download(self, agent_code, source: RemoteFile, target: Path): if not self.settings.file_base_url: raise ValueError("下载文件需要对方提供 SHANXI_FILE_BASE_URL 和文件接口") url = self.settings.file_base_url + relative_path(source.path) token = await self.token(agent_code) target.parent.mkdir(parents=True, exist_ok=True) partial = target.with_name(target.name + ".part") try: for attempt in range(2): async with self.files.stream("GET", url, params=source.params, headers={ "Authorization": f"Bearer {token['access_token']}", "clientid": token["client_id"]}) as resp: if resp.status_code in (401, 403) and attempt == 0: new_token = await self.token(agent_code) if new_token["access_token"] == token["access_token"]: raise GatewayError("文件凭证已失效且网关未下发新凭证", 401) token = new_token continue if resp.status_code != 200: raise GatewayError("文件下载", resp.status_code) if "json" in resp.headers.get("content-type", "").lower(): raise GatewayError("下载接口返回 JSON，未获得文件") digest, size = hashlib.sha256(), 0 with partial.open("wb") as output: async for chunk in resp.aiter_bytes(): size += len(chunk) if size > self.settings.max_file_bytes: raise GatewayError("下载文件超过大小上限") digest.update(chunk) output.write(chunk) if size == 0: raise GatewayError("下载接口返回空文件") partial.replace(target) return {"file_id": source.file_id, "version": source.version, "filename": source.filename, "local_path": str(target), "sha256": digest.hexdigest(), "size": size} except httpx.RequestError: raise GatewayError("文件下载", retryable=True) from None finally: partial.unlink(missing_ok=True) async def events(self, agent_code): """连接成功先产出 None；空 topics 也可用连接校验感知逻辑启停。""" async with self.http.stream("GET", self.settings.base_url + "/internal/events/stream", params={"agentCode": agent_code}, headers=self.headers(agent_code), timeout=httpx.Timeout(self.settings.timeout, read=None)) as resp: if resp.status_code == 403: raise GatewayStopped("事件订阅", 403) if resp.status_code != 200: raise GatewayError("事件订阅", resp.status_code, retryable=True) yield None fields, size = {}, 0 async for line in resp.aiter_lines(): if not line: if fields.get("data"): payload = json.loads("\n".join(fields["data"])) if not isinstance(payload, dict): raise GatewayError("事件载荷必须为 JSON 对象") yield GatewayEvent(fields.get("event", ""), fields.get("id", ""), payload) fields, size = {}, 0 elif not line.startswith(":"): size += len(line) if size > 1024 * 1024: raise GatewayError("事件帧超过大小上限") key, _, value = line.partition(":") if key == "data": fields.setdefault("data", []).append(value.lstrip()) elif key in {"event", "id"}: fields[key] = value.strip()