#!/usr/bin/env python3
"""Smoke-check the Keycloak realms: each realm's OIDC login works (password grant)
and returns its expected identifying claim. The medewerker realm additionally enforces
MFA (S-15c), so its login must be refused without a TOTP code. Stdlib only.
Exits non-zero on failure.
"""
import base64, hashlib, hmac, json, struct, sys, time, urllib.error, urllib.parse, urllib.request

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

# Fixture TOTP secret seeded into every medewerker in infra/keycloak/realms/medewerker-realm.json.
# Keycloak HMACs the raw secret bytes, so no base32 decoding is involved.
OTP_SECRET = b"BIGMEDEWERKEROTPSEED"

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


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


def totp(secret=OTP_SECRET, period=30, digits=6):
    """RFC 6238 code: HMAC-SHA1 over the 30-second counter, dynamically truncated."""
    mac = hmac.new(secret, struct.pack(">Q", int(time.time()) // period), hashlib.sha1).digest()
    o = mac[-1] & 0x0F
    return str((struct.unpack(">I", mac[o:o + 4])[0] & 0x7FFFFFFF) % 10 ** digits).zfill(digits)


def grant(realm, user, **extra):
    data = urllib.parse.urlencode({
        "grant_type": "password", "client_id": CLIENT,
        "username": user, "password": PWD, "scope": "openid", **extra,
    }).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 second_factor_refused(realm, user):
    """The password alone must not yield a token on an MFA-enforced realm."""
    try:
        grant(realm, user)
    except urllib.error.HTTPError as e:
        return e.code in (400, 401)
    return False


def main():
    ok = True
    for realm, user, claim, expect, mfa in CHECKS:
        extra = {}
        if mfa:
            refused = second_factor_refused(realm, user)
            ok = ok and refused
            print(f"{realm:12} {user:18} password-only login refused "
                  f"[{'OK' if refused else 'MFA NOT ENFORCED'}]")
            extra = {"totp": totp()}
        try:
            at = decode(grant(realm, user, **extra)["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__":
    # `check_realms.py otp` prints a current code for the fixture secret — what a human demoing
    # the medewerker portals types at Keycloak's OTP prompt (docs/runbooks/keycloak.md).
    if len(sys.argv) > 1 and sys.argv[1] == "otp":
        print(totp())
    else:
        main()