auth-middleware.py

from functools import wraps
from flask import request, jsonify
import jwt

def require_valid_token(public_key: str):
    """Decorator for endpoint authentication"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            # Extract token from header
            auth_header = request.headers.get('Authorization')
            if not auth_header or not auth_header.startswith('Bearer '):
                return jsonify({'error': 'Missing token'}), 401

            token = auth_header.split(' ')[1]

            try:
                # Validate token
                payload = jwt.decode(
                    token,
                    public_key,
                    algorithms=["ES256"]
                )

                # Check container binding if enforced
                if payload.get('enforce_container_binding') == 'true':
                    container_id = request.headers.get('X-Container-ID')
                    if container_id != payload.get('container_id'):
                        return jsonify({'error': 'Container ID mismatch'}), 403

                # Add payload to request context
                request.jwt_payload = payload

            except jwt.ExpiredSignatureError:
                return jsonify({'error': 'Token expired'}), 401
            except jwt.InvalidTokenError:
                return jsonify({'error': 'Invalid token'}), 401

            return f(*args, **kwargs)
        return decorated_function
    return decorator