vent_agent_2.0/extensions/shaanxi/client.py

"""多业务共享连接池，各请求显式携带平台身份。"""
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()