Coverage for src/lexigram/auth/authn/saml.py: 0%

148 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 12:26 +0800

1"""SAML 2.0 authentication flows""" 

2 

3from __future__ import annotations 

4 

5import threading 

6from typing import Any, Protocol 

7 

8from lexigram.logging import get_logger 

9 

10try: 

11 from saml2 import BINDING_HTTP_POST 

12 

13 # Optional imports that may not be used directly are aliased to avoid unused-import warnings 

14 from saml2 import BINDING_HTTP_REDIRECT as _BINDING_HTTP_REDIRECT 

15 from saml2.client import Saml2Client 

16 from saml2.config import Config as Saml2Config 

17 from saml2.metadata import ( 

18 entity_descriptor as _entity_descriptor, 

19 ) 

20 from saml2.response import AuthnResponse 

21 from saml2.saml import NAMEID_FORMAT_EMAILADDRESS 

22 import xmlsec as _xmlsec 

23 

24 HAS_SAML = True 

25 # Use no-op references to aliased optional imports so static linters don't flag them 

26 _ = _xmlsec 

27 _ = _BINDING_HTTP_REDIRECT 

28 _ = _entity_descriptor 

29except ImportError: 

30 HAS_SAML = False 

31 Saml2Client = None 

32 Saml2Config = None 

33 AuthnResponse = None 

34 # Fallback constant for when SAML libraries are not available 

35 NAMEID_FORMAT_EMAILADDRESS = ( 

36 "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress" 

37 ) 

38 

39 

40logger = get_logger(__name__) 

41 

42 

43# ============================================================================= 

44# SAML Attribute Mapper Registry 

45# ============================================================================= 

46 

47 

48class SAMLAttributeMapper(Protocol): 

49 """Protocol for SAML attribute mappers.""" 

50 

51 def map_attribute( 

52 self, 

53 ava: dict[str, list[str]], 

54 user_info: dict[str, Any], 

55 ) -> None: 

56 """Map SAML attributes to user info.""" 

57 ... 

58 

59 

60class EmailAttributeMapper: 

61 """Maps SAML email attribute to user info.""" 

62 

63 def map_attribute( 

64 self, 

65 ava: dict[str, list[str]], 

66 user_info: dict[str, Any], 

67 ) -> None: 

68 if "email" in ava: 

69 user_info["email"] = ava["email"][0] 

70 

71 

72class NameAttributeMapper: 

73 """Maps SAML name attributes to user info.""" 

74 

75 def map_attribute( 

76 self, 

77 ava: dict[str, list[str]], 

78 user_info: dict[str, Any], 

79 ) -> None: 

80 if "givenName" in ava and "sn" in ava: 

81 user_info["name"] = f"{ava['givenName'][0]} {ava['sn'][0]}" 

82 elif "displayName" in ava: 

83 user_info["name"] = ava["displayName"][0] 

84 

85 

86class FirstNameAttributeMapper: 

87 """Maps SAML firstName/givenName attribute to user info.""" 

88 

89 def map_attribute( 

90 self, 

91 ava: dict[str, list[str]], 

92 user_info: dict[str, Any], 

93 ) -> None: 

94 if "givenName" in ava: 

95 user_info["first_name"] = ava["givenName"][0] 

96 elif "firstName" in ava: 

97 user_info["first_name"] = ava["firstName"][0] 

98 

99 

100class LastNameAttributeMapper: 

101 """Maps SAML lastName/surname attribute to user info.""" 

102 

103 def map_attribute( 

104 self, 

105 ava: dict[str, list[str]], 

106 user_info: dict[str, Any], 

107 ) -> None: 

108 if "sn" in ava: 

109 user_info["last_name"] = ava["sn"][0] 

110 elif "surname" in ava: 

111 user_info["last_name"] = ava["surname"][0] 

112 

113 

114class GroupsAttributeMapper: 

115 """Maps SAML groups attribute to user info.""" 

116 

117 def map_attribute( 

118 self, 

119 ava: dict[str, list[str]], 

120 user_info: dict[str, Any], 

121 ) -> None: 

122 if "groups" in ava: 

123 user_info["groups"] = ava["groups"] 

124 elif "memberOf" in ava: 

125 user_info["groups"] = ava["memberOf"] 

126 

127 

128class SAMLAttributeMapperRegistry: 

129 """Registry for SAML attribute mappers. 

130 

131 Allows dynamic registration of custom attribute mappers for 

132 different SAML identity providers. 

133 """ 

134 

135 def __init__(self) -> None: 

136 self._lock = threading.Lock() 

137 self._mappers: list[SAMLAttributeMapper] = [] 

138 

139 @classmethod 

140 def with_defaults(cls) -> SAMLAttributeMapperRegistry: 

141 """Create a registry pre-loaded with the standard attribute mappers.""" 

142 instance = cls() 

143 instance._register_default_mappers() 

144 return instance 

145 

146 def _register_default_mappers(self) -> None: 

147 """Register the default attribute mappers.""" 

148 self._mappers = [ 

149 EmailAttributeMapper(), 

150 NameAttributeMapper(), 

151 FirstNameAttributeMapper(), 

152 LastNameAttributeMapper(), 

153 GroupsAttributeMapper(), 

154 ] 

155 

156 def register_mapper(self, mapper: SAMLAttributeMapper) -> None: 

157 """Register a custom attribute mapper.""" 

158 with self._lock: 

159 self._mappers.append(mapper) 

160 

161 def clear_mappers(self) -> None: 

162 """Clear all registered mappers.""" 

163 with self._lock: 

164 self._mappers.clear() 

165 

166 def get_mappers(self) -> list[SAMLAttributeMapper]: 

167 """Get all registered mappers.""" 

168 with self._lock: 

169 return list(self._mappers) 

170 

171 def map_attributes( 

172 self, 

173 ava: dict[str, list[str]], 

174 user_info: dict[str, Any], 

175 ) -> None: 

176 """Map all SAML attributes using registered mappers.""" 

177 with self._lock: 

178 mappers = list(self._mappers) 

179 for mapper in mappers: 

180 mapper.map_attribute(ava, user_info) 

181 

182 

183# Global registry instance 

184_saml_attribute_registry = SAMLAttributeMapperRegistry.with_defaults() 

185 

186 

187class SAMLProvider: 

188 """SAML identity provider configuration""" 

189 

190 def __init__( 

191 self, 

192 name: str, 

193 entity_id: str, 

194 sso_url: str, 

195 slo_url: str | None = None, 

196 x509_cert: str | None = None, 

197 name_id_format: str = NAMEID_FORMAT_EMAILADDRESS, 

198 want_assertions_signed: bool = True, 

199 want_response_signed: bool = True, 

200 want_logout_response_signed: bool = False, 

201 want_logout_request_signed: bool = False, 

202 ): 

203 self.name = name 

204 self.entity_id = entity_id 

205 self.sso_url = sso_url 

206 self.slo_url = slo_url 

207 self.x509_cert = x509_cert 

208 self.name_id_format = name_id_format 

209 self.want_assertions_signed = want_assertions_signed 

210 self.want_response_signed = want_response_signed 

211 self.want_logout_response_signed = want_logout_response_signed 

212 self.want_logout_request_signed = want_logout_request_signed 

213 

214 

215class SAMLManager: 

216 """SAML 2.0 authentication manager""" 

217 

218 def __init__( 

219 self, 

220 providers: dict[str, SAMLProvider], 

221 http_client: Any | None = None, 

222 ) -> None: 

223 if not HAS_SAML: 

224 raise ImportError( 

225 "SAML libraries (pysaml2, xmlsec) are required for SAML functionality", 

226 ) 

227 

228 self.providers = providers 

229 self._http_client = http_client 

230 self._clients = {} 

231 

232 # Initialize SAML clients for each provider 

233 for name, provider in providers.items(): 

234 self._clients[name] = self._create_saml_client(provider) 

235 

236 def __repr__(self) -> str: 

237 """Return developer-friendly string representation.""" 

238 return f"SAMLManager(providers={list(self.providers)!r})" 

239 

240 def _create_saml_client(self, provider: SAMLProvider) -> Saml2Client: 

241 """Create SAML client for a provider""" 

242 config = Saml2Config() 

243 

244 # Service provider configuration 

245 config.load( 

246 { 

247 "entityid": f"lexigram-admin-{provider.name}", 

248 "service": { 

249 "sp": { 

250 "name": f"Lexigram Admin - {provider.name}", 

251 "name_id_format": provider.name_id_format, 

252 "want_assertions_signed": provider.want_assertions_signed, 

253 "want_response_signed": provider.want_response_signed, 

254 "want_logout_response_signed": provider.want_logout_response_signed, 

255 "want_logout_request_signed": provider.want_logout_request_signed, 

256 }, 

257 }, 

258 "metadata": { 

259 "remote": [ 

260 { 

261 "url": provider.entity_id, 

262 "cert": provider.x509_cert, 

263 }, 

264 ], 

265 }, 

266 "key_file": None, # SP doesn't need signing key for basic auth 

267 "cert_file": None, 

268 }, 

269 ) 

270 

271 return Saml2Client(config) 

272 

273 async def get_login_url( 

274 self, 

275 provider_name: str, 

276 relay_state: str | None = None, 

277 ) -> str: 

278 """Get SAML login URL for the given provider""" 

279 client = self._clients.get(provider_name) 

280 if not client: 

281 raise ValueError(f"SAML provider '{provider_name}' not configured") 

282 

283 _req_id, info = client.prepare_for_authenticate(relay_state=relay_state) 

284 

285 # Return the login URL 

286 url = info.get("url") 

287 if isinstance(url, str): 

288 return url 

289 if url is not None: 

290 return str(url) 

291 return "" # Fallback to empty string if not available 

292 

293 async def process_assertion( 

294 self, 

295 provider_name: str, 

296 saml_response: str, 

297 relay_state: str | None = None, 

298 ) -> dict[str, Any]: 

299 """Process SAML assertion response""" 

300 client = self._clients.get(provider_name) 

301 if not client: 

302 raise ValueError(f"SAML provider '{provider_name}' not configured") 

303 

304 # Parse and validate the SAML response 

305 authn_response = client.parse_authn_request_response( 

306 saml_response, 

307 BINDING_HTTP_POST, 

308 ) 

309 

310 if authn_response is None: 

311 raise ValueError("Invalid SAML response") 

312 

313 # Extract user information 

314 user_info = { 

315 "name_id": authn_response.name_id, 

316 "name_id_format": authn_response.name_id_format, 

317 "session_index": authn_response.session_index, 

318 "attributes": authn_response.ava, # Attribute value assertions 

319 } 

320 

321 # Map common attributes using the registry 

322 _saml_attribute_registry.map_attributes(authn_response.ava, user_info) 

323 

324 return user_info 

325 

326 async def get_logout_url( 

327 self, 

328 provider_name: str, 

329 name_id: str, 

330 session_index: str | None = None, 

331 ) -> str | None: 

332 """Get SAML logout URL for the given provider""" 

333 client = self._clients.get(provider_name) 

334 if not client: 

335 raise ValueError(f"SAML provider '{provider_name}' not configured") 

336 

337 if not client.slo_service_urls: 

338 return None 

339 

340 # Prepare logout request 

341 _slo_req = client.create_logout_request( 

342 name_id=name_id, 

343 session_index=session_index, 

344 ) 

345 

346 # Get logout URL 

347 _binding, slo_url = next(iter(client.slo_service_urls.items())) 

348 if isinstance(slo_url, str): 

349 return slo_url 

350 if slo_url is not None: 

351 return str(slo_url) 

352 return None 

353 

354 async def process_logout_response( 

355 self, 

356 provider_name: str, 

357 saml_response: str, 

358 ) -> bool: 

359 """Process SAML logout response""" 

360 client = self._clients.get(provider_name) 

361 if not client: 

362 raise ValueError(f"SAML provider '{provider_name}' not configured") 

363 

364 # Parse logout response 

365 logout_response = client.parse_logout_request_response( 

366 saml_response, 

367 BINDING_HTTP_POST, 

368 ) 

369 

370 return logout_response is not None 

371 

372 

373__all__ = [ 

374 "EmailAttributeMapper", 

375 "FirstNameAttributeMapper", 

376 "GroupsAttributeMapper", 

377 "LastNameAttributeMapper", 

378 "NameAttributeMapper", 

379 "SAMLAttributeMapper", 

380 "SAMLAttributeMapperRegistry", 

381 "SAMLManager", 

382 "SAMLProvider", 

383]