Coverage for src/lexigram/admin/controllers/auth/mfa.py: 85%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, Any 

4from urllib.parse import quote_plus 

5 

6from starlette.requests import Request 

7from starlette.responses import HTMLResponse, RedirectResponse 

8 

9from lexigram.admin.controllers.auth.core import ( 

10 _CACHE_CONTROL_NO_STORE, 

11 AuthCoreMixin, 

12 _humanize_error, 

13 logger, 

14) 

15from lexigram.admin.lib.template import render_mfa_challenge_page, render_mfa_setup_page 

16from lexigram.contracts.web import get, post 

17 

18if TYPE_CHECKING: 

19 from lexigram.admin.auth.protocols import ( 

20 AdminAuthServiceProtocol, 

21 AdminCsrfServiceProtocol, 

22 AdminEmailOtpServiceProtocol, 

23 AdminEmailVerificationServiceProtocol, 

24 AdminMfaServiceProtocol, 

25 ) 

26 from lexigram.admin.observability.admin_metrics import AdminMetrics 

27 

28 

29class AuthMfaMixin(AuthCoreMixin): 

30 """AuthController mfa endpoints.""" 

31 

32 _auth_service: AdminAuthServiceProtocol 

33 _csrf_service: AdminCsrfServiceProtocol 

34 _email_otp_service: AdminEmailOtpServiceProtocol | None 

35 _email_verification_service: AdminEmailVerificationServiceProtocol | None 

36 _metrics: AdminMetrics 

37 _mfa_service: AdminMfaServiceProtocol | None 

38 

39 @get("/login/2fa") 

40 async def mfa_challenge_form( 

41 self, request: Request 

42 ) -> HTMLResponse | RedirectResponse: 

43 """Display the standalone TOTP challenge form. 

44 

45 Only reachable when a pending 2FA challenge exists in the session 

46 (parked by ``login_submit``); otherwise the user is redirected to 

47 the login page. A fresh CSRF token is embedded in the form. 

48 

49 Args: 

50 request: Incoming HTTP request. 

51 

52 Returns: 

53 HTMLResponse with the rendered challenge page, or a 

54 RedirectResponse to /admin/ or /admin/login. 

55 """ 

56 user = getattr(request.state, "user", None) 

57 if user and user.user_id != "guest": 

58 return RedirectResponse(url="/admin/", status_code=302) 

59 

60 pending_user_id = request.session.get("mfa_pending_user_id", "") 

61 if not pending_user_id: 

62 return RedirectResponse(url="/admin/login", status_code=302) 

63 

64 email = request.session.get("mfa_pending_email", "") 

65 next_url = request.session.get("mfa_pending_next", "/admin/") 

66 factor = request.session.get("mfa_pending_factor", "totp") 

67 error = _humanize_error(request.query_params.get("error", "")) 

68 notice = request.query_params.get("notice", "") 

69 csrf_token = self._fresh_csrf(request) 

70 

71 html = render_mfa_challenge_page( 

72 email=email, 

73 error=error, 

74 csrf_token=csrf_token, 

75 next_url=next_url, 

76 factor=factor, 

77 resend_notice=notice, 

78 ) 

79 return HTMLResponse(content=html, headers=_CACHE_CONTROL_NO_STORE) 

80 

81 @post("/login/2fa") 

82 async def mfa_challenge_submit(self, request: Request) -> RedirectResponse: 

83 """Complete a 2FA challenge and finish the login. 

84 

85 Verifies the CSRF token, delegates code verification to 

86 ``AdminAuthServiceProtocol.complete_mfa_login``, and on success 

87 populates the session exactly like ``login_submit``. On failure 

88 the challenge page is re-shown with an error message. 

89 

90 Args: 

91 request: Incoming HTTP request carrying form data. 

92 

93 Returns: 

94 RedirectResponse to the pending destination on success, or back 

95 to the challenge page with an error query parameter. 

96 """ 

97 form_data = request.scope.get("admin_form_data") or await request.form() 

98 code = str(form_data.get("code", "")) 

99 csrf_token = str(form_data.get("csrf_token", "")) 

100 

101 pending_user_id = request.session.get("mfa_pending_user_id", "") 

102 if not pending_user_id: 

103 return RedirectResponse(url="/admin/login", status_code=302) 

104 

105 csrf_session_id = request.session.get("csrf_session_id", "") 

106 if not csrf_session_id or not self._csrf_service.validate_token( 

107 csrf_session_id, csrf_token 

108 ): 

109 logger.warning( 

110 "auth.csrf_validation_failed", ip=self._get_client_ip(request) 

111 ) 

112 return RedirectResponse( 

113 url=f"/admin/login/2fa?error={quote_plus('Invalid or expired security token. Please try again.')}", 

114 status_code=302, 

115 ) 

116 

117 if not code: 

118 return RedirectResponse( 

119 url=f"/admin/login/2fa?error={quote_plus('Verification code is required.')}", 

120 status_code=302, 

121 ) 

122 

123 ip = self._get_client_ip(request) 

124 user_agent = request.headers.get("user-agent", "") 

125 

126 result = await self._auth_service.complete_mfa_login( 

127 user_id=pending_user_id, 

128 email=request.session.get("mfa_pending_email", ""), 

129 roles=request.session.get("mfa_pending_roles", []), 

130 code=code, 

131 ip_address=ip, 

132 user_agent=user_agent, 

133 ) 

134 

135 if result.is_err(): 

136 self._metrics.record_login(status="failure") 

137 error_msg = str(result.unwrap_err()) 

138 logger.warning( 

139 "auth.mfa_code_failed", user_id=pending_user_id, reason=error_msg 

140 ) 

141 return RedirectResponse( 

142 url=f"/admin/login/2fa?error={quote_plus(error_msg)}", 

143 status_code=302, 

144 ) 

145 

146 auth_result = result.unwrap() 

147 next_url = request.session.get("mfa_pending_next", "/admin/") 

148 request.session["admin_user_id"] = auth_result.user_id 

149 request.session["admin_user_email"] = auth_result.email 

150 request.session["admin_session_expires_at"] = auth_result.expires_at.isoformat() 

151 request.session["session_id"] = auth_result.session_id 

152 request.session.pop("csrf_session_id", None) 

153 for key in ( 

154 "mfa_pending_user_id", 

155 "mfa_pending_email", 

156 "mfa_pending_roles", 

157 "mfa_pending_next", 

158 "mfa_pending_factor", 

159 ): 

160 request.session.pop(key, None) 

161 self._metrics.record_login(status="success") 

162 logger.info( 

163 "auth.login_success", 

164 user_id=auth_result.user_id, 

165 email=auth_result.email, 

166 redirect=next_url, 

167 ) 

168 return RedirectResponse(url=self._safe_next_url(next_url), status_code=302) 

169 

170 @post("/login/2fa/resend") 

171 async def mfa_challenge_resend(self, request: Request) -> RedirectResponse: 

172 """Resend the email verification code for a pending challenge. 

173 

174 Only valid for the email factor and while a challenge is parked in 

175 the session. Cooldown and delivery errors (e.g. "please wait") are 

176 surfaced on the challenge page. 

177 

178 Args: 

179 request: Incoming HTTP request carrying form data. 

180 

181 Returns: 

182 RedirectResponse to /admin/login/2fa with a notice on success or 

183 an error query parameter on failure. 

184 """ 

185 pending_user_id = request.session.get("mfa_pending_user_id", "") 

186 if not pending_user_id: 

187 return RedirectResponse(url="/admin/login", status_code=302) 

188 

189 form_data = request.scope.get("admin_form_data") or await request.form() 

190 csrf_token = str(form_data.get("csrf_token", "")) 

191 csrf_session_id = request.session.get("csrf_session_id", "") 

192 if not csrf_session_id or not self._csrf_service.validate_token( 

193 csrf_session_id, csrf_token 

194 ): 

195 logger.warning( 

196 "auth.csrf_validation_failed", ip=self._get_client_ip(request) 

197 ) 

198 return RedirectResponse( 

199 url=f"/admin/login/2fa?error={quote_plus('Invalid or expired security token. Please try again.')}", 

200 status_code=302, 

201 ) 

202 

203 factor = request.session.get("mfa_pending_factor", "totp") 

204 if factor != "email" or self._email_otp_service is None: 

205 return RedirectResponse( 

206 url=f"/admin/login/2fa?error={quote_plus('Resending codes is not available for this factor.')}", 

207 status_code=302, 

208 ) 

209 

210 email = str(form_data.get("email", "")) or request.session.get( 

211 "mfa_pending_email", "" 

212 ) 

213 result = await self._email_otp_service.send_otp( 

214 user_id=pending_user_id, 

215 email=email, 

216 user_name=email, 

217 ) 

218 if result.is_err(): 

219 logger.warning( 

220 "auth.email_otp_resend_failed", 

221 user_id=pending_user_id, 

222 reason=str(result.unwrap_err()), 

223 ) 

224 return RedirectResponse( 

225 url=f"/admin/login/2fa?error={quote_plus(_humanize_error(str(result.unwrap_err())))}", 

226 status_code=302, 

227 ) 

228 return RedirectResponse( 

229 url="/admin/login/2fa?notice=" + quote_plus("A new code has been sent."), 

230 status_code=302, 

231 ) 

232 

233 @get("/profile/mfa") 

234 async def mfa_profile_form( 

235 self, request: Request 

236 ) -> HTMLResponse | RedirectResponse: 

237 """Display the authenticated user's 2FA settings page. 

238 

239 Shows the disable form when 2FA is active; otherwise generates a 

240 fresh secret, QR code, and confirm form. Requires an authenticated 

241 request (``request.state.user`` provided by the auth guard). 

242 

243 Args: 

244 request: Incoming HTTP request. 

245 

246 Returns: 

247 HTMLResponse with the rendered setup page, or a RedirectResponse 

248 when unauthenticated or 2FA is disabled in configuration. 

249 """ 

250 user = getattr(request.state, "user", None) 

251 if not user or user.user_id == "guest": 

252 return RedirectResponse( 

253 url="/admin/login?next=/admin/profile/mfa", status_code=302 

254 ) 

255 if self._mfa_service is None: 

256 return RedirectResponse(url="/admin/", status_code=302) 

257 

258 error = _humanize_error(request.query_params.get("error", "")) 

259 notice = request.query_params.get("notice", "") 

260 csrf_token = self._fresh_csrf(request) 

261 user_id = str(user.user_id) 

262 

263 email_verified: bool | None = None 

264 if self._email_verification_service is not None: 

265 email_verified = await self._email_verification_service.is_verified(user_id) 

266 

267 if await self._mfa_service.is_enabled(user_id): 

268 html = render_mfa_setup_page( 

269 enabled=True, 

270 csrf_token=csrf_token, 

271 email_verified=email_verified, 

272 ) 

273 else: 

274 result = await self._mfa_service.start_setup(user_id, str(user.email)) 

275 if result.is_err(): 

276 error = _humanize_error(str(result.unwrap_err())) 

277 html = render_mfa_setup_page( 

278 enabled=False, 

279 csrf_token=csrf_token, 

280 email_verified=email_verified, 

281 ) 

282 else: 

283 secret, _, svg = result.unwrap() 

284 request.session["mfa_pending_secret"] = secret 

285 html = render_mfa_setup_page( 

286 enabled=False, 

287 qr_svg=svg, 

288 secret=secret, 

289 csrf_token=csrf_token, 

290 email_verified=email_verified, 

291 ) 

292 

293 from lexigram.admin.state.context import AdminContextManager 

294 

295 async with AdminContextManager(request) as ctx: 

296 if error: 

297 ctx.add_flash(error, "error") 

298 if notice: 

299 ctx.add_flash(notice, "success") 

300 response = await self.render_admin( 

301 request, 

302 html, 

303 title="Two-Factor Authentication", 

304 breadcrumbs=self.generate_breadcrumbs( 

305 ("Home", "/admin/"), 

306 ("Profile", "/admin/profile"), 

307 current="Two-Factor Authentication", 

308 ), 

309 ) 

310 response.headers.update(_CACHE_CONTROL_NO_STORE) 

311 return response 

312 

313 @post("/profile/mfa/setup") 

314 async def mfa_setup_submit(self, request: Request) -> RedirectResponse: 

315 """Confirm a newly generated TOTP secret. 

316 

317 Validates the code from the setup form against the pending secret 

318 stashed in the session; only a valid code persists the secret. 

319 

320 Args: 

321 request: Incoming HTTP request carrying form data. 

322 

323 Returns: 

324 RedirectResponse to /admin/profile/mfa with a notice on success 

325 or an error query parameter on failure. 

326 """ 

327 user = getattr(request.state, "user", None) 

328 if not user or user.user_id == "guest": 

329 return RedirectResponse( 

330 url="/admin/login?next=/admin/profile/mfa", status_code=302 

331 ) 

332 if self._mfa_service is None: 

333 return RedirectResponse(url="/admin/", status_code=302) 

334 

335 form_data = request.scope.get("admin_form_data") or await request.form() 

336 csrf_token = str(form_data.get("csrf_token", "")) 

337 csrf_session_id = request.session.get("csrf_session_id") or request.session.get( 

338 "admin_user_id", "" 

339 ) 

340 if not csrf_session_id or not self._csrf_service.validate_token( 

341 csrf_session_id, csrf_token 

342 ): 

343 logger.warning( 

344 "auth.csrf_validation_failed", ip=self._get_client_ip(request) 

345 ) 

346 return RedirectResponse( 

347 url=f"/admin/profile/mfa?error={quote_plus('Invalid or expired security token. Please try again.')}", 

348 status_code=302, 

349 ) 

350 

351 code = str(form_data.get("code", "")) 

352 secret = request.session.pop("mfa_pending_secret", "") 

353 if not code or not secret: 

354 return RedirectResponse( 

355 url=f"/admin/profile/mfa?error={quote_plus('Verification code is required.')}", 

356 status_code=302, 

357 ) 

358 

359 result = await self._mfa_service.confirm_setup(str(user.user_id), secret, code) 

360 if result.is_err(): 

361 return RedirectResponse( 

362 url=f"/admin/profile/mfa?error={quote_plus(_humanize_error(str(result.unwrap_err())))}", 

363 status_code=302, 

364 ) 

365 return RedirectResponse( 

366 url="/admin/profile/mfa?notice=" 

367 + quote_plus("Two-factor authentication enabled."), 

368 status_code=302, 

369 ) 

370 

371 @post("/profile/mfa/disable") 

372 async def mfa_disable_submit(self, request: Request) -> RedirectResponse: 

373 """Disable 2FA after validating the current TOTP code. 

374 

375 Requires the user to prove possession of the secret by entering a 

376 valid current code. 

377 

378 Args: 

379 request: Incoming HTTP request carrying form data. 

380 

381 Returns: 

382 RedirectResponse to /admin/profile/mfa with a notice on success 

383 or an error query parameter on failure. 

384 """ 

385 user = getattr(request.state, "user", None) 

386 if not user or user.user_id == "guest": 

387 return RedirectResponse( 

388 url="/admin/login?next=/admin/profile/mfa", status_code=302 

389 ) 

390 if self._mfa_service is None: 

391 return RedirectResponse(url="/admin/", status_code=302) 

392 

393 form_data = request.scope.get("admin_form_data") or await request.form() 

394 csrf_token = str(form_data.get("csrf_token", "")) 

395 csrf_session_id = request.session.get("csrf_session_id") or request.session.get( 

396 "admin_user_id", "" 

397 ) 

398 if not csrf_session_id or not self._csrf_service.validate_token( 

399 csrf_session_id, csrf_token 

400 ): 

401 logger.warning( 

402 "auth.csrf_validation_failed", ip=self._get_client_ip(request) 

403 ) 

404 return RedirectResponse( 

405 url=f"/admin/profile/mfa?error={quote_plus('Invalid or expired security token. Please try again.')}", 

406 status_code=302, 

407 ) 

408 

409 code = str(form_data.get("code", "")) 

410 if not code: 

411 return RedirectResponse( 

412 url=f"/admin/profile/mfa?error={quote_plus('Verification code is required.')}", 

413 status_code=302, 

414 ) 

415 

416 result = await self._mfa_service.disable(str(user.user_id), code) 

417 if result.is_err(): 

418 return RedirectResponse( 

419 url=f"/admin/profile/mfa?error={quote_plus(_humanize_error(str(result.unwrap_err())))}", 

420 status_code=302, 

421 ) 

422 request.session.pop("mfa_pending_secret", None) 

423 return RedirectResponse( 

424 url="/admin/profile/mfa?notice=" 

425 + quote_plus("Two-factor authentication disabled."), 

426 status_code=302, 

427 )