Authorization token

import { createPublicKey, verify } from "node:crypto";

type Jwk = { kid: string; kty: string; crv: string; x: string; y: string };
let jwks = { keys: [] as Jwk[], expiresAt: 0, maxAge: 300 };
let refetchAfter = 0;

// The obligation types of TLP 1.0. Reject a token that carries one you do not handle.
const OBLIGATIONS = new Set([
  "generation_check", "hold_until_release", "draft_watermark", "generation_instructions",
  "disclosure", "attribution", "c2pa_manifest", "no_training", "delete_material", "usage_report",
]);

async function load() {
  const res = await fetch("https://app.reflections.ai/.well-known/jwks.json");
  if (!res.ok) throw new Error(`JWKS fetch failed: ${res.status}`);
  const maxAge = Number(/max-age=(\d+)/.exec(res.headers.get("Cache-Control") ?? "")?.[1] ?? 300);
  jwks = { keys: (await res.json()).keys, expiresAt: Date.now() + maxAge * 1000, maxAge };
}

async function signingKey(kid: string): Promise<Jwk | undefined> {
  if (Date.now() >= jwks.expiresAt) await load();
  let key = jwks.keys.find((k) => k.kid === kid);
  if (!key && Date.now() >= refetchAfter) {
    refetchAfter = Date.now() + jwks.maxAge * 1000; // at most one refetch per max-age (TLP-TOK-9)
    await load();
    key = jwks.keys.find((k) => k.kid === kid);
  }
  return key;
}

export async function verifyAuthorization(token: string, platform: string, model: string, live: boolean) {
  const [h, p, s] = token.split(".");
  if (!h || !p || !s) throw new Error("malformed token");
  const header = JSON.parse(Buffer.from(h, "base64url").toString("utf8"));
  if (header.alg !== "ES256" || header.typ !== "tlp-auth+jwt") throw new Error("not an authorization token");
  if (live && String(header.kid).startsWith("rfl-test-")) throw new Error("test token in live mode");

  const jwk = await signingKey(header.kid);
  if (!jwk) throw new Error("unknown kid");
  const key = createPublicKey({ key: jwk, format: "jwk" });
  const signed = Buffer.from(`${h}.${p}`);
  if (!verify("sha256", signed, { key, dsaEncoding: "ieee-p1363" }, Buffer.from(s, "base64url"))) {
    throw new Error("bad signature");
  }

  const claims = JSON.parse(Buffer.from(p, "base64url").toString("utf8"));
  if (claims.iss !== "https://app.reflections.ai") throw new Error("wrong issuer");
  if (claims.aud !== platform) throw new Error("token is for another platform");
  const now = Date.now() / 1000;
  if (now > claims.exp + 60) throw new Error("token expired");
  if (now < claims.iat - 60) throw new Error("token issued in the future");
  if (String(claims.tlp?.ver).split(".")[0] !== "1") throw new Error("unsupported protocol version");
  if (claims.tlp.decision !== "allow") throw new Error("not an allow");
  if (claims.tlp.execution?.model !== model) throw new Error("token is for another model");
  if (!claims.tlp.obligations.every((t: string) => OBLIGATIONS.has(t))) throw new Error("unknown obligation");
  return claims; // Then count generations for claims.jti against claims.tlp.max_generations.
}

import time
import jwt  # pip install "pyjwt[crypto]"

JWKS = jwt.PyJWKClient("https://app.reflections.ai/.well-known/jwks.json", lifespan=3600)
refetch_after = 0.0
# The obligation types of TLP 1.0. Reject a token that carries one you do not handle.
OBLIGATIONS = {
    "generation_check", "hold_until_release", "draft_watermark", "generation_instructions",
    "disclosure", "attribution", "c2pa_manifest", "no_training", "delete_material", "usage_report",
}

def signing_key(token: str):
    """The token's key. Refetches the JWKS for an unknown kid at most once an hour (TLP-TOK-9)."""
    global refetch_after
    kid = jwt.get_unverified_header(token).get("kid")
    keys = {k.key_id: k.key for k in JWKS.get_signing_keys()}
    if kid not in keys and time.time() >= refetch_after:
        refetch_after = time.time() + 3600
        keys = {k.key_id: k.key for k in JWKS.get_signing_keys(refresh=True)}
    if kid not in keys:
        raise PermissionError("unknown_kid")
    return keys[kid]

def verify_authorization(token: str, platform: str, model: str, live: bool) -> dict:
    header = jwt.get_unverified_header(token)
    if header.get("alg") != "ES256" or header.get("typ") != "tlp-auth+jwt":
        raise ValueError("not an authorization token")
    if live and str(header.get("kid", "")).startswith("rfl-test-"):
        raise ValueError("test token in live mode")

    claims = jwt.decode(
        token,
        signing_key(token),
        algorithms=["ES256"],
        audience=platform,
        issuer="https://app.reflections.ai",
        leeway=60,
        options={"require": ["iss", "aud", "exp", "iat", "jti"]},  # also rejects an iat in the future
    )
    tlp = claims["tlp"]
    if str(tlp["ver"]).split(".")[0] != "1":
        raise ValueError("unsupported protocol version")
    if tlp["decision"] != "allow":
        raise ValueError("not an allow")
    if tlp["execution"]["model"] != model:
        raise ValueError("token is for another model")
    if not set(tlp["obligations"]) <= OBLIGATIONS:
        raise ValueError("unknown obligation")
    return claims  # Then count generations for claims["jti"] against claims["tlp"]["max_generations"].