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