dieya_risk_prediction_backend/scripts/call_external_risk_report.py

63 lines
2.6 KiB
Python

"""POST an assessment, print its JSON response, and save data.reportMarkdown."""
import argparse
import base64
import hashlib
import hmac
import json
import os
import secrets
import time
from pathlib import Path
from urllib.parse import urlsplit
import httpx
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--url", required=True, help="Full /api/v1/integrations/risk-assessments/report URL")
parser.add_argument("--client-id", default="dieya-plm")
parser.add_argument(
"--body-file", type=Path, required=True, help="UTF-8 JSON assessment; repost same bytes to retry"
)
parser.add_argument("--output", type=Path, default=Path("risk-report.md"))
args = parser.parse_args()
secret = os.environ.get("RISK_INTEGRATION_CLIENT_SECRET", "")
if len(secret) < 32:
parser.error("Set RISK_INTEGRATION_CLIENT_SECRET to the system's secret (at least 32 characters)")
url = urlsplit(args.url)
if url.scheme not in {"https", "http"} or url.query or url.fragment:
parser.error("--url must be an HTTP(S) endpoint without query or fragment")
timestamp, nonce = str(int(time.time())), secrets.token_urlsafe(16)
message = timestamp.encode() + b"\n" + nonce.encode() + b"\n"
headers = {"X-Client-Id": args.client_id, "X-Timestamp": timestamp, "X-Nonce": nonce}
body = args.body_file.read_bytes()
data = json.loads(body)
headers.update({"Content-Type": "application/json", "Idempotency-Key": data["requestId"]})
message += body
headers["X-Signature"] = base64.b64encode(hmac.new(secret.encode(), message, hashlib.sha256).digest()).decode()
try:
response = httpx.post(args.url, content=body, headers=headers, timeout=120)
except httpx.HTTPError as exc:
print(f"Request did not complete ({exc.__class__.__name__}); repost the same requestId and original body.")
return 1
if response.status_code != 200:
print(f"HTTP {response.status_code}: {response.text}")
if response.headers.get("Retry-After"):
print(f"Retry-After: {response.headers['Retry-After']} seconds")
return 1
result = response.json()
args.output.write_bytes(result["data"]["reportMarkdown"].encode("utf-8"))
print(json.dumps(result, ensure_ascii=False))
print(f"Saved {args.output.resolve()}")
print(
f"Assessment: {response.headers.get('X-Assessment-Id')}, "
f"score: {response.headers.get('X-Risk-Score')}, level: {response.headers.get('X-Risk-Level')}"
)
return 0
if __name__ == "__main__":
raise SystemExit(main())