query-orchestration/scripts/manual_tests/test_two_pool_auth.py

#!/usr/bin/env python3
"""End-to-end manual test for the two-pool service-account auth path.

Validates plans/api.key.support.cognito.two-pool.plan.3.claude.md
against a running queryAPI configured with both pool A
(COGNITO_USER_POOL_ID) and pool B (COGNITO_SVC_USER_POOL_ID), and
auth enabled.

What this script does (top to bottom):
  Step 1. Sanity-check queryAPI is reachable (GET /health, no auth).
  Step 2. InitiateAuth against pool B with the pre-baked test
          service-account credentials (svc-test-admin-1). Returns
          AccessToken / IdToken / RefreshToken.
  Step 3. Decode the AccessToken locally (no signature check) to
          confirm sub, iss, token_use match the values from the
          provisioning log.
  Step 4. Call GET /admin/eula with Authorization: Bearer
          <access_token>. /admin/eula is a super_admin-only
          endpoint, so a 200 proves: (a) queryAPI's multi-issuer
          verifier accepted the pool B token; (b) the service
          account has super_admin in Permit; (c) the env-check
          dispatched to pool B.
  Step 5. Call GET /admin/eula/agreements with the same token.
          Second super_admin-only endpoint — covers the "pick at
          least 2" requirement from the goal statement.
  Step 6. Print PASS/FAIL summary.

Zero third-party deps — stdlib only, like
cmd/auth_related/service.accounts/test.scripts/test.service.account.auth.py.

Prerequisites:
  - queryAPI is reachable at --base-url (default http://localhost:8080).
  - queryAPI was started with both COGNITO_USER_POOL_ID and
    COGNITO_SVC_USER_POOL_ID configured.
  - DISABLE_AUTH is NOT set (we want real Cognito JWT validation).
  - The test service account already has super_admin in the Permit
    instance backing the target queryAPI. This script does not
    provision Permit — that is an environment-setup concern.
  - SVC_TEST_PASSWORD is supplied via env (the repo-root .env.example
    documents the variable). Any .env between this file and the
    filesystem root is auto-loaded. Everything else (pool ID, app
    client ID, username, sub, region) is hardcoded below — public
    identifiers, not secrets — so the script "just works" without
    further shell setup.
  - INSECURE_TLS=1 disables queryAPI TLS cert verification (Cognito
    is left alone). Needed when queryAPI is fronted by an ALB with
    a self-signed cert. Off by default.

Usage:
  python3 scripts/manual_tests/test_two_pool_auth.py
  python3 scripts/manual_tests/test_two_pool_auth.py --base-url http://localhost:8080
  INSECURE_TLS=1 python3 scripts/manual_tests/test_two_pool_auth.py --base-url https://...
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import ssl
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict

HTTP_TIMEOUT = 15


def _load_dotenv() -> None:
    """Walk from this file up to the filesystem root, loading every
    .env encountered. First-wins precedence: closer (script-adjacent)
    .envs win over deeper (repo-root) ones, and existing shell vars
    win over both. Loading every .env on the way up — rather than
    stopping at the first hit — means a script-local .env cannot
    hide the repo-root .env where SVC_TEST_PASSWORD lives.
    """
    here = Path(__file__).resolve()
    for parent in [here.parent, *here.parents]:
        candidate = parent / ".env"
        if not candidate.is_file():
            continue
        for raw in candidate.read_text().splitlines():
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            if "=" not in line:
                continue
            k, v = line.split("=", 1)
            k = k.strip()
            v = v.strip().strip('"').strip("'")
            if k and k not in os.environ:
                os.environ[k] = v


def _required(name: str) -> str:
    v = os.environ.get(name, "").strip()
    if not v:
        print(f"FAIL: required env var {name} is unset. See .env.example.")
        sys.exit(2)
    return v


# Hardcoded service-account identifiers. Pool ID, app client ID,
# username, and sub are public values, not secrets — hardcoding them
# makes the script share-able: anyone with SVC_TEST_PASSWORD in their
# .env can run it with no other prerequisites.
USER_POOL_ID = "us-east-2_Qom1i9qIG"
APP_CLIENT_ID = "6r7chkvjotgs2bsfbd5ibaaft7"
USERNAME = "svc-test-admin-1"
EXPECTED_SUB = "c16b0520-80e1-7076-eeb2-1e520d81e781"

# Region is encoded in the pool ID prefix. Derive it rather than
# trust the shell — devbox commonly exports AWS_REGION=us-east-1,
# which routes InitiateAuth to a region where this pool's app client
# does not exist and yields ResourceNotFoundException.
AWS_REGION = USER_POOL_ID.split("_", 1)[0]

_load_dotenv()

# Strip LocalStack endpoint overrides for hygiene. urllib (used
# below) does not honor AWS_ENDPOINT_URL, but the env stays clean for
# anything spawned downstream from this process.
for key in ("AWS_ENDPOINT_URL", "AWS_ENDPOINT_URL_COGNITO_IDP"):
    os.environ.pop(key, None)

PASSWORD = _required("SVC_TEST_PASSWORD")

# INSECURE_TLS=1 disables certificate verification for queryAPI calls
# only (Cognito is left alone — its cert is real). Needed when queryAPI
# is fronted by an ALB with a self-signed cert. Off by default so the
# bypass cannot slip in silently.
INSECURE_TLS = os.environ.get("INSECURE_TLS", "").strip().lower() in ("1", "true", "yes", "y")
_TLS_CTX = ssl._create_unverified_context() if INSECURE_TLS else None

ISSUER_URL = f"https://cognito-idp.{AWS_REGION}.amazonaws.com/{USER_POOL_ID}"
COGNITO_IDP_URL = f"https://cognito-idp.{AWS_REGION}.amazonaws.com/"


def section(title: str) -> None:
    print()
    print("=" * 72)
    print(title)
    print("=" * 72)


def cognito_call(operation: str, body: Dict[str, Any]) -> Dict[str, Any]:
    """POST a JSON-over-HTTPS request to Cognito Identity Provider.

    Matches the wire format used by boto3 internally so we don't
    need any AWS credentials for the unauthenticated InitiateAuth
    operation.
    """
    payload = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        COGNITO_IDP_URL,
        data=payload,
        method="POST",
        headers={
            "Content-Type": "application/x-amz-json-1.1",
            "X-Amz-Target": f"AWSCognitoIdentityProviderService.{operation}",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace")
        try:
            err = json.loads(body_text)
            raise RuntimeError(
                f"Cognito {operation}: HTTP {e.code} {err.get('__type', '?')}: "
                f"{err.get('message', body_text)}"
            ) from None
        except json.JSONDecodeError:
            raise RuntimeError(f"Cognito {operation}: HTTP {e.code} {body_text}") from None


def decode_jwt_payload(token: str) -> Dict[str, Any]:
    """Base64url-decode the payload segment of a JWT (no signature check)."""
    seg = token.split(".")[1]
    seg += "=" * (-len(seg) % 4)
    return json.loads(base64.urlsafe_b64decode(seg))


def queryapi_get(base_url: str, path: str, access_token: str = "") -> tuple[int, str]:
    """GET <base_url><path> with optional bearer auth. Returns (status, body)."""
    url = base_url.rstrip("/") + path
    headers = {}
    if access_token:
        headers["Authorization"] = "Bearer " + access_token
    req = urllib.request.Request(url, headers=headers, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT, context=_TLS_CTX) as r:
            return r.status, r.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        return e.code, body


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default="http://localhost:8080")
    args = parser.parse_args()
    base = args.base_url

    if INSECURE_TLS:
        print("WARNING: INSECURE_TLS=1 — queryAPI TLS verification is DISABLED.")

    section("Step 1: queryAPI health")
    code, body = queryapi_get(base, "/health")
    print(f"GET {base}/health -> {code}")
    if code != 200:
        print(f"FAIL: queryAPI not healthy at {base}: {body}")
        return 1

    section("Step 2: InitiateAuth against pool B")
    try:
        resp = cognito_call(
            "InitiateAuth",
            {
                "AuthFlow": "USER_PASSWORD_AUTH",
                "ClientId": APP_CLIENT_ID,
                "AuthParameters": {"USERNAME": USERNAME, "PASSWORD": PASSWORD},
            },
        )
    except RuntimeError as e:
        print(f"FAIL: {e}")
        return 1
    if "AuthenticationResult" not in resp:
        print(f"FAIL: unexpected InitiateAuth response: {resp}")
        return 1
    access_token = resp["AuthenticationResult"]["AccessToken"]
    print(f"  AccessToken length : {len(access_token)} chars")
    print(f"  ExpiresIn          : {resp['AuthenticationResult']['ExpiresIn']}s")

    section("Step 3: AccessToken claims")
    claims = decode_jwt_payload(access_token)
    print(json.dumps(claims, indent=2, sort_keys=True))
    failures = []
    if claims.get("iss") != ISSUER_URL:
        failures.append(f"iss mismatch (got {claims.get('iss')}, want {ISSUER_URL})")
    if claims.get("sub") != EXPECTED_SUB:
        failures.append(f"sub mismatch (got {claims.get('sub')}, want {EXPECTED_SUB})")
    if claims.get("token_use") != "access":
        failures.append(f"token_use != access (got {claims.get('token_use')})")
    if failures:
        for f in failures:
            print(f"FAIL: {f}")
        return 1

    # We pick two super_admin-only endpoints that read from Postgres
    # only (no downstream AWS calls). /admin/users would also work
    # against real AWS but hits Cognito ListUsers, which localstack
    # community edition does not implement, so we use the EULA admin
    # endpoints instead — both are gated on the same super_admin role
    # in Permit.

    section("Step 4: GET /admin/eula (super_admin only)")
    code, body = queryapi_get(base, "/admin/eula", access_token)
    print(f"GET {base}/admin/eula -> {code}")
    if code != 200:
        print("FAIL: /admin/eula did not return 200 — service account either")
        print("  did not authenticate, or is not registered with super_admin in Permit.")
        print(f"  body: {body[:400]}")
        return 1
    print(f"  body preview: {body[:200]}")

    section("Step 5: GET /admin/eula/agreements (super_admin only)")
    code, body = queryapi_get(base, "/admin/eula/agreements", access_token)
    print(f"GET {base}/admin/eula/agreements -> {code}")
    if code != 200:
        print("FAIL: /admin/eula/agreements did not return 200.")
        print(f"  body: {body[:400]}")
        return 1
    print(f"  body preview: {body[:200]}")

    section("Result")
    print("PASS — two-pool service-account auth is wired end-to-end.")
    print(
        f"  Service account {USERNAME} (sub={EXPECTED_SUB})\n"
        "  authenticated against pool B and successfully called two\n"
        "  super_admin-only endpoints on queryAPI."
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())