#!/usr/bin/env python3
"""CVE-2026-24207 — non-destructive patch-status check for the
NVIDIA Triton SageMaker HTTP restriction bypass.

Probes the SageMaker management endpoint at GET /models with no auth
header. On a vulnerable build, the response will be 200 (or 404
"model not found"). On a patched build with --http-restricted-api set,
the response will be 403 with body "This API is restricted".

Verified empirically on 2026-05-23 against the official NVIDIA NGC
containers `nvcr.io/nvidia/tritonserver:26.02-py3` (returns 200) and
`:26.03-py3` (returns 403 with body "This API is restricted, expecting
header 'X-SM-Auth'"). See docs/root-cause.md.
"""
import sys

import requests
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


def check(target):
    """Probe the SageMaker /models endpoint without an auth header."""
    url = f"http://{target}/models"
    try:
        r = requests.get(url, verify=False, timeout=10,
                         allow_redirects=False)
    except requests.exceptions.ConnectTimeout:
        return "UNKNOWN", "connect timeout (host unreachable)"
    except requests.exceptions.ConnectionError:
        return "UNKNOWN", "connection refused (port closed or SageMaker not enabled)"
    except requests.exceptions.RequestException as e:
        return "UNKNOWN", f"{type(e).__name__}"

    if r.status_code == 403 and "restricted" in r.text.lower():
        return "PATCHED", "endpoint enforces restriction (HTTP 403)"
    if r.status_code in (200, 404):
        # 200 = listed models without auth; 404 = endpoint reachable but no
        # model named; both indicate the restriction is not enforced.
        return "VULNERABLE", f"endpoint reachable without auth (HTTP {r.status_code})"
    return "UNKNOWN", f"unexpected response (HTTP {r.status_code})"


def main():
    if len(sys.argv) != 2:
        print(f"usage: {sys.argv[0]} <host:sagemaker_port>",
              file=sys.stderr)
        print(f"  e.g. {sys.argv[0]} 10.0.0.5:8080", file=sys.stderr)
        sys.exit(2)
    verdict, detail = check(sys.argv[1])
    icon = {"PATCHED": "[+]", "VULNERABLE": "[!]", "UNKNOWN": "[?]"}[verdict]
    print(f"{icon} {sys.argv[1]}: {verdict} — {detail}")
    sys.exit({"PATCHED": 0, "VULNERABLE": 1, "UNKNOWN": 2}[verdict])


if __name__ == "__main__":
    main()