Sign in with Aurora Account

"""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"


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")


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


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

        # 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})

        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)

        # 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)

    def callback(self, example: Example, query: dict):
        # 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']}.")

        verifier, nonce, _ = pending
        try:
            # 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()

            # 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)

            # 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.")
        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()

// Sign in with Aurora Account: a complete example in Node.js, 20 or later.
//
// 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 { createHash, randomBytes, randomUUID, timingSafeEqual } from "node:crypto";
import { readFileSync } from "node:fs";
import { createServer as createHttpServer } from "node:http";
import { createServer as createHttpsServer } from "node:https";
import { pathToFileURL } from "node:url";
import { SignJWT, createRemoteJWKSet, importPKCS8, jwtVerify } from "jose";

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

// The client assertion that proves a token request is yours (section 5): signed with
// your private key, valid for a minute, and accepted once.
export async function makeAssertion(clientId, issuer, privateKey) {
  const now = Math.floor(Date.now() / 1000);

  const claims = {
    iss: clientId, // you
    sub: clientId, // you, again
    aud: issuer, // the issuer itself, as one string
    iat: now,
    exp: now + 60, // a minute is plenty
    jti: randomUUID(), // new every time: an assertion works once
  };

  return new SignJWT(claims)
    .setProtectedHeader({ alg: "RS256", typ: "JWT" })
    .sign(privateKey);
}

// The ID token's claims, once it is shown to be Aurora Account's, for you, and for this
// sign-in (section 6). Throws otherwise.
export async function verifyIdToken(idToken, keys, issuer, clientId, nonce) {
  // The signature by the key set, RS256 only; then iss, aud and exp. The tolerance
  // forgives a clock a little off Aurora Account's, as Aurora Account forgives yours.
  const { payload: claims } = await jwtVerify(idToken, keys, {
    algorithms: ["RS256"],
    issuer,
    audience: clientId,
    requiredClaims: ["exp", "sub", "nonce"],
    clockTolerance: 30,
  });

  // aud must be your client_id as a string, not a list that holds it.
  if (typeof claims.aud !== "string")
    throw new Error("The ID token's aud is not a string.");

  // And the nonce this sign-in sent.
  if (!same(claims.nonce, nonce))
    throw new Error("The ID token's nonce is not the one sent.");

  return claims;
}

function setting(name) {
  if (!process.env[name]) throw new Error(`Set ${name}: README.md says how.`);
  return process.env[name];
}

function same(a, b) {
  const [x, y] = [Buffer.from(String(a)), Buffer.from(String(b))];
  return x.length === y.length && timingSafeEqual(x, y);
}

// One cookie from the Cookie header, read by hand: on localhost every dev server's
// cookies arrive here, and none of them may get in the way.
function cookieOf(req, name) {
  for (const part of (req.headers.cookie ?? "").split(";")) {
    const at = part.indexOf("=");
    if (at !== -1 && part.slice(0, at).trim() === name)
      return part.slice(at + 1).trim();
  }
  return "";
}

const escape = (text) =>
  String(text).replace(/[&<>"']/g, (c) => `&#${c.charCodeAt(0)};`);
const shown = (value) => `<pre>${escape(JSON.stringify(value, null, 2))}</pre>`;
const page = (title, body) =>
  `<!doctype html><html lang="en"><meta charset="utf-8">` +
  `<meta name="viewport" content="width=device-width, initial-scale=1">` +
  `<title>${title}</title><h1>${title}</h1>${body}`;

async function main() {
  const { default: express } = await import("express"); // only the server needs it

  const issuer = process.env.AURORA_ISSUER || "https://account.aurora-ws.com/o";
  const clientId = setting("CLIENT_ID");
  const redirectUri = setting("REDIRECT_URI");

  // A key that cannot sign fails now, not at the first sign-in.
  const pem = readFileSync(setting("PRIVATE_KEY_PATH"), "utf8");
  const privateKey = await importPKCS8(pem, "RS256").catch((error) => {
    throw new Error(
      `PRIVATE_KEY_PATH holds no RSA private key that signs: ${error.message}`,
    );
  });

  const headers = { "User-Agent": USER_AGENT };

  // The endpoints come from the discovery document, which names the issuer.
  const found = await fetch(`${issuer}/.well-known/openid-configuration`, {
    headers,
    signal: AbortSignal.timeout(10_000),
  });
  if (!found.ok) throw new Error(`The discovery document answered ${found.status}.`);
  const meta = await found.json();

  // It must name the issuer you configured, exactly.
  if (meta.issuer !== issuer)
    throw new Error(`The discovery document names ${meta.issuer}.`);

  // The key set: cached, and fetched again for a token whose key it does not hold
  // (section 6).
  const keys = createRemoteJWKSet(new URL(meta.jwks_uri), {
    headers,
    timeoutDuration: 10_000,
  });

  const pending = new Map(); // state -> { verifier, nonce, started }
  const cookie = {
    httpOnly: true,
    sameSite: "lax",
    secure: redirectUri.startsWith("https://"),
  };

  const app = express();
  app.disable("x-powered-by");
  app.use((req, res, next) => {
    res.set({
      "Cache-Control": "no-store",
      "Content-Security-Policy": "default-src 'none'",
    });
    next();
  });

  const fail = (res, message) => {
    const body = `<p>${escape(message)}</p><p><a href="/login">Sign in again</a></p>`;
    res
      .status(400)
      .clearCookie("aurora_state", cookie)
      .send(page("Not signed in", body));
  };

  app.get("/", (req, res) => {
    res.send(
      page("Example", `<p><a href="/login">Sign in with Aurora Account</a></p>`),
    );
  });

  app.get("/login", (req, res) => {
    const now = Date.now();
    for (const [state, { started }] of pending)
      if (now - started > PENDING_FOR) pending.delete(state);

    // Three fresh random values for this sign-in, kept on the server.
    const [state, nonce, verifier] = [1, 2, 3].map(() =>
      randomBytes(32).toString("base64url"),
    );
    pending.set(state, { verifier, nonce, started: now });

    // PKCE: the challenge is the verifier's SHA-256, base64url without padding.
    const codeChallenge = createHash("sha256").update(verifier).digest("base64url");

    const url = new URL(meta.authorization_endpoint);
    url.search = new URLSearchParams({
      response_type: "code",
      client_id: clientId,
      redirect_uri: redirectUri,
      scope: SCOPE,
      state,
      nonce,
      code_challenge: codeChallenge,
      code_challenge_method: "S256",
    }).toString();

    // The state is tied to this browser by a cookie, checked on the way back.
    res
      .cookie("aurora_state", state, { ...cookie, maxAge: PENDING_FOR })
      .redirect(url.href);
  });

  app.get("/callback", async (req, res) => {
    // The state must be the one this browser was given, and still pending.
    const query = new URL(req.originalUrl, "http://localhost").searchParams;
    const state = query.get("state") ?? "";
    const started = same(cookieOf(req, "aurora_state"), state)
      ? pending.get(state)
      : undefined;
    if (started) pending.delete(state); // used once
    if (!started || Date.now() - started.started > PENDING_FOR) {
      return fail(res, "This sign-in did not start here, or expired.");
    }

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

    // access_denied: the person pressed Cancel (section 9).
    if (query.has("error"))
      return fail(res, `Aurora Account answered ${query.get("error")}.`);

    try {
      // The code, the PKCE verifier and a new assertion, form-encoded. Each assertion
      // is accepted once, so every request makes its own.
      const token = await fetch(meta.token_endpoint, {
        method: "POST",
        headers,
        signal: AbortSignal.timeout(10_000),
        body: new URLSearchParams({
          grant_type: "authorization_code",
          code: query.get("code") ?? "",
          redirect_uri: redirectUri,
          code_verifier: started.verifier,
          client_id: clientId,
          client_assertion_type: JWT_BEARER,
          client_assertion: await makeAssertion(clientId, issuer, privateKey),
        }),
      });
      if (!token.ok)
        return fail(res, `The token endpoint answered ${await token.text()}`);

      const tokens = await token.json();

      // The ID token, checked before anything in it is believed (section 6).
      const claims = await verifyIdToken(
        tokens.id_token,
        keys,
        issuer,
        clientId,
        started.nonce,
      );

      // The access token goes in the Authorization header, never in the URL.
      const info = await fetch(meta.userinfo_endpoint, {
        headers: { ...headers, Authorization: `Bearer ${tokens.access_token}` },
        signal: AbortSignal.timeout(10_000),
      });
      if (!info.ok) return fail(res, `Userinfo answered ${await info.text()}`);
      const userinfo = await info.json();

      // The same person as the ID token's.
      if (userinfo.sub !== claims.sub)
        return fail(res, "Userinfo names another person.");

      // Here your service starts its own session, keyed on sub, never on the email
      // (section 1).
      const body =
        `<p>Aurora Account says you are <code>${escape(claims.sub)}</code>.</p>` +
        `<h2>The ID token's claims</h2>${shown(claims)}` +
        `<h2>Userinfo's answer</h2>${shown(userinfo)}` +
        `<p><a href="/login">Sign in again</a></p>`;
      res.clearCookie("aurora_state", cookie).send(page("Signed in", body));
    } catch (problem) {
      fail(res, `The sign-in failed: ${problem}`);
    }
  });

  const port = Number(process.env.PORT || 8443);
  const tls = Boolean(process.env.TLS_CERT || process.env.TLS_KEY);
  const certificate = () => ({
    cert: readFileSync(setting("TLS_CERT")),
    key: readFileSync(setting("TLS_KEY")),
  });

  const server = tls ? createHttpsServer(certificate(), app) : createHttpServer(app);
  server.on("error", (error) => {
    console.error(`Cannot serve on port ${port}: ${error.message}`);
    process.exit(1);
  });

  // Loopback only: in production, a reverse proxy in front of it serves it to others.
  server.listen(port, "127.0.0.1", () => {
    console.log(
      `Open ${tls ? "https" : "http"}://localhost:${port}/ to sign in with ${issuer}`,
    );
  });
}

// Started as a program (`node app.mjs`), it serves; imported, it lends makeAssertion.
if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) {
  main().catch((problem) => {
    console.error(problem.message);
    process.exit(1);
  });
}