Security & Authentication at Scale

# JWT Authentication in Python
import jwt
import time
from datetime import datetime, timedelta
from functools import wraps
from flask import Flask, request, jsonify

app = Flask(__name__)
SECRET_KEY = "your-secret-key"  # Use env var in production!

def create_token(user_id: str, role: str,
                  expires_minutes=30) -> str:
    """Create a signed JWT token."""
    payload = {
        "sub": user_id,           # Subject (who)
        "role": role,             # Custom claim
        "iat": datetime.utcnow(), # Issued at
        "exp": datetime.utcnow()  # Expiration
              + timedelta(minutes=expires_minutes),
        "iss": "myapp.com",       # Issuer
    }
    return jwt.encode(payload, SECRET_KEY, algorithm="HS256")

def verify_token(token: str) -> dict:
    """Verify and decode a JWT token."""
    try:
        payload = jwt.decode(
            token, SECRET_KEY,
            algorithms=["HS256"],
            issuer="myapp.com"
        )
        return payload
    except jwt.ExpiredSignatureError:
        raise AuthError("Token expired")
    except jwt.InvalidTokenError:
        raise AuthError("Invalid token")

def require_auth(allowed_roles=None):
    """Decorator for protecting endpoints."""
    def decorator(f):
        @wraps(f)
        def wrapped(*args, **kwargs):
            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]
            payload = verify_token(token)

            # Check authorization (role-based)
            if allowed_roles and payload["role"] not in allowed_roles:
                return jsonify({"error": "Forbidden"}), 403

            request.user = payload
            return f(*args, **kwargs)
        return wrapped
    return decorator

# Usage
@app.route("/api/admin/users")
@require_auth(allowed_roles=["admin"])
def list_users():
    return jsonify({"users": ["..."]})

# OAuth 2.0 Authorization Code Flow (server-side)
import requests
import secrets
from flask import Flask, redirect, request, session

app = Flask(__name__)
app.secret_key = secrets.token_hex(32)

OAUTH_CONFIG = {
    "client_id": "your-client-id",
    "client_secret": "your-client-secret",
    "auth_url": "https://auth.example.com/authorize",
    "token_url": "https://auth.example.com/token",
    "redirect_uri": "https://myapp.com/callback",
}

@app.route("/login")
def login():
    # Step 1: Generate state for CSRF protection
    state = secrets.token_urlsafe(32)
    session["oauth_state"] = state

    # Step 2: Redirect to authorization server
    auth_url = (
        f"{OAUTH_CONFIG['auth_url']}?"
        f"response_type=code&"
        f"client_id={OAUTH_CONFIG['client_id']}&"
        f"redirect_uri={OAUTH_CONFIG['redirect_uri']}&"
        f"scope=openid profile email&"
        f"state={state}"
    )
    return redirect(auth_url)

@app.route("/callback")
def callback():
    # Step 3: Verify state to prevent CSRF
    if request.args.get("state") != session.get("oauth_state"):
        return "State mismatch - possible CSRF", 403

    # Step 4: Exchange auth code for tokens
    code = request.args.get("code")
    token_response = requests.post(
        OAUTH_CONFIG["token_url"],
        data={
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": OAUTH_CONFIG["redirect_uri"],
            "client_id": OAUTH_CONFIG["client_id"],
            "client_secret": OAUTH_CONFIG["client_secret"],
        }
    )

    tokens = token_response.json()
    session["access_token"] = tokens["access_token"]
    session["id_token"] = tokens["id_token"]

    return redirect("/dashboard")