rabbitmq/auth_backend/server.py

# 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")