security/rbac.py

"""
Role-Based Access Control (RBAC) for AegisLM

Defines roles, permissions, and access control logic for the multi-tenant system.
"""

import uuid
from enum import Enum
from typing import Dict, List, Optional, Set
from dataclasses import dataclass, field


class Role(str, Enum):
    """User roles in the system."""
    ADMIN = "admin"
    ENGINEER = "engineer"
    VIEWER = "viewer"
    API_CLIENT = "api_client"


class Permission(str, Enum):
    """Available permissions in the system."""
    # Job permissions
    CREATE_JOB = "create_job"
    VIEW_JOB = "view_job"
    CANCEL_JOB = "cancel_job"
    DELETE_JOB = "delete_job"
    
    # Report permissions
    VIEW_REPORT = "view_report"
    EXPORT_REPORT = "export_report"
    CREATE_REPORT = "create_report"
    
    # API Key permissions
    MANAGE_API_KEYS = "manage_api_keys"
    VIEW_API_KEYS = "view_api_keys"
    
    # Monitoring permissions
    VIEW_MONITORING = "view_monitoring"
    CONFIGURE_ALERTS = "configure_alerts"
    
    # User management
    MANAGE_USERS = "manage_users"
    VIEW_USERS = "view_users"
    
    # Tenant management
    MANAGE_TENANT = "manage_tenant"
    VIEW_TENANT = "view_tenant"
    
    # Release validation
    APPROVE_RELEASE = "approve_release"
    VIEW_RELEASE = "view_release"
    
    # Model/Dataset management
    REGISTER_MODEL = "register_model"
    REGISTER_DATASET = "register_dataset"
    
    # Worker management
    MANAGE_WORKERS = "manage_workers"


# Permission matrix: Role -> Set of Permissions
ROLE_PERMISSIONS: Dict[Role, Set[Permission]] = {
    Role.ADMIN: {
        # Full control
        Permission.CREATE_JOB,
        Permission.VIEW_JOB,
        Permission.CANCEL_JOB,
        Permission.DELETE_JOB,
        Permission.VIEW_REPORT,
        Permission.EXPORT_REPORT,
        Permission.CREATE_REPORT,
        Permission.MANAGE_API_KEYS,
        Permission.VIEW_API_KEYS,
        Permission.VIEW_MONITORING,
        Permission.CONFIGURE_ALERTS,
        Permission.MANAGE_USERS,
        Permission.VIEW_USERS,
        Permission.MANAGE_TENANT,
        Permission.VIEW_TENANT,
        Permission.APPROVE_RELEASE,
        Permission.VIEW_RELEASE,
        Permission.REGISTER_MODEL,
        Permission.REGISTER_DATASET,
        Permission.MANAGE_WORKERS,
    },
    Role.ENGINEER: {
        # Run benchmarks, view reports
        Permission.CREATE_JOB,
        Permission.VIEW_JOB,
        Permission.CANCEL_JOB,
        Permission.VIEW_REPORT,
        Permission.EXPORT_REPORT,
        Permission.VIEW_API_KEYS,
        Permission.VIEW_MONITORING,
        Permission.VIEW_USERS,
        Permission.VIEW_TENANT,
        Permission.VIEW_RELEASE,
        Permission.REGISTER_MODEL,
        Permission.REGISTER_DATASET,
    },
    Role.VIEWER: {
        # View dashboards only
        Permission.VIEW_JOB,
        Permission.VIEW_REPORT,
        Permission.VIEW_MONITORING,
        Permission.VIEW_USERS,
        Permission.VIEW_TENANT,
        Permission.VIEW_RELEASE,
    },
    Role.API_CLIENT: {
        # API-only access (limited)
        Permission.CREATE_JOB,
        Permission.VIEW_JOB,
    },
}


@dataclass
class RBACContext:
    """
    Context for RBAC checks containing user and tenant information.
    """
    user_id: uuid.UUID
    tenant_id: uuid.UUID
    role: Role
    permissions: Set[Permission] = field(init=False)
    
    def __post_init__(self):
        self.permissions = ROLE_PERMISSIONS.get(self.role, set())
    
    def has_permission(self, permission: Permission) -> bool:
        """Check if the user has a specific permission."""
        return permission in self.permissions
    
    def has_any_permission(self, permissions: List[Permission]) -> bool:
        """Check if the user has any of the specified permissions."""
        return any(p in self.permissions for p in permissions)
    
    def has_all_permissions(self, permissions: List[Permission]) -> bool:
        """Check if the user has all of the specified permissions."""
        return all(p in self.permissions for p in permissions)


class RBAC:
    """
    Role-Based Access Control system.
    
    Provides methods to check permissions and enforce access control.
    """
    
    @staticmethod
    def get_permissions_for_role(role: Role) -> Set[Permission]:
        """Get all permissions for a given role."""
        return ROLE_PERMISSIONS.get(role, set()).copy()
    
    @staticmethod
    def get_role_permissions(role: Role) -> List[str]:
        """Get all permission names for a given role."""
        permissions = ROLE_PERMISSIONS.get(role, set())
        return [p.value for p in permissions]
    
    @staticmethod
    def has_permission(role: Role, permission: Permission) -> bool:
        """Check if a role has a specific permission."""
        return permission in ROLE_PERMISSIONS.get(role, set())
    
    @staticmethod
    def can_create_job(role: Role) -> bool:
        """Check if role can create evaluation jobs."""
        return RBAC.has_permission(role, Permission.CREATE_JOB)
    
    @staticmethod
    def can_view_job(role: Role) -> bool:
        """Check if role can view jobs."""
        return RBAC.has_permission(role, Permission.VIEW_JOB)
    
    @staticmethod
    def can_cancel_job(role: Role) -> bool:
        """Check if role can cancel jobs."""
        return RBAC.has_permission(role, Permission.CANCEL_JOB)
    
    @staticmethod
    def can_delete_job(role: Role) -> bool:
        """Check if role can delete jobs."""
        return RBAC.has_permission(role, Permission.DELETE_JOB)
    
    @staticmethod
    def can_export_report(role: Role) -> bool:
        """Check if role can export reports."""
        return RBAC.has_permission(role, Permission.EXPORT_REPORT)
    
    @staticmethod
    def can_manage_api_keys(role: Role) -> bool:
        """Check if role can manage API keys."""
        return RBAC.has_permission(role, Permission.MANAGE_API_KEYS)
    
    @staticmethod
    def can_view_monitoring(role: Role) -> bool:
        """Check if role can view monitoring data."""
        return RBAC.has_permission(role, Permission.VIEW_MONITORING)
    
    @staticmethod
    def can_manage_users(role: Role) -> bool:
        """Check if role can manage users."""
        return RBAC.has_permission(role, Permission.MANAGE_USERS)
    
    @staticmethod
    def can_approve_release(role: Role) -> bool:
        """Check if role can approve releases."""
        return RBAC.has_permission(role, Permission.APPROVE_RELEASE)


def rbac_check(
    role: Role,
    required_permission: Permission,
) -> bool:
    """
    Decorator/function to check if a role has a specific permission.
    
    Usage:
        @rbac_check(Role.ENGINEER, Permission.CREATE_JOB)
        async def create_job(...):
            ...
    """
    return RBAC.has_permission(role, required_permission)


def require_permission(permission: Permission):
    """
    Decorator factory to require a specific permission.
    
    Usage:
        @require_permission(Permission.CREATE_JOB)
        async def create_job(request, ...):
            ...
    """
    def decorator(func):
        async def wrapper(*args, **kwargs):
            # This would be used in FastAPI dependency injection
            # The actual implementation would check request.state.rbac_context
            return await func(*args, **kwargs)
        return wrapper
    return decorator