OAuth2 Audience Validation — Ensuring Tokens Reach the Right Resource Server

import jwt, datetime, os
from flask import Flask, request, jsonify, g
from functools import wraps

app = Flask(__name__)
JWKS_URI = os.environ.get("JWKS_URI", "https://auth.dodatech.com/.well-known/jwks.json")

# Each resource server knows its own audience
RESOURCE_SERVER_AUDIENCE = "https://api.durga-antivirus.com/v1"

def validate_audience(token):
    """Validate the token's audience claim."""
    try:
        # In production, fetch JWKS and verify properly
        payload = jwt.decode(
            token,
            os.environ.get("AUTH_PUBLIC_KEY", ""),
            algorithms=["RS256"],
            audience=RESOURCE_SERVER_AUDIENCE,
            options={"verify_aud": True}
        )
        return payload
    except jwt.InvalidAudienceError:
        return None  # Audience mismatch
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

@app.before_request
def authenticate():
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        return jsonify({"error": "Unauthorized"}), 401

    payload = validate_audience(auth[7:])
    if not payload:
        return jsonify({
            "error": "invalid_token",
            "error_description": "Token audience does not match this resource server"
        }), 401

    g.current_user = payload

@app.route("/api/v1/threats")
def list_threats():
    return jsonify({"threats": [], "audience_validated": True})

# Token with multiple audiences
# aud: ["https://api.durga-antivirus.com/v1", "https://api.durga-antivirus.com/v2"]

MULTI_AUDIENCE_SERVERS = [
    "https://api.durga-antivirus.com/v1",
    "https://api.durga-antivirus.com/v2",
    "https://api.durga-antivirus.com/admin"
]

def validate_multi_audience(token):
    """Validate token against multiple allowed audiences."""
    try:
        # Decode without audience verification first
        payload = jwt.decode(
            token,
            os.environ.get("AUTH_PUBLIC_KEY", ""),
            algorithms=["RS256"],
            options={"verify_aud": False}
        )

        # Manual audience check
        token_aud = payload.get("aud", [])
        if isinstance(token_aud, str):
            token_aud = [token_aud]

        # Check if any of our audiences match
        server_audiences = [
            request.host_url.rstrip("/"),
            f"{request.scheme}://{request.host}{request.path}"
        ]

        for aud in token_aud:
            if any(server_aud.startswith(aud) for server_aud in server_audiences):
                return payload

        return None

    except jwt.InvalidTokenError:
        return None

# Custom decorator for specific audience
def require_audience(expected_audience):
    """Decorator to enforce specific audience."""
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            auth = request.headers.get("Authorization", "")
            payload = validate_multi_audience(auth[7:])
            if not payload:
                return jsonify({"error": "invalid_audience"}), 401

            token_aud = payload.get("aud", [])
            if isinstance(token_aud, str):
                token_aud = [token_aud]

            if expected_audience not in token_aud:
                return jsonify({
                    "error": "insufficient_audience",
                    "required": expected_audience,
                    "token_audiences": token_aud
                }), 403

            g.current_user = payload
            return f(*args, **kwargs)
        return decorated
    return decorator

@app.route("/api/v1/admin/config")
@require_audience("https://api.durga-antivirus.com/admin")
def admin_config():
    return jsonify({"config": "sensitive"})

def issue_service_token(service_name, audience, scopes):
    """Issue a token with specific audience."""
    now = datetime.datetime.utcnow()
    token = jwt.encode({
        "iss": "auth.dodatech.com",
        "sub": service_name,
        "aud": audience,
        "scope": " ".join(scopes),
        "iat": now,
        "exp": now + datetime.timedelta(hours=1),
        "jti": secrets.token_hex(16)
    }, SECRET, algorithm="HS256")
    return token

# Auth server issues tokens with specific audience
@app.route("/oauth/token", methods=["POST"])
def token():
    client_id = request.form.get("client_id")
    audience = request.form.get("audience", "default")

    # Validate client
    client = clients.get(client_id)
    if not client:
        return jsonify({"error": "invalid_client"}), 401

    # Validate audience is allowed for this client
    if audience not in client["allowed_audiences"]:
        return jsonify({"error": "invalid_target"}), 400

    token = issue_service_token(client_id, audience, client["scopes"])
    return jsonify({
        "access_token": token,
        "token_type": "Bearer",
        "expires_in": 3600,
        "audience": audience
    })