293 lines
13 KiB
Python
293 lines
13 KiB
Python
"""Exercise the PLM POST/GET protocol and save the real report without logging credentials."""
|
||
|
||
import argparse
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import math
|
||
import os
|
||
import re
|
||
import secrets
|
||
import time
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from urllib.parse import urlencode, urlsplit
|
||
|
||
import httpx
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
DEFAULT_PAYLOAD = ROOT / "docs" / "examples" / "plm-risk-request-20261009.json"
|
||
DEFAULT_URL = "http://192.168.9.28:18474/api/v1/integrations/risk-assessments/report"
|
||
LEVELS = {
|
||
"low": "低风险",
|
||
"moderately_low": "较低风险",
|
||
"medium": "中等风险",
|
||
"high": "高风险",
|
||
"very_high": "极高风险",
|
||
}
|
||
SAFE_HEADERS = (
|
||
"content-type",
|
||
"content-encoding",
|
||
"x-request-id",
|
||
"x-assessment-id",
|
||
"x-model-version",
|
||
"x-risk-score",
|
||
"x-risk-level",
|
||
"retry-after",
|
||
)
|
||
|
||
|
||
def signed_headers(method: str, url: str, body: bytes, secret: str, client_id: str, request_id: str) -> dict:
|
||
timestamp, nonce = str(int(time.time())), secrets.token_urlsafe(16)
|
||
message = timestamp.encode() + b"\n" + nonce.encode() + b"\n"
|
||
headers = {"X-Client-Id": client_id, "X-Timestamp": timestamp, "X-Nonce": nonce}
|
||
if method == "GET":
|
||
target = urlsplit(url)
|
||
message = b"GET\n" + target.path.encode() + b"\n" + target.query.encode() + b"\n" + message
|
||
else:
|
||
message += body
|
||
headers.update({"Content-Type": "application/json", "Idempotency-Key": request_id})
|
||
headers["X-Signature"] = base64.b64encode(hmac.new(secret.encode(), message, hashlib.sha256).digest()).decode()
|
||
return headers
|
||
|
||
|
||
def report_metadata(response: httpx.Response, request_id: str) -> dict:
|
||
if response.status_code != 200:
|
||
raise ValueError(f"Expected HTTP 200, received {response.status_code}")
|
||
if response.headers.get("content-type", "").split(";", 1)[0].strip().lower() != "text/markdown":
|
||
raise ValueError("Successful response must be text/markdown")
|
||
if response.headers.get("x-request-id") != request_id:
|
||
raise ValueError("X-Request-Id does not match this request")
|
||
score = float(response.headers.get("x-risk-score", "nan"))
|
||
if not math.isfinite(score) or not 0 <= score <= 100:
|
||
raise ValueError("Invalid X-Risk-Score")
|
||
level = response.headers.get("x-risk-level")
|
||
if level not in LEVELS:
|
||
raise ValueError("Invalid X-Risk-Level")
|
||
assessment_id, model = response.headers.get("x-assessment-id"), response.headers.get("x-model-version")
|
||
if not assessment_id or not model:
|
||
raise ValueError("Missing assessment ID or model version")
|
||
# httpx already decompresses gzip/deflate. Decode the resulting bytes exactly once.
|
||
if len(response.content) > 1024 * 1024:
|
||
raise ValueError("Decompressed report exceeds 1MiB")
|
||
markdown = response.content.decode("utf-8")
|
||
if not markdown.strip():
|
||
raise ValueError("Empty report")
|
||
dimensions = sum(
|
||
bool(re.match(r"^\|\s*(供应商表现|物料品类|产品品类|加工厂质量|标准符合性|节点阶段|缺陷等级)\s*\|", line))
|
||
for line in markdown.splitlines()
|
||
)
|
||
if dimensions != 7:
|
||
raise ValueError(f"Expected 7 dimension rows, received {dimensions}")
|
||
return {
|
||
"requestId": request_id,
|
||
"assessmentId": assessment_id,
|
||
"score": score,
|
||
"level": level,
|
||
"levelLabel": LEVELS[level],
|
||
"modelVersion": model,
|
||
"dimensionCount": dimensions,
|
||
"reportSha256": hashlib.sha256(response.content).hexdigest(),
|
||
}
|
||
|
||
|
||
class PLMReportTest:
|
||
def __init__(
|
||
self, client: httpx.Client, url: str, secret: str, payload: dict, output: Path, total_timeout: float = 180
|
||
):
|
||
self.client, self.url, self.secret, self.payload, self.output = client, url, secret, payload, output
|
||
self.request_id = payload["requestId"]
|
||
self.body = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||
self.query_url = url + "?" + urlencode({"requestId": self.request_id})
|
||
self.events = []
|
||
self.deadline = time.monotonic() + total_timeout
|
||
|
||
def send(self, method: str) -> httpx.Response:
|
||
remaining = self.deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
raise TimeoutError("Overall test deadline exceeded; use the saved snapshot to resume")
|
||
url = self.query_url if method == "GET" else self.url
|
||
body = b"" if method == "GET" else self.body
|
||
headers = signed_headers(method, url, body, self.secret, self.payload["sourceSystem"], self.request_id)
|
||
started = time.monotonic()
|
||
try:
|
||
response = self.client.request(
|
||
method,
|
||
url,
|
||
content=body,
|
||
headers=headers,
|
||
timeout=httpx.Timeout(min(130, remaining), connect=min(10, remaining)),
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
self.events.append(
|
||
{
|
||
"method": method,
|
||
"transportError": exc.__class__.__name__,
|
||
"seconds": round(time.monotonic() - started, 2),
|
||
}
|
||
)
|
||
raise
|
||
event = {
|
||
"method": method,
|
||
"status": response.status_code,
|
||
"seconds": round(time.monotonic() - started, 2),
|
||
"headers": {key: response.headers[key] for key in SAFE_HEADERS if key in response.headers},
|
||
}
|
||
if response.status_code != 200:
|
||
try:
|
||
event["error"] = response.json()
|
||
except ValueError:
|
||
event["error"] = {"message": "Non-JSON error response"}
|
||
self.events.append(event)
|
||
print(f"{method} HTTP {response.status_code} ({event['seconds']}s)", flush=True)
|
||
return response
|
||
|
||
def wait(self, response: httpx.Response):
|
||
try:
|
||
delay = max(1, min(60, int(response.headers.get("retry-after", "5"))))
|
||
except ValueError:
|
||
delay = 5
|
||
if time.monotonic() + delay >= self.deadline:
|
||
raise TimeoutError("Retry-After exceeds remaining test deadline")
|
||
# Short waits keep progress visible even if rate limiting requests a long delay.
|
||
print(f"Wait {delay}s before retry", flush=True)
|
||
while delay:
|
||
step = min(5, delay)
|
||
time.sleep(step)
|
||
delay -= step
|
||
|
||
def resolve(self, response: httpx.Response) -> httpx.Response:
|
||
for _ in range(40):
|
||
if response.status_code == 200:
|
||
return response
|
||
if response.status_code not in {202, 429}:
|
||
raise ValueError(f"HTTP {response.status_code}; inspect response.json for platform diagnostics")
|
||
if response.status_code == 202:
|
||
data = response.json()
|
||
if data.get("code") != 202 or (data.get("data") or {}).get("requestId") != self.request_id:
|
||
raise ValueError("Processing response does not match this request")
|
||
self.wait(response)
|
||
response = self.send("GET")
|
||
raise TimeoutError("Polling limit reached; use the saved snapshot to resume")
|
||
|
||
def run(self) -> dict:
|
||
self.output.mkdir(parents=True, exist_ok=True)
|
||
(self.output / "request.json").write_bytes(self.body)
|
||
checks = []
|
||
first = self.send("GET")
|
||
while first.status_code == 429:
|
||
self.wait(first)
|
||
first = self.send("GET")
|
||
if first.status_code in {404, 503}:
|
||
try:
|
||
first = self.send("POST")
|
||
except httpx.HTTPError:
|
||
# A timed-out POST may already have persisted the request.
|
||
first = self.send("GET")
|
||
first = self.resolve(first)
|
||
metadata = report_metadata(first, self.request_id)
|
||
(self.output / "report.md").write_bytes(first.content)
|
||
checks.append("报告为 UTF-8 Markdown，编号、分数、等级及七个维度校验通过")
|
||
repeated = self.send("POST")
|
||
if repeated.status_code == 429:
|
||
self.wait(repeated)
|
||
repeated = self.send("POST")
|
||
repeated = self.resolve(repeated)
|
||
if report_metadata(repeated, self.request_id) != metadata or repeated.content != first.content:
|
||
raise ValueError("Repeated POST differs from the original report")
|
||
checks.append("同编号同内容 POST 返回原报告，幂等校验通过")
|
||
queried = self.resolve(self.send("GET"))
|
||
if report_metadata(queried, self.request_id) != metadata or queried.content != first.content:
|
||
raise ValueError("GET report differs from the original report")
|
||
checks.append("GET 报告正文、评估 ID、分数、等级及模型版本一致")
|
||
return {
|
||
"success": True,
|
||
"endpoint": self.url,
|
||
**metadata,
|
||
"checks": checks,
|
||
"requestSha256": hashlib.sha256(self.body).hexdigest(),
|
||
"events": self.events,
|
||
}
|
||
|
||
|
||
def load_secret(client_id: str, use_local_config: bool) -> str:
|
||
secret = os.environ.get("RISK_CLIENT_SECRET") or os.environ.get("RISK_INTEGRATION_CLIENT_SECRET")
|
||
if not secret and use_local_config:
|
||
# Explicit opt-in for running this tool in the risk backend workspace.
|
||
from dotenv import dotenv_values
|
||
|
||
raw = os.environ.get("RISK_INTEGRATION_CLIENTS") or dotenv_values(ROOT / ".env").get("RISK_INTEGRATION_CLIENTS")
|
||
secret = json.loads(raw or "{}").get(client_id)
|
||
if not isinstance(secret, str) or len(secret) < 32:
|
||
raise ValueError("Set RISK_CLIENT_SECRET or use --local-config to read this backend's .env")
|
||
return secret
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--url", default=DEFAULT_URL)
|
||
parser.add_argument("--body-file", type=Path, default=DEFAULT_PAYLOAD)
|
||
parser.add_argument("--output-dir", type=Path)
|
||
parser.add_argument("--local-config", action="store_true", help="Read secret from backend .env without printing it")
|
||
parser.add_argument(
|
||
"--new-request", action="store_true", help="Use a unique requestId to generate a new assessment"
|
||
)
|
||
args = parser.parse_args()
|
||
stamp = datetime.now(timezone(timedelta(hours=8))).strftime("%Y%m%d-%H%M%S")
|
||
output = args.output_dir or ROOT / ".runtime-uploads" / "integration-tests" / (stamp + "-" + secrets.token_hex(3))
|
||
output.mkdir(parents=True, exist_ok=True)
|
||
runner = None
|
||
try:
|
||
target = urlsplit(args.url)
|
||
if target.scheme not in {"http", "https"} or target.query or target.fragment or target.username:
|
||
raise ValueError("Endpoint must be HTTP(S), without credentials, query or fragment")
|
||
payload = json.loads(args.body_file.read_text(encoding="utf-8-sig"))
|
||
if args.new_request:
|
||
payload["requestId"] = f"PLM-TEST-{payload['assessment']['subjectId']}-{stamp}-{secrets.token_hex(4)}"
|
||
secret = load_secret(payload["sourceSystem"], args.local_config)
|
||
with httpx.Client(follow_redirects=False, trust_env=False) as client:
|
||
runner = PLMReportTest(client, args.url, secret, payload, output)
|
||
result = runner.run()
|
||
except (ValueError, KeyError, OSError, TimeoutError, httpx.HTTPError) as exc:
|
||
result = {
|
||
"success": False,
|
||
"errorType": exc.__class__.__name__,
|
||
"message": str(exc)
|
||
if not isinstance(exc, (httpx.HTTPError, OSError))
|
||
else "Transport or file operation failed",
|
||
"events": runner.events if runner else [],
|
||
}
|
||
(output / "response.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
if result["success"]:
|
||
summary = [
|
||
"# PLM 风险接口测试结果",
|
||
"",
|
||
f"- 请求编号：{result['requestId']}",
|
||
f"- 平台评估 ID：{result['assessmentId']}",
|
||
f"- 综合安全得分：{result['score']:.2f}",
|
||
f"- 风险等级：{result['levelLabel']}",
|
||
f"- 模型版本：{result['modelVersion']}",
|
||
f"- 维度数量：{result['dimensionCount']}",
|
||
"",
|
||
"报告：[report.md](report.md)",
|
||
"",
|
||
]
|
||
summary.extend(f"- {check}" for check in result["checks"])
|
||
summary.extend(
|
||
[
|
||
"",
|
||
"本次测试沿用接口现有评分规则；HTTP 200 不代表全部维度拥有真实数据。",
|
||
"当品类为护发素时，当前规则可能无法归类，报告会注明品类维度默认分。",
|
||
"",
|
||
]
|
||
)
|
||
(output / "summary.md").write_text("\n".join(summary), encoding="utf-8")
|
||
print(f"Score: {result['score']:.2f}; level: {result['level']}; assessment: {result['assessmentId']}")
|
||
print(f"Outputs: {output.resolve()}")
|
||
return 0 if result["success"] else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|