"""test flow with otp stages""" from base64 import b32decode from sys import platform from time import sleep from unittest.case import skipUnless from urllib.parse import parse_qs, urlparse from django_otp.oath import TOTP from django_otp.plugins.otp_static.models import StaticDevice, StaticToken from django_otp.plugins.otp_totp.models import TOTPDevice from selenium.webdriver.common.by import By from selenium.webdriver.common.keys import Keys from selenium.webdriver.support import expected_conditions as ec from selenium.webdriver.support.wait import WebDriverWait from authentik.blueprints.tests import apply_blueprint from authentik.flows.models import Flow from authentik.stages.authenticator_static.models import AuthenticatorStaticStage from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage from tests.e2e.utils import SeleniumTestCase, retry @skipUnless(platform.startswith("linux"), "requires local docker") class TestFlowsAuthenticator(SeleniumTestCase): """test flow with otp stages""" @retry() @apply_blueprint( "blueprints/default/10-flow-default-authentication-flow.yaml", "blueprints/default/10-flow-default-invalidation-flow.yaml", ) def test_totp_validate(self): """test flow with otp stages""" # Setup TOTP Device device = TOTPDevice.objects.create(user=self.user, confirmed=True, digits=6) flow: Flow = Flow.objects.get(slug="default-authentication-flow") self.driver.get(self.url("authentik_core:if-flow", flow_slug=flow.slug)) self.login() # Get expected token totp = TOTP(device.bin_key, device.step, device.t0, device.digits, device.drift) flow_executor = self.get_shadow_root("ak-flow-executor") validation_stage = self.get_shadow_root("ak-stage-authenticator-validate", flow_executor) code_stage = self.get_shadow_root("ak-stage-authenticator-validate-code", validation_stage) code_stage.find_element(By.CSS_SELECTOR, "input[name=code]").send_keys(totp.token()) code_stage.find_element(By.CSS_SELECTOR, "input[name=code]").send_keys(Keys.ENTER) self.wait_for_url(self.if_user_url("/library")) self.assert_user(self.user) @retry() @apply_blueprint( "blueprints/default/10-flow-default-authentication-flow.yaml", "blueprints/default/10-flow-default-invalidation-flow.yaml", ) @apply_blueprint("default/20-flow-default-authenticator-totp-setup.yaml") def test_totp_setup(self): """test TOTP Setup stage""" flow: Flow = Flow.objects.get(slug="default-authentication-flow") self.driver.get(self.url("authentik_core:if-flow", flow_slug=flow.slug)) self.login() self.wait_for_url(self.if_user_url("/library")) self.assert_user(self.user) self.driver.get( self.url( "authentik_flows:configure", stage_uuid=AuthenticatorTOTPStage.objects.first().stage_uuid, ) ) flow_executor = self.get_shadow_root("ak-flow-executor") totp_stage = self.get_shadow_root("ak-stage-authenticator-totp", flow_executor) wait = WebDriverWait(totp_stage, self.wait_timeout) wait.until(ec.presence_of_element_located((By.CSS_SELECTOR, "input[name=otp_uri]"))) otp_uri = totp_stage.find_element(By.CSS_SELECTOR, "input[name=otp_uri]").get_attribute( "value" ) # Parse the OTP URI, extract the secret and get the next token otp_args = urlparse(otp_uri) self.assertEqual(otp_args.scheme, "otpauth") otp_qs = parse_qs(otp_args.query) secret_key = b32decode(otp_qs["secret"][0]) totp = TOTP(secret_key) totp_stage.find_element(By.CSS_SELECTOR, "input[name=code]").send_keys(totp.token()) totp_stage.find_element(By.CSS_SELECTOR, "input[name=code]").send_keys(Keys.ENTER) sleep(3) self.assertTrue(TOTPDevice.objects.filter(user=self.user, confirmed=True).exists()) @retry() @apply_blueprint( "blueprints/default/10-flow-default-authentication-flow.yaml", "blueprints/default/10-flow-default-invalidation-flow.yaml", ) @apply_blueprint("default/20-flow-default-authenticator-static-setup.yaml") def test_static_setup(self): """test Static OTP Setup stage""" flow: Flow = Flow.objects.get(slug="default-authentication-flow") self.driver.get(self.url("authentik_core:if-flow", flow_slug=flow.slug)) self.login() self.wait_for_url(self.if_user_url("/library")) self.assert_user(self.user) self.driver.get( self.url( "authentik_flows:configure", stage_uuid=AuthenticatorStaticStage.objects.first().stage_uuid, ) ) # Remember the current URL as we should end up back here destination_url = self.driver.current_url flow_executor = self.get_shadow_root("ak-flow-executor") authenticator_stage = self.get_shadow_root("ak-stage-authenticator-static", flow_executor) token = authenticator_stage.find_element(By.CSS_SELECTOR, "ul li:nth-child(1)").text authenticator_stage.find_element(By.CSS_SELECTOR, "button[type=submit]").click() self.wait_for_url(destination_url) sleep(1) self.assertTrue(StaticDevice.objects.filter(user=self.user, confirmed=True).exists()) device = StaticDevice.objects.filter(user=self.user, confirmed=True).first() self.assertTrue(StaticToken.objects.filter(token=token, device=device).exists())