Authentication & Security

Role

You are an elite security architect. You design and implement authentication and

authorization systems with defense-in-depth, following OWASP best practices, proper

token handling, and secure secrets management.

Part 1: JWT Authentication

JWT Token Service

from datetime import datetime, timedelta, timezone
from typing import Optional
import jwt
from pydantic import BaseModel

SECRET_KEY = "your-secret-key-from-env"  # Load from env/secrets manager
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE = timedelta(minutes=15)
REFRESH_TOKEN_EXPIRE = timedelta(days=7)


class TokenPayload(BaseModel):
    sub: str          # User ID
    exp: datetime     # Expiration
    iat: datetime     # Issued at
    scope: str = ""   # Permissions scope
    jti: str = ""     # Token ID (for revocation)


def create_access_token(user_id: str, scopes: list[str] = None) -> str:
    now = datetime.now(timezone.utc)
    payload = {
        "sub": user_id,
        "exp": now + ACCESS_TOKEN_EXPIRE,
        "iat": now,
        "scope": " ".join(scopes or []),
        "type": "access",
    }
    return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)


def create_refresh_token(user_id: str) -> str:
    import uuid
    now = datetime.now(timezone.utc)
    payload = {
        "sub": user_id,
        "exp": now + REFRESH_TOKEN_EXPIRE,
        "iat": now,
        "jti": str(uuid.uuid4()),
        "type": "refresh",
    }
    return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)


def verify_token(token: str) -> Optional[TokenPayload]:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return TokenPayload(**payload)
    except jwt.ExpiredSignatureError:
        return None  # Token expired
    except jwt.InvalidTokenError:
        return None  # Invalid token


# FastAPI dependency
from fastapi import Depends, HTTPException, Security
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

security = HTTPBearer()

async def get_current_user(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> TokenPayload:
    token_data = verify_token(credentials.credentials)
    if not token_data:
        raise HTTPException(401, "Invalid or expired token")
    return token_data

Part 2: OAuth2 / OIDC

OAuth2 Authorization Code Flow

from fastapi import APIRouter, Request
from fastapi.responses import RedirectResponse
import httpx
import secrets

router = APIRouter(prefix="/auth")

GOOGLE_CLIENT_ID = "your-client-id"
GOOGLE_CLIENT_SECRET = "your-client-secret"
REDIRECT_URI = "https://app.example.com/auth/callback"

@router.get("/login/google")
async def google_login(request: Request):
    state = secrets.token_urlsafe(32)
    request.session["oauth_state"] = state

    auth_url = (
        "https://accounts.google.com/o/oauth2/v2/auth?"
        f"client_id={GOOGLE_CLIENT_ID}"
        f"&redirect_uri={REDIRECT_URI}"
        f"&response_type=code"
        f"&scope=openid email profile"
        f"&state={state}"
        f"&access_type=offline"
    )
    return RedirectResponse(auth_url)

@router.get("/callback")
async def google_callback(request: Request, code: str, state: str):
    # Verify state to prevent CSRF
    if state != request.session.get("oauth_state"):
        raise HTTPException(400, "Invalid state parameter")

    # Exchange code for tokens
    async with httpx.AsyncClient() as client:
        token_resp = await client.post("https://oauth2.googleapis.com/token", data={
            "code": code,
            "client_id": GOOGLE_CLIENT_ID,
            "client_secret": GOOGLE_CLIENT_SECRET,
            "redirect_uri": REDIRECT_URI,
            "grant_type": "authorization_code",
        })
        tokens = token_resp.json()

        # Get user info
        userinfo_resp = await client.get(
            "https://www.googleapis.com/oauth2/v2/userinfo",
            headers={"Authorization": f"Bearer {tokens['access_token']}"},
        )
        user_info = userinfo_resp.json()

    # Create or update user in database
    user = await get_or_create_user(
        email=user_info["email"],
        name=user_info["name"],
        picture=user_info.get("picture"),
    )

    # Issue our own JWT
    access_token = create_access_token(str(user.id))
    return {"access_token": access_token, "token_type": "bearer"}

Part 3: RBAC (Role-Based Access Control)

from enum import Enum
from functools import wraps
from fastapi import HTTPException

class Role(str, Enum):
    ADMIN = "admin"
    EDITOR = "editor"
    VIEWER = "viewer"

# Permission matrix
PERMISSIONS = {
    Role.ADMIN: {"read", "write", "delete", "manage_users", "manage_settings"},
    Role.EDITOR: {"read", "write"},
    Role.VIEWER: {"read"},
}

def has_permission(role: Role, permission: str) -> bool:
    return permission in PERMISSIONS.get(role, set())

def require_permission(*permissions: str):
    """FastAPI dependency for permission checking."""
    async def check(user: TokenPayload = Depends(get_current_user)):
        user_role = Role(user.scope.split()[0]) if user.scope else Role.VIEWER
        for perm in permissions:
            if not has_permission(user_role, perm):
                raise HTTPException(
                    403,
                    f"Permission denied: requires '{perm}', user role is '{user_role.value}'"
                )
        return user
    return Depends(check)

# Usage
@router.delete("/users/{user_id}")
async def delete_user(user_id: int, user=require_permission("delete", "manage_users")):
    await db.execute("DELETE FROM users WHERE id = $1", user_id)
    return {"deleted": user_id}

@router.get("/reports")
async def get_reports(user=require_permission("read")):
    return await db.fetch_all("SELECT * FROM reports")

Part 4: API Key Management

import hashlib
import secrets
from datetime import datetime, timezone

def generate_api_key(prefix: str = "axe") -> tuple[str, str]:
    """Generate API key and its hash. Return (raw_key, hash)."""
    raw_key = f"{prefix}_{secrets.token_urlsafe(32)}"
    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
    return raw_key, key_hash

async def verify_api_key(api_key: str) -> dict | None:
    """Verify API key against stored hash."""
    key_hash = hashlib.sha256(api_key.encode()).hexdigest()
    record = await db.fetch_one(
        "SELECT * FROM api_keys WHERE key_hash = $1 AND revoked_at IS NULL",
        key_hash,
    )
    if record:
        # Update last used timestamp
        await db.execute(
            "UPDATE api_keys SET last_used_at = $1 WHERE id = $2",
            datetime.now(timezone.utc), record["id"],
        )
    return record

# FastAPI dependency
from fastapi.security import APIKeyHeader

api_key_header = APIKeyHeader(name="X-API-Key")

async def validate_api_key(api_key: str = Security(api_key_header)):
    record = await verify_api_key(api_key)
    if not record:
        raise HTTPException(401, "Invalid API key")
    # Check rate limits per key
    if record.get("rate_limit"):
        await check_rate_limit(f"apikey:{record['id']}", record["rate_limit"])
    return record

# Endpoints
@router.post("/api-keys")
async def create_key(name: str, user=require_permission("manage_settings")):
    raw_key, key_hash = generate_api_key()
    await db.execute(
        "INSERT INTO api_keys (name, key_hash, user_id, created_at) VALUES ($1, $2, $3, $4)",
        name, key_hash, user.sub, datetime.now(timezone.utc),
    )
    return {"api_key": raw_key, "note": "Save this key — it cannot be retrieved again"}

Part 5: Security Headers

from starlette.middleware.base import BaseHTTPMiddleware

class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)

        # Prevent MIME sniffing
        response.headers["X-Content-Type-Options"] = "nosniff"

        # Clickjacking protection
        response.headers["X-Frame-Options"] = "DENY"

        # XSS protection (legacy browsers)
        response.headers["X-XSS-Protection"] = "1; mode=block"

        # Strict Transport Security
        response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"

        # Content Security Policy
        response.headers["Content-Security-Policy"] = (
            "default-src 'self'; "
            "script-src 'self' 'unsafe-inline'; "
            "style-src 'self' 'unsafe-inline'; "
            "img-src 'self' data: https:; "
            "font-src 'self'; "
            "connect-src 'self' https://api.example.com; "
            "frame-ancestors 'none'; "
        )

        # Referrer policy
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"

        # Permissions policy
        response.headers["Permissions-Policy"] = (
            "camera=(), microphone=(), geolocation=(), payment=()"
        )

        return response

# CORS configuration
from fastapi.middleware.cors import CORSMiddleware

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://app.example.com",
        "https://admin.example.com",
    ],
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
    allow_headers=["Authorization", "Content-Type", "X-API-Key"],
    expose_headers=["X-Request-ID"],
    max_age=3600,
)

Part 6: TOTP / 2FA

import pyotp
import qrcode
import io
import base64

def setup_totp(user_email: str, issuer: str = "AXE") -> dict:
    """Generate TOTP secret and QR code for 2FA setup."""
    secret = pyotp.random_base32()
    totp = pyotp.TOTP(secret)

    # Generate provisioning URI
    uri = totp.provisioning_uri(name=user_email, issuer_name=issuer)

    # Generate QR code as base64
    qr = qrcode.make(uri)
    buf = io.BytesIO()
    qr.save(buf, format="PNG")
    qr_b64 = base64.b64encode(buf.getvalue()).decode()

    return {
        "secret": secret,  # Store encrypted in database
        "qr_code": f"data:image/png;base64,{qr_b64}",
        "manual_entry_key": secret,
    }

def verify_totp(secret: str, code: str) -> bool:
    """Verify a TOTP code with 30-second window tolerance."""
    totp = pyotp.TOTP(secret)
    return totp.verify(code, valid_window=1)  # Allows 1 period before/after

# Login with 2FA
@router.post("/login")
async def login(email: str, password: str, totp_code: str = None):
    user = await authenticate_user(email, password)
    if not user:
        raise HTTPException(401, "Invalid credentials")

    if user.totp_enabled:
        if not totp_code:
            return {"requires_2fa": True}
        if not verify_totp(user.totp_secret, totp_code):
            raise HTTPException(401, "Invalid 2FA code")

    token = create_access_token(str(user.id))
    return {"access_token": token}

Part 7: Input Validation & Sanitization

import re
import html
from pydantic import BaseModel, field_validator, Field

class SafeUserInput(BaseModel):
    """Input model with security validations."""
    name: str = Field(..., min_length=1, max_length=200)
    email: str = Field(..., pattern=r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$')
    bio: str = Field(default="", max_length=2000)
    website: str = Field(default="", max_length=500)

    @field_validator("name")
    @classmethod
    def sanitize_name(cls, v: str) -> str:
        # Remove control characters
        v = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', v)
        return html.escape(v.strip())

    @field_validator("bio")
    @classmethod
    def sanitize_bio(cls, v: str) -> str:
        return html.escape(v.strip())

    @field_validator("website")
    @classmethod
    def validate_website(cls, v: str) -> str:
        if v and not v.startswith(("https://", "http://")):
            raise ValueError("Website must start with https:// or http://")
        # Prevent SSRF: block internal IPs
        if v:
            from urllib.parse import urlparse
            hostname = urlparse(v).hostname
            blocked = ["localhost", "127.0.0.1", "0.0.0.0", "169.254.169.254"]
            if hostname in blocked or (hostname and hostname.startswith("10.")):
                raise ValueError("Internal URLs are not allowed")
        return v

# SQL injection prevention — always use parameterized queries
# NEVER: f"SELECT * FROM users WHERE id = {user_id}"
# ALWAYS: await db.fetch_one("SELECT * FROM users WHERE id = $1", user_id)

Part 8: Secrets Management

import json
from pathlib import Path
from functools import lru_cache

class SecretsManager:
    """Local secrets manager — loads from encrypted files."""
    def __init__(self, secrets_dir: str = "~/.axe/secrets"):
        self.secrets_dir = Path(secrets_dir).expanduser()

    @lru_cache(maxsize=32)
    def get(self, name: str) -> dict:
        """Load a secret by name."""
        secret_file = self.secrets_dir / f"{name}.json"
        if not secret_file.exists():
            raise FileNotFoundError(f"Secret '{name}' not found")

        # Verify file permissions
        mode = oct(secret_file.stat().st_mode)[-3:]
        if mode != "600":
            raise PermissionError(f"Secret file {name} has insecure permissions: {mode} (expected 600)")

        with open(secret_file) as f:
            return json.load(f)

    def set(self, name: str, data: dict):
        """Save a secret."""
        secret_file = self.secrets_dir / f"{name}.json"
        self.secrets_dir.mkdir(parents=True, exist_ok=True)
        with open(secret_file, "w") as f:
            json.dump(data, f, indent=2)
        secret_file.chmod(0o600)  # Owner read/write only
        self.get.cache_clear()

secrets = SecretsManager()

# Usage
db_creds = secrets.get("database")
api_key = secrets.get("openai")["api_key"]

# Environment-based fallback
import os

def get_secret(name: str, key: str = None) -> str:
    """Get secret from env var or secrets file."""
    env_key = f"{name.upper()}{'_' + key.upper() if key else ''}"
    env_val = os.environ.get(env_key)
    if env_val:
        return env_val
    try:
        data = secrets.get(name)
        return data[key] if key else data
    except (FileNotFoundError, KeyError):
        raise ValueError(f"Secret '{name}' not found in env or secrets file")

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)