69 lines
2.9 KiB
Python
69 lines
2.9 KiB
Python
"""POST an assessment or GET a persisted report using the documented HMAC protocol."""
|
|
|
|
import argparse
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import os
|
|
import secrets
|
|
import time
|
|
from pathlib import Path
|
|
from urllib.parse import urlencode, 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")
|
|
inputs = parser.add_mutually_exclusive_group(required=True)
|
|
inputs.add_argument("--body-file", type=Path, help="UTF-8 JSON assessment to POST")
|
|
inputs.add_argument("--request-id", help="Existing requestId to GET")
|
|
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}
|
|
if args.body_file:
|
|
import json
|
|
|
|
body = args.body_file.read_bytes()
|
|
data = json.loads(body)
|
|
headers.update({"Content-Type": "application/json", "Idempotency-Key": data["requestId"]})
|
|
message += body
|
|
method, target = "POST", args.url
|
|
else:
|
|
body = b""
|
|
query = urlencode({"requestId": args.request_id})
|
|
message = b"GET\n" + url.path.encode() + b"\n" + query.encode() + b"\n" + message
|
|
method, target = "GET", args.url + "?" + query
|
|
headers["X-Signature"] = base64.b64encode(hmac.new(secret.encode(), message, hashlib.sha256).digest()).decode()
|
|
try:
|
|
response = httpx.request(method, target, content=body, headers=headers, timeout=120)
|
|
except httpx.HTTPError as exc:
|
|
print(f"Request did not complete ({exc.__class__.__name__}); query the same requestId before retrying.")
|
|
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
|
|
args.output.write_text(response.text, encoding="utf-8")
|
|
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())
|