auth-service.py

#!/usr/bin/env python3
import jwt
import time
import secrets
from typing import List, Optional
from datetime import datetime, timedelta

class AuthService:
    def __init__(self, private_key_path: str, key_id: str):
        with open(private_key_path, 'r') as f:
            self.private_key = f.read()
        self.key_id = key_id

    def generate_container_token(self,
                                 container_id: str,
                                 user_id: str,
                                 allowed_hosts: List[str],
                                 duration_hours: int = 4,
                                 enforce_binding: bool = True) -> str:
        """
        Generate JWT for container access

        CRITICAL: Set enforce_container_binding=true for production
        """
        now = int(time.time())

        payload = {
            # Standard claims
            "iss": "your-service-egress-control",
            "iat": now,
            "exp": now + (duration_hours * 3600),
            "jti": secrets.token_urlsafe(32),  # Unique token ID

            # Custom claims
            "container_id": container_id,
            "organization_uuid": user_id,
            "allowed_hosts": ",".join(allowed_hosts),

            # Security flags
            "enforce_container_binding": "true" if enforce_binding else "false",
            "use_egress_gateway": "true",
            "enforce_centralized_egress": "true",

            # Compliance flags (customize per user)
            "is_hipaa_regulated": "false",
            "is_pci_compliant": "false"
        }

        token = jwt.encode(
            payload,
            self.private_key,
            algorithm="ES256",
            headers={"kid": self.key_id}
        )

        return token

    def validate_token(self, token: str, public_key: str) -> Optional[dict]:
        """Validate JWT token"""
        try:
            payload = jwt.decode(
                token,
                public_key,
                algorithms=["ES256"],
                options={"verify_exp": True}
            )
            return payload
        except jwt.ExpiredSignatureError:
            return None
        except jwt.InvalidTokenError:
            return None

# Usage
if __name__ == '__main__':
    auth = AuthService('private-key.pem', 'your-key-id')

    token = auth.generate_container_token(
        container_id="sandbox_session_def456",
        user_id="user_789",
        allowed_hosts=["api.example.com", "github.com"],
        enforce_binding=True
    )

    print(f"Token: {token}")