JSON Web Key Set (JWKS) Management¶

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives import serialization
import json

class JWKSManager:
    def __init__(self, config):
        self.config = config
        self.keys = {}
        self.current_key_id = None

    async def initialize(self):
        """Initialize JWKS with first key"""
        if not self.config.JWKS_ENABLED:
            return

        # Generate initial key pair
        await self.rotate_keys()

    async def rotate_keys(self):
        """Generate new key pair and add to JWKS"""
        # Generate Ed25519 key pair (default algorithm)
        private_key = Ed25519PrivateKey.generate()

        # Generate key ID with timestamp
        key_id = f"arcp-{datetime.now().strftime('%Y%m%d-%H%M%S')}"

        # Store key pair
        self.keys[key_id] = {
            "private_key": private_key,
            "public_key": private_key.public_key(),
            "created_at": datetime.now(),
            "algorithm": self.config.JWKS_ALGORITHM
        }

        # Set as current signing key
        self.current_key_id = key_id

        # Clean up old keys
        await self._cleanup_old_keys()

        logger.info(f"Rotated to new key: {key_id}")

    def get_jwks(self):
        """Return public JWKS for clients"""
        keys = []

        for kid, key_data in self.keys.items():
            public_key = key_data["public_key"]
            algorithm = key_data["algorithm"]

            if algorithm == "EdDSA":
                # Export Ed25519 public key
                public_bytes = public_key.public_bytes(
                    encoding=serialization.Encoding.Raw,
                    format=serialization.PublicFormat.Raw
                )

                keys.append({
                    "kty": "OKP",
                    "crv": "Ed25519",
                    "use": "sig",
                    "kid": kid,
                    "alg": "EdDSA",
                    "x": self._base64url_encode(public_bytes)
                })
            elif algorithm == "ES256":
                # Export EC P-256 public key
                public_numbers = public_key.public_numbers()

                keys.append({
                    "kty": "EC",
                    "crv": "P-256",
                    "use": "sig",
                    "kid": kid,
                    "alg": "ES256",
                    "x": self._base64url_encode(public_numbers.x),
                    "y": self._base64url_encode(public_numbers.y)
                })

        return {"keys": keys}

    def sign_token(self, payload):
        """Sign JWT with current private key"""
        key_data = self.keys[self.current_key_id]

        # Add kid to header
        headers = {"kid": self.current_key_id}

        token = jwt.encode(
            payload,
            key_data["private_key"],
            algorithm=key_data["algorithm"],
            headers=headers
        )

        return token

from fastapi import APIRouter

router = APIRouter()

@router.get("/.well-known/jwks.json")
async def get_jwks():
    """Public JWKS endpoint for token verification"""
    jwks = jwks_manager.get_jwks()

    return JSONResponse(
        content=jwks,
        headers={
            "Cache-Control": "public, max-age=3600",
            "Content-Type": "application/json"
        }
    )

import jwt
import requests
from jwt import PyJWKClient

class ARCPClient:
    def __init__(self, base_url):
        self.base_url = base_url
        self.jwks_client = PyJWKClient(
            f"{base_url}/.well-known/jwks.json",
            cache_keys=True,
            max_cached_keys=10
        )

    def verify_token(self, token):
        """Verify JWT token using JWKS"""
        try:
            # Decode header to get kid
            header = jwt.get_unverified_header(token)
            kid = header.get("kid")

            if not kid:
                raise ValueError("No kid in token header")

            # Fetch signing key from JWKS
            signing_key = self.jwks_client.get_signing_key_from_jwt(token)

            # Verify token
            payload = jwt.decode(
                token,
                signing_key.key,
                algorithms=["EdDSA", "ES256"],  # Support both EdDSA and ES256
                audience="arcp-agent",
                options={"verify_exp": True}
            )

            return payload

        except jwt.ExpiredSignatureError:
            raise TokenExpired()
        except jwt.InvalidTokenError as e:
            raise TokenInvalid(str(e))

async def _rotation_task(self):
    """Background task for automatic key rotation"""
    while True:
        # Wait for rotation interval (convert days to seconds)
        await asyncio.sleep(self.config.JWKS_ROTATION_DAYS * 86400)

        try:
            await self.rotate_keys()
            logger.info("Automatic key rotation successful")
        except Exception as e:
            logger.error(f"Key rotation failed: {e}")

async def _cleanup_old_keys(self):
    """Remove keys older than overlap period"""
    now = datetime.now()
    overlap = timedelta(days=self.config.JWKS_OVERLAP_DAYS)

    keys_to_remove = []

    for kid, key_data in self.keys.items():
        # Keep current key
        if kid == self.current_key_id:
            continue

        # Remove if older than overlap period
        age = now - key_data["created_at"]
        if age > overlap:
            keys_to_remove.append(kid)

    for kid in keys_to_remove:
        del self.keys[kid]
        logger.info(f"Removed expired key: {kid}")