authentik.stages.authenticator_totp.stage

 1"""TOTP Setup stage"""
 2
 3from urllib.parse import quote
 4
 5from django.http import HttpRequest, HttpResponse
 6from django.http.request import QueryDict
 7from django.utils.translation import gettext_lazy as _
 8from rest_framework.fields import CharField
 9from rest_framework.serializers import ValidationError
10
11from authentik.flows.challenge import (
12    Challenge,
13    ChallengeResponse,
14    WithUserInfoChallenge,
15)
16from authentik.flows.stage import ChallengeStageView
17from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage, TOTPDevice
18from authentik.stages.authenticator_totp.settings import OTP_TOTP_ISSUER
19
20SESSION_TOTP_DEVICE = "totp_device"
21
22
23class AuthenticatorTOTPChallenge(WithUserInfoChallenge):
24    """TOTP Setup challenge"""
25
26    config_url = CharField()
27    component = CharField(default="ak-stage-authenticator-totp")
28
29
30class AuthenticatorTOTPChallengeResponse(ChallengeResponse):
31    """TOTP Challenge response, device is set by get_response_instance"""
32
33    device: TOTPDevice
34
35    code = CharField()
36    component = CharField(default="ak-stage-authenticator-totp")
37
38    def validate_code(self, code: str) -> str:
39        """Validate totp code"""
40        if not self.device:
41            raise ValidationError(_("Code does not match"))
42        if not self.device.verify_token(code):
43            self.device.confirmed = False
44            raise ValidationError(_("Code does not match"))
45        return code
46
47
48class AuthenticatorTOTPStageView(ChallengeStageView):
49    """OTP totp Setup stage"""
50
51    response_class = AuthenticatorTOTPChallengeResponse
52
53    def get_challenge(self, *args, **kwargs) -> Challenge:
54        device: TOTPDevice = self.request.session[SESSION_TOTP_DEVICE]
55        return AuthenticatorTOTPChallenge(
56            data={
57                "config_url": device.config_url.replace(
58                    OTP_TOTP_ISSUER, quote(self.request.brand.branding_title)
59                ),
60            }
61        )
62
63    def get_response_instance(self, data: QueryDict) -> ChallengeResponse:
64        response = super().get_response_instance(data)
65        response.device = self.request.session.get(SESSION_TOTP_DEVICE)
66        return response
67
68    def get(self, request: HttpRequest, *args, **kwargs) -> HttpResponse:
69        user = self.get_pending_user()
70        if not user.is_authenticated:
71            self.logger.debug("No pending user, continuing")
72            return self.executor.stage_ok()
73
74        stage: AuthenticatorTOTPStage = self.executor.current_stage
75
76        if SESSION_TOTP_DEVICE not in self.request.session:
77            device = TOTPDevice(
78                user=user, confirmed=False, digits=stage.digits, name="TOTP Authenticator"
79            )
80
81            self.request.session[SESSION_TOTP_DEVICE] = device
82        return super().get(request, *args, **kwargs)
83
84    def challenge_valid(self, response: ChallengeResponse) -> HttpResponse:
85        """TOTP Token is validated by challenge"""
86        device: TOTPDevice = self.request.session[SESSION_TOTP_DEVICE]
87        device.confirmed = True
88        device.save()
89        del self.request.session[SESSION_TOTP_DEVICE]
90        return self.executor.stage_ok()