# rabbitmq_auth_backend.py import json import logging import os import re import urllib.error import urllib.request from typing import Dict, List, Optional from urllib.parse import urlencode from fastapi import FastAPI, Form, HTTPException from pydantic import BaseModel from starlette.responses import PlainTextResponse # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI(title="RabbitMQ HTTP Auth Backend") # Configuration YOUR_API_BASE_URL = os.environ["BACKEND_URL"] def get_default_headers(token: str = None, url: str = None) -> Dict[str, str]: """ Get default headers for API requests. """ headers = { "Content-Type": "application/json", "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br", "Connection": "keep-alive", } if token: headers["Authorization"] = f"Bearer {token}" if url: headers["Referer"] = url # Extract host from URL from urllib.parse import urlparse parsed_url = urlparse(url) headers["Host"] = parsed_url.netloc return headers async def validate_token_and_get_user(token: str) -> dict | None: """ Validate the Keycloak token and get user info. This replicates what your get_current_token_payload does. """ try: url = f"{YOUR_API_BASE_URL}/user" headers = get_default_headers(token=token, url=url) # Call your API's /me endpoint to validate token req = urllib.request.Request( url, headers=headers, method="GET", ) with urllib.request.urlopen(req, timeout=10) as response: if response.status == 200: data = json.loads(response.read().decode()) return data else: logger.error(f"Token validation failed: {response.status}") return None except Exception as e: logger.error(f"Error validating token: {e}") return None async def check_system_access(system_id: int, user_info: dict) -> bool: """ Check if the user has access to the computational system. Uses your existing API endpoints. """ try: params = urlencode({"system_id": system_id}) url = f"{YOUR_API_BASE_URL}/machine/system?{params}" headers = get_default_headers(token=user_info.get("token"), url=url) req = urllib.request.Request( url, headers=headers, method="GET", ) with urllib.request.urlopen(req, timeout=10) as response: if response.status == 200: return True elif response.status == 403: logger.warning( f"User {user_info.get('keycloak_id')} denied access to system {system_id}" ) return False else: logger.error(f"Error checking system access: {response.status}") return False except urllib.error.HTTPError as e: if e.code == 403: logger.warning( f"User {user_info.get('keycloak_id')} denied access to system {system_id}" ) return False else: logger.error(f"Error checking system access: {e.code}") return False except Exception as e: logger.error(f"Error checking system access: {e}") return False async def get_system_teams(system_id: int) -> List[Dict]: """ Get all teams that this system is a member of. """ try: url = f"{YOUR_API_BASE_URL}/machine/system/team?system_id={system_id}" headers = get_default_headers(url=url) req = urllib.request.Request( url, headers=headers, method="GET", ) with urllib.request.urlopen(req, timeout=10) as response: if response.status == 200: data = json.loads(response.read().decode()) return data else: return [] except Exception as e: logger.error(f"Error getting system teams: {e}") return [] @app.post("/rabbit/auth/user") async def authenticate_user( username: str = Form(...), # Changed from AuthRequest to form parameters password: str = Form(...), ): """ RabbitMQ calls this to authenticate a user. Username format: "system_{system_id}" Password: Keycloak token """ logger.info(f"Auth request for username: {username}") try: system_id = int(username) except (IndexError, ValueError): logger.warning(f"Invalid system_id in username: {username}") return PlainTextResponse("deny") # Validate the token and get user info user_info = await validate_token_and_get_user(password) if not user_info: logger.warning(f"Invalid token for system {system_id}") return PlainTextResponse("deny") # Store token in user_info for subsequent checks user_info["token"] = password # Check if user has access to the system if not await check_system_access(system_id, user_info): return PlainTextResponse("deny") logger.info( f"Authentication successful for system {system_id}, user {user_info.get('keycloak_id')}" ) # Return success - RabbitMQ will allow connection return PlainTextResponse("allow") @app.post("/rabbit/auth/vhost") async def authorize_vhost( username: str = Form(...), vhost: str = Form(...), ip: Optional[str] = Form(None) ): """ RabbitMQ calls this to check vhost permissions. Only allow specific vhosts. """ logger.info(f"VHost check for user: {username}, vhost: {vhost}") # Only allow specific vhosts allowed_vhosts = ["/"] if vhost not in allowed_vhosts: logger.warning(f"Unauthorized vhost access attempt: {vhost}") return PlainTextResponse("deny") logger.info(f"VHost access granted for {username} to {vhost}") return PlainTextResponse("allow") @app.post("/rabbit/auth/resource") async def authorize_resource( username: str = Form(...), vhost: str = Form(...), resource: str = Form(...), name: str = Form(...), permission: str = Form(...), ip: Optional[str] = Form(None), ): """ RabbitMQ calls this to check resource permissions. """ logger.info( f"Resource check for {username} on vhost {vhost}, " f"resource: {resource}, name: {name}, " f"permission: {permission}" ) try: system_id = int(username) except (IndexError, ValueError): return PlainTextResponse("deny") # Permission logic based on vhost and resource type if vhost == "/": # On heartbeat vhost, allow publishing to "heartbeat" exchange if resource == "exchange": if (name == "heartbeat" or name == "progress_report") and ( permission == "write" or permission == "configure" ): logger.info(f"Publish permission granted for heartbeat exchange") return PlainTextResponse("allow") else: logger.warning(f"Unauthorized exchange access: {name}") # Handle queue operations on team vhost if resource == "queue": # Pattern: team_{team_id}.qubits_{N} team_queue_pattern = r"^team_(\w+)\.qubits_(\d+)$" match = re.match(team_queue_pattern, name) if not match: return PlainTextResponse("deny") team_id = match.group(1) qubits = int(match.group(2)) # Get teams this system belongs to system_teams = await get_system_teams(system_id) logger.info(system_teams) team = list( filter(lambda x: int(x.get("team_id")) == int(team_id), system_teams) ) # Check if system is a member of this team if len(team) == 0: logger.warning(f"System {system_id} not a member of team {team_id}") return PlainTextResponse("deny") team = team[0] if team.get("qubits_given") < qubits: logger.warning(f"System {system_id} does not have that many qubits") return PlainTextResponse("deny") # Allow configure and read permissions on team qubit queues if permission in ["configure", "read"]: logger.info(f"System {system_id} granted {permission} on queue {name}") return PlainTextResponse("allow") return PlainTextResponse("deny") # Default deny return PlainTextResponse("deny")