Coverage for src/alprina_cli/auth_system.py: 82%

1""" 

2Authentication & Authorization System 

3 

4Provides: 

5- User authentication (API key based) 

6- Role-Based Access Control (RBAC) 

7- Audit logging for security operations 

8- Session management 

9 

10Context Engineering: 

11- Lightweight auth that doesn't inflate context 

12- Fast permission checks (< 1ms) 

13- Audit logs for compliance 

14""" 

15 

16from typing import Optional, List, Dict, Any 

17from datetime import datetime, timedelta 

18from enum import Enum 

19import hashlib 

20import secrets 

21from pydantic import BaseModel, Field 

22from loguru import logger 

23 

24 

25class Role(str, Enum): 

26 """User roles for RBAC""" 

27 ADMIN = "admin" # Full access to all operations 

28 SECURITY_ANALYST = "security_analyst" # Can run all security tools 

29 PENTESTER = "pentester" # Can run offensive tools (red team, exploit) 

30 DEFENDER = "defender" # Can run defensive tools (blue team, DFIR) 

31 AUDITOR = "auditor" # Read-only access, can view reports 

32 USER = "user" # Basic access to scan and recon 

33 

34 

35class Permission(str, Enum): 

36 """Fine-grained permissions""" 

37 # Tool permissions 

38 SCAN = "scan" 

39 RECON = "recon" 

40 VULN_SCAN = "vuln_scan" 

41 EXPLOIT = "exploit" 

42 RED_TEAM = "red_team" 

43 BLUE_TEAM = "blue_team" 

44 DFIR = "dfir" 

45 ANDROID_SAST = "android_sast" 

46 

47 # Administrative permissions 

48 MANAGE_USERS = "manage_users" 

49 VIEW_AUDIT_LOGS = "view_audit_logs" 

50 MANAGE_ROLES = "manage_roles" 

51 

52 # Report permissions 

53 GENERATE_REPORTS = "generate_reports" 

54 VIEW_REPORTS = "view_reports" 

55 

56 

57# Role to permissions mapping 

58ROLE_PERMISSIONS: Dict[Role, List[Permission]] = { 

59 Role.ADMIN: list(Permission), # All permissions 

60 Role.SECURITY_ANALYST: [ 

61 Permission.SCAN, 

62 Permission.RECON, 

63 Permission.VULN_SCAN, 

64 Permission.BLUE_TEAM, 

65 Permission.DFIR, 

66 Permission.ANDROID_SAST, 

67 Permission.GENERATE_REPORTS, 

68 Permission.VIEW_REPORTS, 

69 Permission.VIEW_AUDIT_LOGS, 

70 ], 

71 Role.PENTESTER: [ 

72 Permission.SCAN, 

73 Permission.RECON, 

74 Permission.VULN_SCAN, 

75 Permission.EXPLOIT, 

76 Permission.RED_TEAM, 

77 Permission.VIEW_REPORTS, 

78 ], 

79 Role.DEFENDER: [ 

80 Permission.SCAN, 

81 Permission.RECON, 

82 Permission.BLUE_TEAM, 

83 Permission.DFIR, 

84 Permission.GENERATE_REPORTS, 

85 Permission.VIEW_REPORTS, 

86 ], 

87 Role.AUDITOR: [ 

88 Permission.VIEW_REPORTS, 

89 Permission.VIEW_AUDIT_LOGS, 

90 ], 

91 Role.USER: [ 

92 Permission.SCAN, 

93 Permission.RECON, 

94 Permission.VIEW_REPORTS, 

95 ], 

96} 

97 

98 

99class User(BaseModel): 

100 """User model""" 

101 user_id: str 

102 username: str 

103 email: str 

104 role: Role 

105 api_key_hash: Optional[str] = None 

106 created_at: datetime = Field(default_factory=datetime.utcnow) 

107 last_login: Optional[datetime] = None 

108 is_active: bool = True 

109 

110 

111class AuditLogEntry(BaseModel): 

112 """Audit log entry for security operations""" 

113 timestamp: datetime = Field(default_factory=datetime.utcnow) 

114 user_id: str 

115 username: str 

116 operation: str 

117 tool_name: str 

118 target: Optional[str] = None 

119 success: bool 

120 details: Optional[Dict[str, Any]] = None 

121 ip_address: Optional[str] = None 

122 

123 

124class AuthenticationService: 

125 """ 

126 Authentication service for API key management. 

127 

128 Context: Fast, stateless authentication using API keys. 

129 """ 

130 

131 def __init__(self, use_database: bool = True): 

132 # In-memory storage (fallback when database unavailable) 

133 self._users: Dict[str, User] = {} 

134 self._api_keys: Dict[str, str] = {} # api_key -> user_id 

135 

136 # Database integration 

137 self._use_database = use_database 

138 self._db_client = None 

139 

140 if use_database: 

141 try: 

142 from alprina_cli.database.neon_client import get_database_client 

143 self._db_client = get_database_client() 

144 logger.info("AuthenticationService using database backend") 

145 except Exception as e: 

146 logger.warning(f"Database unavailable, using in-memory storage: {e}") 

147 self._use_database = False 

148 

149 def generate_api_key(self) -> str: 

150 """Generate a secure API key""" 

151 return f"alprina_{secrets.token_urlsafe(32)}" 

152 

153 def hash_api_key(self, api_key: str) -> str: 

154 """Hash API key for storage""" 

155 return hashlib.sha256(api_key.encode()).hexdigest() 

156 

157 def create_user( 

158 self, 

159 username: str, 

160 email: str, 

161 role: Role = Role.USER 

162 ) -> tuple[User, str]: 

163 """ 

164 Create a new user and return user object + API key. 

165 

166 Args: 

167 username: Username 

168 email: Email address 

169 role: User role (default: USER) 

170 

171 Returns: 

172 Tuple of (User, api_key) 

173 """ 

174 user_id = f"user_{secrets.token_hex(8)}" 

175 api_key = self.generate_api_key() 

176 api_key_hash = self.hash_api_key(api_key) 

177 

178 user = User( 

179 user_id=user_id, 

180 username=username, 

181 email=email, 

182 role=role, 

183 api_key_hash=api_key_hash 

184 ) 

185 

186 # Store in-memory (always keep for backward compatibility) 

187 self._users[user_id] = user 

188 self._api_keys[api_key] = user_id 

189 

190 logger.info(f"Created user {username} ({user_id}) with role {role}") 

191 

192 return user, api_key 

193 

194 async def authenticate(self, api_key: str) -> Optional[User]: 

195 """ 

196 Authenticate user by API key. 

197 

198 Args: 

199 api_key: API key to authenticate 

200 

201 Returns: 

202 User object if authenticated, None otherwise 

203 """ 

204 # Try database first 

205 if self._use_database and self._db_client: 

206 try: 

207 user_data = await self._db_client.authenticate_api_key(api_key) 

208 if user_data: 

209 return User( 

210 user_id=user_data['id'], 

211 username=user_data.get('name', user_data.get('username', 'unknown')), 

212 email=user_data.get('email', ''), 

213 role=Role(user_data.get('role', 'user')), 

214 last_login=datetime.utcnow(), 

215 is_active=True 

216 ) 

217 except Exception as e: 

218 logger.warning(f"Database authentication failed, falling back to in-memory: {e}") 

219 

220 # Fallback to in-memory 

221 user_id = self._api_keys.get(api_key) 

222 if not user_id: 

223 logger.warning("Authentication failed: Invalid API key") 

224 return None 

225 

226 user = self._users.get(user_id) 

227 if not user or not user.is_active: 

228 logger.warning(f"Authentication failed: User {user_id} not found or inactive") 

229 return None 

230 

231 # Update last login 

232 user.last_login = datetime.utcnow() 

233 

234 logger.debug(f"Authenticated user {user.username} ({user_id})") 

235 return user 

236 

237 def revoke_api_key(self, api_key: str) -> bool: 

238 """ 

239 Revoke an API key. 

240 

241 Args: 

242 api_key: API key to revoke 

243 

244 Returns: 

245 True if revoked, False if not found 

246 """ 

247 if api_key in self._api_keys: 

248 user_id = self._api_keys[api_key] 

249 del self._api_keys[api_key] 

250 logger.info(f"Revoked API key for user {user_id}") 

251 return True 

252 return False 

253 

254 def deactivate_user(self, user_id: str) -> bool: 

255 """ 

256 Deactivate a user account. 

257 

258 Args: 

259 user_id: User ID to deactivate 

260 

261 Returns: 

262 True if deactivated, False if not found 

263 """ 

264 user = self._users.get(user_id) 

265 if user: 

266 user.is_active = False 

267 logger.info(f"Deactivated user {user.username} ({user_id})") 

268 return True 

269 return False 

270 

271 

272class AuthorizationService: 

273 """ 

274 Authorization service for RBAC. 

275 

276 Context: Fast permission checks without context overhead. 

277 """ 

278 

279 def __init__(self): 

280 self.role_permissions = ROLE_PERMISSIONS 

281 

282 def has_permission(self, user: User, permission: Permission) -> bool: 

283 """ 

284 Check if user has a specific permission. 

285 

286 Args: 

287 user: User to check 

288 permission: Permission to check 

289 

290 Returns: 

291 True if user has permission, False otherwise 

292 """ 

293 user_permissions = self.role_permissions.get(user.role, []) 

294 has_perm = permission in user_permissions 

295 

296 if not has_perm: 

297 logger.warning( 

298 f"Permission denied: User {user.username} ({user.role}) " 

299 f"does not have {permission} permission" 

300 ) 

301 

302 return has_perm 

303 

304 def require_permission(self, user: User, permission: Permission) -> None: 

305 """ 

306 Require a permission, raise exception if not granted. 

307 

308 Args: 

309 user: User to check 

310 permission: Required permission 

311 

312 Raises: 

313 PermissionError: If user doesn't have permission 

314 """ 

315 if not self.has_permission(user, permission): 

316 raise PermissionError( 

317 f"User {user.username} does not have {permission} permission" 

318 ) 

319 

320 def get_user_permissions(self, user: User) -> List[Permission]: 

321 """Get all permissions for a user""" 

322 return self.role_permissions.get(user.role, []) 

323 

324 def can_use_tool(self, user: User, tool_name: str) -> bool: 

325 """ 

326 Check if user can use a specific tool. 

327 

328 Args: 

329 user: User to check 

330 tool_name: Name of tool to check 

331 

332 Returns: 

333 True if user can use tool, False otherwise 

334 """ 

335 # Map tool names to permissions 

336 tool_permission_map = { 

337 "ScanTool": Permission.SCAN, 

338 "ReconTool": Permission.RECON, 

339 "VulnScanTool": Permission.VULN_SCAN, 

340 "ExploitTool": Permission.EXPLOIT, 

341 "RedTeamTool": Permission.RED_TEAM, 

342 "BlueTeamTool": Permission.BLUE_TEAM, 

343 "DFIRTool": Permission.DFIR, 

344 "AndroidSASTTool": Permission.ANDROID_SAST, 

345 } 

346 

347 permission = tool_permission_map.get(tool_name) 

348 if not permission: 

349 # Unknown tool - deny by default 

350 logger.warning(f"Unknown tool: {tool_name}") 

351 return False 

352 

353 return self.has_permission(user, permission) 

354 

355 

356class AuditLogger: 

357 """ 

358 Audit logger for security operations. 

359 

360 Context: Compliance-ready audit logging. 

361 """ 

362 

363 def __init__(self, max_entries: int = 10000): 

364 # In-memory storage (replace with database in production) 

365 self.audit_log: List[AuditLogEntry] = [] 

366 self.max_entries = max_entries 

367 

368 def log( 

369 self, 

370 user: User, 

371 operation: str, 

372 tool_name: str, 

373 target: Optional[str] = None, 

374 success: bool = True, 

375 details: Optional[Dict[str, Any]] = None, 

376 ip_address: Optional[str] = None 

377 ) -> None: 

378 """ 

379 Log a security operation. 

380 

381 Args: 

382 user: User performing operation 

383 operation: Operation name (e.g., "scan", "exploit") 

384 tool_name: Name of tool used 

385 target: Target of operation (e.g., IP, domain) 

386 success: Whether operation succeeded 

387 details: Additional details 

388 ip_address: IP address of user 

389 """ 

390 entry = AuditLogEntry( 

391 user_id=user.user_id, 

392 username=user.username, 

393 operation=operation, 

394 tool_name=tool_name, 

395 target=target, 

396 success=success, 

397 details=details or {}, 

398 ip_address=ip_address 

399 ) 

400 

401 self.audit_log.append(entry) 

402 

403 # Trim log if it gets too large 

404 if len(self.audit_log) > self.max_entries: 

405 self.audit_log = self.audit_log[-self.max_entries:] 

406 

407 logger.info( 

408 f"AUDIT: {user.username} {operation} {tool_name} " 

409 f"target={target} success={success}" 

410 ) 

411 

412 def get_logs( 

413 self, 

414 user_id: Optional[str] = None, 

415 tool_name: Optional[str] = None, 

416 start_time: Optional[datetime] = None, 

417 end_time: Optional[datetime] = None, 

418 limit: int = 100 

419 ) -> List[AuditLogEntry]: 

420 """ 

421 Query audit logs with filters. 

422 

423 Args: 

424 user_id: Filter by user ID 

425 tool_name: Filter by tool name 

426 start_time: Filter by start time 

427 end_time: Filter by end time 

428 limit: Maximum number of results 

429 

430 Returns: 

431 List of matching audit log entries 

432 """ 

433 results = self.audit_log 

434 

435 if user_id: 

436 results = [e for e in results if e.user_id == user_id] 

437 

438 if tool_name: 

439 results = [e for e in results if e.tool_name == tool_name] 

440 

441 if start_time: 

442 results = [e for e in results if e.timestamp >= start_time] 

443 

444 if end_time: 

445 results = [e for e in results if e.timestamp <= end_time] 

446 

447 # Return most recent first 

448 results = sorted(results, key=lambda e: e.timestamp, reverse=True) 

449 

450 return results[:limit] 

451 

452 def get_user_activity(self, user_id: str, days: int = 7) -> List[AuditLogEntry]: 

453 """Get user activity for the past N days""" 

454 start_time = datetime.utcnow() - timedelta(days=days) 

455 return self.get_logs(user_id=user_id, start_time=start_time) 

456 

457 

458# Global instances (singleton pattern) 

459_auth_service: Optional[AuthenticationService] = None 

460_authz_service: Optional[AuthorizationService] = None 

461_audit_logger: Optional[AuditLogger] = None 

462 

463 

464def get_auth_service() -> AuthenticationService: 

465 """Get global authentication service instance""" 

466 global _auth_service 

467 if _auth_service is None: 

468 _auth_service = AuthenticationService() 

469 return _auth_service 

470 

471 

472def get_authz_service() -> AuthorizationService: 

473 """Get global authorization service instance""" 

474 global _authz_service 

475 if _authz_service is None: 

476 _authz_service = AuthorizationService() 

477 return _authz_service 

478 

479 

480def get_audit_logger() -> AuditLogger: 

481 """Get global audit logger instance""" 

482 global _audit_logger 

483 if _audit_logger is None: 

484 _audit_logger = AuditLogger() 

485 return _audit_logger