"""Sign in with Aurora Account: a complete example in Python.

A small web server that signs a person in as https://aurora-ws.com/developers
describes: the authorization code flow with PKCE, a private_key_jwt client assertion
for the token request, the ID token checked against Aurora Account's key set, and
userinfo. It keeps the sign-ins under way in memory: an example to read and run, not
a production server. README.md says how to run it. MIT licence: LICENSE.
"""
import base64
import hashlib
import hmac
import html
import json
import os
import secrets
import ssl
import sys
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlencode, urlsplit

import jwt  # PyJWT, with its cryptography extra
import requests

# Name your own service here: Cloudflare in front of Aurora Account turns away some
# libraries' default User-Agent (section 10).
USER_AGENT = "AuroraSignInExample/1.0 (Python; +https://aurora-ws.com/developers)"
SCOPE = "openid email profile"
PENDING_FOR = 600  # seconds a sign-in may take, from /login to /callback
JWT_BEARER = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"


# [START assertion]
def make_assertion(client_id: str, issuer: str, private_pem: bytes) -> str:
    """The client assertion that proves a token request is yours (section 5): signed
    with your private key, valid for a minute, and accepted once."""
    now = int(time.time())

    claims = {
        "iss": client_id,  # you
        "sub": client_id,  # you, again
        "aud": issuer,  # the issuer itself, as one string
        "iat": now,
        "exp": now + 60,  # a minute is plenty
        "jti": str(uuid.uuid4()),  # new every time: an assertion works once
    }

    return jwt.encode(claims, private_pem, algorithm="RS256")
# [END assertion]


# [START id-token]
def verify_id_token(id_token: str, keys: jwt.PyJWKClient, issuer: str, client_id: str,
                    nonce: str) -> dict:
    """The ID token's claims, once it is shown to be Aurora Account's, for you, and for
    this sign-in (section 6). Raises jwt.InvalidTokenError otherwise."""
    # The signature, by the key the token's header names, RS256 only.
    key = keys.get_signing_key_from_jwt(id_token)

    # Then iss, aud and exp. strict_aud refuses an aud that is a list, and the leeway
    # forgives a clock a little off Aurora Account's, as Aurora Account forgives yours.
    claims = jwt.decode(id_token, key.key, algorithms=["RS256"], issuer=issuer,
                        audience=client_id, leeway=30,
                        options={"require": ["iss", "aud", "exp", "sub", "nonce"],
                                 "strict_aud": True})

    # And the nonce this sign-in sent.
    if not hmac.compare_digest(str(claims["nonce"]).encode(), nonce.encode()):
        raise jwt.InvalidTokenError("The ID token's nonce is not the one sent.")

    return claims
# [END id-token]


def setting(name: str) -> str:
    if not os.environ.get(name):
        sys.exit(f"Set {name}: README.md says how.")
    return os.environ[name]


def cookie_of(header: str, name: str) -> str:
    """One cookie from a Cookie header, read by hand: http.cookies gives up on a whole
    header for one cookie it does not like, and on localhost every dev server's
    cookies arrive here."""
    for part in header.split(";"):
        key, _, value = part.strip().partition("=")
        if key == name:
            return value
    return ""


def page(title: str, body: str) -> bytes:
    return ('<!doctype html><html lang="en"><meta charset="utf-8">'
            '<meta name="viewport" content="width=device-width, initial-scale=1">'
            f"<title>{title}</title><h1>{title}</h1>{body}").encode()


def shown(value) -> str:
    return f"<pre>{html.escape(json.dumps(value, indent=2, ensure_ascii=False))}</pre>"


class Example:
    """The configuration, Aurora Account's endpoints, and the sign-ins under way."""

    def __init__(self):
        self.issuer = os.environ.get("AURORA_ISSUER", "https://account.aurora-ws.com/o")
        self.client_id = setting("CLIENT_ID")
        self.redirect_uri = setting("REDIRECT_URI")

        with open(setting("PRIVATE_KEY_PATH"), "rb") as file:
            self.private_pem = file.read()
        try:  # a key that cannot sign fails now, not at the first sign-in
            make_assertion(self.client_id, self.issuer, self.private_pem)
        except Exception as error:
            sys.exit(f"PRIVATE_KEY_PATH holds no RSA private key that signs: {error}")

        self.http = requests.Session()
        self.http.headers["User-Agent"] = USER_AGENT

        # [START discovery]
        # The endpoints come from the discovery document, which names the issuer.
        found = self.http.get(f"{self.issuer}/.well-known/openid-configuration",
                              timeout=10)
        found.raise_for_status()
        self.meta = found.json()

        # It must name the issuer you configured, exactly.
        if self.meta["issuer"] != self.issuer:
            sys.exit(f"The discovery document names the issuer {self.meta['issuer']}.")

        # The key set: cached, and fetched again for a token whose key it does not
        # hold (section 6).
        self.keys = jwt.PyJWKClient(self.meta["jwks_uri"],
                                    headers={"User-Agent": USER_AGENT})
        # [END discovery]

        self.pending = {}  # state -> (code_verifier, nonce, when it started)

    def cookie(self, state: str, max_age: int) -> str:
        secure = "; Secure" if self.redirect_uri.startswith("https://") else ""
        return (f"aurora_state={state}; Max-Age={max_age}; Path=/; HttpOnly; "
                f"SameSite=Lax{secure}")


class Handler(BaseHTTPRequestHandler):
    timeout = 10  # seconds a client may stay silent

    def handle(self):
        # Over HTTPS, the handshake happens here, in this request's own thread and under
        # its timeout, so that one silent client cannot hold up the others.
        if isinstance(self.connection, ssl.SSLSocket):
            try:
                self.connection.do_handshake()
            except OSError as error:
                return self.log_error("TLS handshake failed: %s", error)
        super().handle()

    def log_request(self, code="-", size="-"):
        # The request line without its query: the callback's carries the code and state.
        line = self.requestline.split("?")[0]
        self.log_message('"%s" %s', line, getattr(code, "value", code))

    def do_GET(self):
        url = urlsplit(self.path)
        query = {name: values[0] for name, values in parse_qs(url.query).items()}
        routes = {"/": self.home, "/login": self.login, "/callback": self.callback}
        if url.path not in routes:
            return self.answer(404, page("Not found", ""))
        routes[url.path](self.server.example, query)

    def answer(self, status: int, body: bytes, location: str = "", cookie: str = ""):
        self.send_response(status)
        headers = {"Location": location, "Set-Cookie": cookie,
                   "Content-Type": "text/html; charset=utf-8",
                   "Content-Length": str(len(body)), "Cache-Control": "no-store",
                   "Content-Security-Policy": "default-src 'none'"}
        for name, value in headers.items():
            if value:
                self.send_header(name, value)
        self.end_headers()
        self.wfile.write(body)

    def fail(self, example: Example, message: str):
        body = f'<p>{html.escape(message)}</p><p><a href="/login">Sign in again</a></p>'
        self.answer(400, page("Not signed in", body), cookie=example.cookie("", 0))

    def home(self, example: Example, query: dict):
        body = '<p><a href="/login">Sign in with Aurora Account</a></p>'
        self.answer(200, page("Example", body))

    def login(self, example: Example, query: dict):
        now = time.time()
        for state, (_, _, started) in list(example.pending.items()):
            if now - started > PENDING_FOR:
                example.pending.pop(state, None)

        # [START authorize]
        # Three fresh random values for this sign-in, kept on the server.
        state, nonce, verifier = (secrets.token_urlsafe(32) for _ in range(3))
        example.pending[state] = (verifier, nonce, now)

        # PKCE: the challenge is the verifier's SHA-256, base64url without padding.
        challenge = hashlib.sha256(verifier.encode("ascii")).digest()
        code_challenge = base64.urlsafe_b64encode(challenge).rstrip(b"=").decode()

        location = example.meta["authorization_endpoint"] + "?" + urlencode({
            "response_type": "code",
            "client_id": example.client_id,
            "redirect_uri": example.redirect_uri,
            "scope": SCOPE,
            "state": state,
            "nonce": nonce,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        })

        # The state is tied to this browser by a cookie, checked on the way back.
        cookie = example.cookie(state, PENDING_FOR)
        self.answer(302, b"", location=location, cookie=cookie)
        # [END authorize]

    def callback(self, example: Example, query: dict):
        # [START callback]
        # The state must be the one this browser was given, and still pending.
        state = query.get("state", "")
        kept = cookie_of(self.headers.get("Cookie", ""), "aurora_state")
        pending = None
        if kept and hmac.compare_digest(kept.encode(), state.encode()):
            pending = example.pending.pop(state, None)  # used once
        if pending is None or time.time() - pending[2] > PENDING_FOR:
            return self.fail(example, "This sign-in did not start here, or expired.")

        # iss comes back on every answer, errors too (RFC 9207).
        if query.get("iss") != example.issuer:
            return self.fail(example, "The answer does not name the issuer.")

        # access_denied: the person pressed Cancel (section 9).
        if "error" in query:
            return self.fail(example, f"Aurora Account answered {query['error']}.")
        # [END callback]

        verifier, nonce, _ = pending
        try:
            # [START token]
            # A new assertion for every request: each one is accepted once.
            assertion = make_assertion(example.client_id, example.issuer,
                                       example.private_pem)

            # The code, the PKCE verifier and the assertion, form-encoded.
            found = example.http.post(example.meta["token_endpoint"], timeout=10, data={
                "grant_type": "authorization_code",
                "code": query.get("code", ""),
                "redirect_uri": example.redirect_uri,
                "code_verifier": verifier,
                "client_id": example.client_id,
                "client_assertion_type": JWT_BEARER,
                "client_assertion": assertion,
            })
            if found.status_code != 200:
                return self.fail(example, f"The token endpoint answered {found.text}")

            tokens = found.json()
            # [END token]

            # The ID token, checked before anything in it is believed (section 6).
            claims = verify_id_token(tokens["id_token"], example.keys, example.issuer,
                                     example.client_id, nonce)

            # [START userinfo]
            # The access token goes in the Authorization header, never in the URL.
            bearer = {"Authorization": f"Bearer {tokens['access_token']}"}
            found = example.http.get(example.meta["userinfo_endpoint"], timeout=10,
                                     headers=bearer)
            found.raise_for_status()
            userinfo = found.json()

            # The same person as the ID token's.
            if userinfo["sub"] != claims["sub"]:
                return self.fail(example, "Userinfo names another person.")
            # [END userinfo]
        except (requests.RequestException, jwt.PyJWTError, KeyError, ValueError) as e:
            return self.fail(example, f"The sign-in failed: {e!r}")

        # Here your service starts its own session, keyed on sub, never on the email
        # (section 1).
        you = html.escape(claims["sub"])
        body = (f"<p>Aurora Account says you are <code>{you}</code>.</p>"
                f"<h2>The ID token's claims</h2>{shown(claims)}"
                f"<h2>Userinfo's answer</h2>{shown(userinfo)}"
                '<p><a href="/login">Sign in again</a></p>')
        self.answer(200, page("Signed in", body), cookie=example.cookie("", 0))


def main():
    example = Example()
    port = int(os.environ.get("PORT", "8443"))

    # Loopback only: in production, a reverse proxy in front of it serves it to others.
    server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
    server.example = example

    tls = bool(os.environ.get("TLS_CERT") or os.environ.get("TLS_KEY"))
    if tls:
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        context.load_cert_chain(setting("TLS_CERT"), setting("TLS_KEY"))
        server.socket = context.wrap_socket(server.socket, server_side=True,
                                            do_handshake_on_connect=False)

    scheme = "https" if tls else "http"
    print(f"Open {scheme}://localhost:{port}/ to sign in with {example.issuer}")
    server.serve_forever()


if __name__ == "__main__":
    main()