#!/usr/bin/env python3
"""Smoke-check the Keycloak realms: each realm's OIDC login works (password grant)
and returns its expected identifying claim. Stdlib only. Exits non-zero on failure.
"""
import base64, json, sys, urllib.error, urllib.parse, urllib.request

BASE = "http://localhost:8180"
CLIENT = "big-portal"
PWD = "test123"

# realm, user, claim ("__roles__" => check realm_access.roles), expected-contains
CHECKS = [
    ("digid", "jan-burger", "bsn", "123456782"),
    ("eherkenning", "acme-ondernemer", "kvk", "12345678"),
    ("eidas", "pierre-dupont", "eidas_id", "FR/NL"),
    ("medewerker", "merel-behandelaar", "__roles__", "behandelaar"),
]


def decode(jwt):
    p = jwt.split(".")[1]
    p += "=" * (-len(p) % 4)
    return json.loads(base64.urlsafe_b64decode(p))


def grant(realm, user):
    data = urllib.parse.urlencode({
        "grant_type": "password", "client_id": CLIENT,
        "username": user, "password": PWD, "scope": "openid",
    }).encode()
    req = urllib.request.Request(
        f"{BASE}/realms/{realm}/protocol/openid-connect/token", data=data,
        headers={"Content-Type": "application/x-www-form-urlencoded"})
    with urllib.request.urlopen(req, timeout=20) as r:
        return json.loads(r.read())


def main():
    ok = True
    for realm, user, claim, expect in CHECKS:
        try:
            at = decode(grant(realm, user)["access_token"])
            if claim == "__roles__":
                val = at.get("realm_access", {}).get("roles", [])
                good = expect in val
            else:
                val = at.get(claim)
                good = val is not None and expect in str(val)
            print(f"{realm:12} {user:18} login OK | {claim} = {val} "
                  f"[{'OK' if good else 'UNEXPECTED'}]")
            ok = ok and good
        except urllib.error.HTTPError as e:
            ok = False
            print(f"{realm:12} {user:18} LOGIN FAILED {e.code}: {e.read()[:200]!r}")
    print("keycloak smoke OK" if ok else "keycloak smoke FAILED")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()