223 lines
9.5 KiB
Python
223 lines
9.5 KiB
Python
"""Client for the NDJSON RPC contract shipped in WeComAgent's Python SDK.
|
|
|
|
Only accepts runtime credentials explicitly inherited from the SDK host.
|
|
It does not extract credentials, launch the agent, change client switches, or retry calls.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import socket
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Mapping
|
|
|
|
MAX_FRAME_BYTES = 4 * 1024 * 1024
|
|
|
|
|
|
class SdkRpcError(RuntimeError):
|
|
def __init__(self, code, message, *, may_have_executed=False):
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.may_have_executed = may_have_executed
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RuntimeEndpoint:
|
|
host: str
|
|
port: int
|
|
token: str = field(repr=False)
|
|
|
|
def __post_init__(self):
|
|
if self.host not in ("127.0.0.1", "::1", "localhost"):
|
|
raise SdkRpcError("invalid_endpoint", "SDK 连接地址必须是本机回环地址")
|
|
if isinstance(self.port, bool) or not isinstance(self.port, int) or not 1 <= self.port <= 65535:
|
|
raise SdkRpcError("invalid_endpoint", "SDK 端口无效")
|
|
if not isinstance(self.token, str) or not self.token or len(self.token) > 65536:
|
|
raise SdkRpcError("invalid_endpoint", "SDK 授权参数无效")
|
|
|
|
@classmethod
|
|
def from_environment(cls, environ: Mapping[str, str] | None = None):
|
|
environ = os.environ if environ is None else environ
|
|
host = environ.get("WECOMAGENT_RPC_HOST", "")
|
|
port = environ.get("WECOMAGENT_RPC_PORT", "")
|
|
token = environ.get("WECOMAGENT_RPC_TOKEN", "")
|
|
if not all((host, port, token)):
|
|
raise SdkRpcError(
|
|
"sdk_runtime_unavailable",
|
|
"尚未获得企业微信 SDK 运行时连接参数；需要由客户端正常启动代理并在其 Python 工具环境中执行。",
|
|
)
|
|
try:
|
|
port = int(port)
|
|
except (TypeError, ValueError):
|
|
raise SdkRpcError("invalid_endpoint", "SDK 端口无效") from None
|
|
return cls(host, port, token)
|
|
|
|
|
|
class WeComSdkClient:
|
|
"""Serial request/response RPC. An uncertain send must be reconciled, never replayed."""
|
|
|
|
def __init__(self, endpoint: RuntimeEndpoint, *, timeout=8.0, max_frame_bytes=MAX_FRAME_BYTES):
|
|
if not math.isfinite(float(timeout)) or not 0 < float(timeout) <= 120:
|
|
raise ValueError("timeout must be within (0, 120]")
|
|
if not 256 <= max_frame_bytes <= MAX_FRAME_BYTES:
|
|
raise ValueError("invalid frame limit")
|
|
self.endpoint = endpoint
|
|
self.timeout = float(timeout)
|
|
self.max_frame_bytes = max_frame_bytes
|
|
self._socket = None
|
|
self._reader = None
|
|
self._lock = threading.Lock()
|
|
self._next_id = 1
|
|
self.server_version = ""
|
|
|
|
def _close(self):
|
|
reader, sock = self._reader, self._socket
|
|
self._reader = self._socket = None
|
|
if reader is not None:
|
|
try:
|
|
reader.close()
|
|
except OSError:
|
|
pass
|
|
if sock is not None:
|
|
try:
|
|
sock.close()
|
|
except OSError:
|
|
pass
|
|
|
|
def close(self):
|
|
with self._lock:
|
|
self._close()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_):
|
|
self.close()
|
|
|
|
def _encode(self, value):
|
|
try:
|
|
frame = (json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")) + "\n").encode("utf-8")
|
|
except (ValueError, TypeError, UnicodeError):
|
|
raise SdkRpcError("invalid_request", "SDK 请求不能序列化为 JSON") from None
|
|
if len(frame) > self.max_frame_bytes:
|
|
raise SdkRpcError("request_too_large", "SDK 请求超过大小限制")
|
|
return frame
|
|
|
|
def _read(self):
|
|
raw = self._reader.readline(self.max_frame_bytes + 1)
|
|
if not raw:
|
|
raise SdkRpcError("connection_closed", "SDK 连接已关闭")
|
|
if len(raw) > self.max_frame_bytes or not raw.endswith(b"\n"):
|
|
raise SdkRpcError("invalid_frame", "SDK 响应帧不完整或超过大小限制")
|
|
try:
|
|
result = json.loads(raw.decode("utf-8"))
|
|
except (ValueError, UnicodeError):
|
|
raise SdkRpcError("invalid_frame", "SDK 返回了无效的 JSON 响应") from None
|
|
if not isinstance(result, dict):
|
|
raise SdkRpcError("invalid_frame", "SDK 响应必须是 JSON 对象")
|
|
return result
|
|
|
|
def _connect(self):
|
|
try:
|
|
host = "127.0.0.1" if self.endpoint.host == "localhost" else self.endpoint.host
|
|
self._socket = socket.create_connection((host, self.endpoint.port), self.timeout)
|
|
self._socket.settimeout(self.timeout)
|
|
self._reader = self._socket.makefile("rb")
|
|
self._socket.sendall(self._encode({"type": "hello", "token": self.endpoint.token}))
|
|
welcome = self._read()
|
|
if welcome.get("type") != "welcome" or not isinstance(welcome.get("server_version"), str):
|
|
raise SdkRpcError("authorization_failed", "SDK 未接受运行时授权握手")
|
|
self.server_version = welcome["server_version"]
|
|
except SdkRpcError:
|
|
self._close()
|
|
raise
|
|
except (OSError, TimeoutError):
|
|
self._close()
|
|
raise SdkRpcError("connection_failed", "无法连接 SDK 运行时或握手超时") from None
|
|
|
|
def call_cli(self, argv: list[str]):
|
|
if not isinstance(argv, list) or not argv or not all(isinstance(value, str) and "\x00" not in value for value in argv):
|
|
raise SdkRpcError("invalid_request", "CLI 参数必须是非空字符串数组")
|
|
if not self._lock.acquire(timeout=self.timeout):
|
|
raise SdkRpcError("client_busy", "SDK 正在处理另一条请求")
|
|
attempted = False
|
|
try:
|
|
request_id = self._next_id
|
|
frame = self._encode({"id": request_id, "method": "wecom.exec", "params": {"argv": argv}})
|
|
if self._socket is None:
|
|
self._connect()
|
|
self._next_id += 1
|
|
# sendall may fail after partially writing. Mark uncertain before sending.
|
|
attempted = True
|
|
self._socket.sendall(frame)
|
|
response = self._read()
|
|
if type(response.get("id")) is not int or response["id"] != request_id:
|
|
raise SdkRpcError("response_mismatch", "SDK 响应 ID 与请求不一致")
|
|
if response.get("error") is not None:
|
|
# Do not include arbitrary server text: it may echo credentials/content.
|
|
# A generic remote error does not prove a business action was not executed.
|
|
raise SdkRpcError("remote_error", "SDK 返回调用错误，需要核对实际结果")
|
|
if "result" not in response:
|
|
raise SdkRpcError("invalid_frame", "SDK 响应缺少调用结果")
|
|
return response["result"]
|
|
except SdkRpcError as exc:
|
|
if attempted:
|
|
self._close()
|
|
raise SdkRpcError(exc.code, str(exc), may_have_executed=True) from None
|
|
raise
|
|
except (OSError, TimeoutError):
|
|
self._close()
|
|
raise SdkRpcError("rpc_interrupted", "SDK 请求中断或超时，需要核对实际结果",
|
|
may_have_executed=attempted) from None
|
|
finally:
|
|
self._lock.release()
|
|
|
|
|
|
def probe_sdk_runtime(*, environ=None, timeout=3.0):
|
|
"""Read-only handshake + CLI help. Never calls message send or business mutations."""
|
|
report = {
|
|
"available": False, "rpcVerified": False, "cliHelpVerified": False,
|
|
"sendVerified": False, "receiptVerified": False, "loginVerified": False,
|
|
}
|
|
try:
|
|
endpoint = RuntimeEndpoint.from_environment(environ)
|
|
report["available"] = True
|
|
with WeComSdkClient(endpoint, timeout=timeout) as client:
|
|
result = client.call_cli(["--help"])
|
|
report.update(rpcVerified=True, serverVersion=client.server_version)
|
|
# Preserve the response type only in ordinary diagnostics.
|
|
report["responseType"] = type(result).__name__
|
|
text = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
|
|
error = isinstance(result, dict) and (
|
|
result.get("error") is not None or result.get("errcode", 0) != 0
|
|
or result.get("exit_code", 0) != 0)
|
|
report["cliHelpVerified"] = not error and ("Usage:" in text or "用法" in text)
|
|
report.update(status="sdk_connected", message="SDK RPC 已连通；仍需核对消息接口、在线身份与回执。")
|
|
if not report["cliHelpVerified"]:
|
|
report.update(status="sdk_contract_unverified", message="SDK RPC 已返回，CLI 帮助格式或业务结果仍待核对。")
|
|
except SdkRpcError as exc:
|
|
report.update(status=exc.code, message=str(exc))
|
|
return report
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--output", help="Save read-only runtime inspection JSON")
|
|
args = parser.parse_args()
|
|
report = probe_sdk_runtime()
|
|
if args.output:
|
|
from wxwork_protocol_probe import export_probe_report
|
|
export_probe_report(report, args.output)
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
return 0 if report["rpcVerified"] else 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
if hasattr(sys.stdout, "reconfigure"):
|
|
sys.stdout.reconfigure(encoding="utf-8")
|
|
raise SystemExit(main())
|