#!/usr/bin/env python3
"""Log in with HeartBank® from Python, with Authlib (served at /developers/docs/python_authlib.py; test/e2e.test.mjs runs this exact file against login).
Authlib builds PKCE and the private_key_jwt assertion; it has no PAR or DPoP helper, so those are written out here.
Reads ISSUER, CLIENT_ID, REDIRECT_URI, CLIENT_JWK (private, JSON) from the environment. Prints the authorization URL, reads the
callback URL back on stdin, verifies the ID token, and prints its claims as JSON."""
import json, os, secrets, sys, time, uuid
from urllib.parse import parse_qs, urlparse
import requests
from authlib.common.security import generate_token
from authlib.jose import JsonWebKey, jwt
from authlib.oauth2.rfc7636 import create_s256_code_challenge

ISSUER, CLIENT_ID, REDIRECT = os.environ["ISSUER"], os.environ["CLIENT_ID"], os.environ["REDIRECT_URI"]
BASE = os.environ.get("BASE_URL", ISSUER)  # only the test harness sets this; in real use it is the issuer
HEADERS = json.loads(os.environ.get("EXTRA_HEADERS", "{}"))
client_key = json.loads(os.environ["CLIENT_JWK"])
here = lambda url: url.replace(ISSUER, BASE)

meta = requests.get(here(f"{ISSUER}/.well-known/openid-configuration"), headers=HEADERS).json()

def client_assertion():
    now = int(time.time())
    return jwt.encode({"alg": "ES256", "kid": client_key["kid"]},
                      {"iss": CLIENT_ID, "sub": CLIENT_ID, "aud": meta["issuer"], "jti": uuid.uuid4().hex, "iat": now, "exp": now + 60},
                      client_key).decode()

def auth_fields():
    return {"client_id": CLIENT_ID, "client_assertion": client_assertion(),
            "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"}

# PAR: the request goes server to server, the browser carries only a reference.
verifier, nonce, state = generate_token(64), secrets.token_urlsafe(16), secrets.token_urlsafe(16)
par = requests.post(here(meta["pushed_authorization_request_endpoint"]), headers=HEADERS, data={
    **auth_fields(), "response_type": "code", "scope": "openid", "redirect_uri": REDIRECT, "nonce": nonce, "state": state,
    "code_challenge": create_s256_code_challenge(verifier), "code_challenge_method": "S256"})
par.raise_for_status()
print(f'{meta["authorization_endpoint"]}?client_id={CLIENT_ID}&request_uri={requests.utils.quote(par.json()["request_uri"], safe="")}', flush=True)

# ...the person logs in, and the browser comes back to REDIRECT with code, state and iss.
back = parse_qs(urlparse(sys.stdin.readline().strip()).query)
assert back["state"][0] == state and back["iss"][0] == meta["issuer"], "mix-up or forgery"

# DPoP: a proof signed by a key this app holds, naming this request.
dpop_key = JsonWebKey.generate_key("EC", "P-256", is_private=True)
dpop = jwt.encode({"alg": "ES256", "typ": "dpop+jwt", "jwk": dpop_key.as_dict(is_private=False)},
                  {"htm": "POST", "htu": meta["token_endpoint"], "jti": uuid.uuid4().hex, "iat": int(time.time())}, dpop_key).decode()
tok = requests.post(here(meta["token_endpoint"]), headers={**HEADERS, "DPoP": dpop}, data={
    **auth_fields(), "grant_type": "authorization_code", "code": back["code"][0], "redirect_uri": REDIRECT, "code_verifier": verifier})
tok.raise_for_status()

keys = JsonWebKey.import_key_set(requests.get(here(meta["jwks_uri"]), headers=HEADERS).json())
claims = jwt.decode(tok.json()["id_token"], keys, claims_options={"iss": {"essential": True, "value": meta["issuer"]},
                                                                  "aud": {"essential": True, "value": CLIENT_ID},
                                                                  "nonce": {"essential": True, "value": nonce}})
claims.validate()
print(json.dumps({"sub": claims["sub"], "amr": claims["amr"], "token_type": tok.json()["token_type"]}), flush=True)