Coverage for core/rbac.py: 95.06%

1import logging 

2from functools import wraps 

3 

4from fastapi import HTTPException, status 

5from sqlalchemy import func, select 

6from sqlalchemy.ext.asyncio import AsyncSession 

7 

8from core.config import settings 

9from models.role_attributes import RoleAttributes 

10from models.role_attributes_mapper import RoleAttributesMapper 

11from models.role_mapper import RoleMapper 

12from models.roles import Roles 

13from utils.custom_exception import ServerException 

14 

15logger = logging.getLogger("rbac") 

16 

17 

18def is_super_admin_role_name(role_name: str | None) -> bool: 

19 """True when role_name is the ENV-configured system super-admin role.""" 

20 return bool(role_name) and role_name == settings.DEFAULT_SUPER_ADMIN_ROLE 

21 

22 

23async def get_user_role_level(user_id: str, db: AsyncSession) -> int: 

24 """Return the user's highest role level, or 0 when the user has no role.""" 

25 try: 

26 result = await db.execute( 

27 select(func.max(Roles.level)) 

28 .join(RoleMapper, Roles.id == RoleMapper.role_id) 

29 .where(RoleMapper.user_id == user_id) 

30 ) 

31 level = result.scalar_one_or_none() 

32 return int(level) if level is not None else 0 

33 except Exception as e: 

34 raise ServerException(f"Failed to get user role level: {e}") 

35 

36 

37async def get_user_role_id(user_id: str, db: AsyncSession) -> str | None: 

38 """Return the user's primary role id (highest level), or None when unassigned.""" 

39 try: 

40 result = await db.execute( 

41 select(Roles.id) 

42 .join(RoleMapper, Roles.id == RoleMapper.role_id) 

43 .where(RoleMapper.user_id == user_id) 

44 .order_by(Roles.level.desc(), Roles.name.asc()) 

45 .limit(1) 

46 ) 

47 return result.scalar_one_or_none() 

48 except Exception as e: 

49 raise ServerException(f"Failed to get user role id: {e}") 

50 

51 

52async def user_has_role(user_id: str, role_id: str, db: AsyncSession) -> bool: 

53 """True when the user is currently assigned the given role.""" 

54 try: 

55 result = await db.execute( 

56 select(RoleMapper.role_id).where( 

57 RoleMapper.user_id == user_id, 

58 RoleMapper.role_id == role_id, 

59 ) 

60 ) 

61 return result.scalar_one_or_none() is not None 

62 except Exception as e: 

63 raise ServerException(f"Failed to check user role assignment: {e}") 

64 

65 

66async def get_user_attributes(user_id: str, db: AsyncSession) -> dict[str, bool]: 

67 """Get user attributes""" 

68 try: 

69 # Check if user has super admin role first 

70 if await check_user_has_super_role(user_id, db): 

71 # Super admin has all attributes - get all available attributes 

72 all_attributes_result = await db.execute(select(RoleAttributes.name)) 

73 all_attributes = [row.name for row in all_attributes_result] 

74 return {attr: True for attr in all_attributes} 

75 

76 result = await db.execute( 

77 select(RoleAttributes.name, RoleAttributesMapper.value) 

78 .join(RoleAttributesMapper, RoleAttributes.id == RoleAttributesMapper.attributes_id) 

79 .join(RoleMapper, RoleMapper.role_id == RoleAttributesMapper.role_id) 

80 .where(RoleMapper.user_id == user_id) 

81 ) 

82 

83 attributes = {} 

84 for row in result: 

85 attr_name = row.name 

86 attr_value = row.value 

87 

88 if attr_name in attributes: 

89 attributes[attr_name] = attributes[attr_name] or attr_value 

90 else: 

91 attributes[attr_name] = attr_value 

92 

93 return attributes 

94 except Exception as e: 

95 raise ServerException(f"Failed to get user attributes: {e}") 

96 

97 

98async def check_user_has_super_role(user_id: str, db: AsyncSession) -> bool: 

99 """Check if user has super admin role""" 

100 try: 

101 result = await db.execute( 

102 select(Roles.name) 

103 .join(RoleMapper, Roles.id == RoleMapper.role_id) 

104 .where(RoleMapper.user_id == user_id) 

105 ) 

106 

107 user_roles = [row.name for row in result] 

108 return settings.DEFAULT_SUPER_ADMIN_ROLE in user_roles 

109 except Exception as e: 

110 raise ServerException(f"Failed to check super role: {e}") 

111 

112 

113def require_permission(required_attributes: list[str]): 

114 """Permission check decorator""" 

115 

116 def decorator(func): 

117 @wraps(func) 

118 async def wrapper(*args, **kwargs): 

119 token = kwargs.get("token") 

120 db = kwargs.get("db") 

121 

122 if not db: 

123 raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) 

124 

125 user_id = token.get("sub") 

126 

127 # Check if user has super admin role first 

128 if await check_user_has_super_role(user_id, db): 

129 return await func(*args, **kwargs) 

130 

131 # If not super admin, check specific permissions 

132 user_attributes = await get_user_attributes(user_id, db) 

133 

134 # Convert Permission enum to string value if needed 

135 def get_attr_value(attr): 

136 """Convert Permission enum to string value, or return as-is if already a string""" 

137 if hasattr(attr, "value"): 

138 return attr.value 

139 return str(attr) if attr else attr 

140 

141 # Check if the user has at least one of the required permissions 

142 attr_values = [get_attr_value(attr) for attr in required_attributes] 

143 has_permission = any( 

144 user_attributes.get(attr_value, False) for attr_value in attr_values 

145 ) 

146 

147 if not has_permission: 

148 logger.warning( 

149 f"Permission denied for user {user_id}. " 

150 f"Required: {attr_values}, " 

151 f"User has: {[k for k, v in user_attributes.items() if v]}" 

152 ) 

153 raise HTTPException( 

154 status_code=status.HTTP_403_FORBIDDEN, detail="Permission denied" 

155 ) 

156 

157 return await func(*args, **kwargs) 

158 

159 return wrapper 

160 

161 return decorator