Security

// npm install jose
import { createRemoteJWKSet, jwtVerify } from "jose";

// Honor the JWKS cache policy and the TLP-TOK-9 refetch limit.
// The planned 30-day key prepublication policy is not a sandbox service guarantee.
const JWKS = createRemoteJWKSet(new URL("https://app.reflections.ai/.well-known/jwks.json"), {
  cacheMaxAge: 3_600_000,
  cooldownDuration: 3_600_000,
});
// 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",
]);

type Tlp = {
  ver: string;
  decision: string;
  max_generations: number;
  execution: { provider: string; model: string };
  obligations: string[];
};

export async function verifyAuthorization(token: string, model: string) {
  const { payload } = await jwtVerify<{ tlp: Tlp }>(token, JWKS, {
    algorithms: ["ES256"],
    typ: "tlp-auth+jwt",
    issuer: "https://app.reflections.ai",
    audience: "acct_7p4r6s9t0wxyhk5zba2q8v1m3n", // your platform
    clockTolerance: 60,
    maxTokenAge: "15m", // also rejects an iat in the future
    requiredClaims: ["exp", "iat", "jti"],
  });
  const { tlp } = payload;
  if (tlp.ver.split(".")[0] !== "1") throw new Error("unsupported protocol version");
  if (!tlp.obligations.every((type) => OBLIGATIONS.has(type))) throw new Error("unknown obligation");
  if (tlp.execution.model !== model) throw new Error("token is for a different model");
  if (!(await reserve(payload.jti!, 1, tlp.max_generations, payload.exp!))) {
    throw new Error("token used up");
  }
  return payload;
}

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, model: str) -> dict:
    if jwt.get_unverified_header(token).get("typ") != "tlp-auth+jwt":
        raise PermissionError("not an authorization token")
    claims = jwt.decode(
        token,
        signing_key(token),
        algorithms=["ES256"],
        issuer="https://app.reflections.ai",
        audience="acct_7p4r6s9t0wxyhk5zba2q8v1m3n",  # your platform
        leeway=60,
        options={"require": ["exp", "iat", "jti"]},  # also rejects an iat in the future
    )
    tlp = claims["tlp"]
    if str(tlp["ver"]).split(".")[0] != "1":
        raise PermissionError("unsupported protocol version")
    if not set(tlp["obligations"]) <= OBLIGATIONS:
        raise PermissionError("unknown obligation")
    if tlp["execution"]["model"] != model:
        raise PermissionError("token is for a different model")
    if not reserve(claims["jti"], 1, tlp["max_generations"], claims["exp"]):
        raise PermissionError("token used up")
    return claims