rhodecode/apps/login/tests/test_2fa.py

import pytest

from rhodecode.model.meta import Session
from rhodecode.tests.fixture import Fixture
from rhodecode.tests.routes import route_path
from rhodecode.model.settings import SettingsModel

fixture = Fixture()


@pytest.mark.usefixtures('app')
class Test2FA(object):
    @classmethod
    def setup_class(cls):
        cls.password = 'valid-one'

    @classmethod
    def teardown_class(cls):
        SettingsModel().create_or_update_setting('auth_rhodecode_global_2fa', False)

    def test_redirect_to_2fa_setup_if_enabled_for_user(self, user_util):
        user = user_util.create_user(password=self.password)
        user.has_enabled_2fa = True
        self.app.post(
            route_path('login'),
            {'username': user.username,
             'password': self.password})

        response = self.app.get('/')
        assert response.status_code == 302
        assert response.location.endswith(route_path('setup_2fa'))

    def test_redirect_to_2fa_check_if_2fa_configured(self, user_util):
        user = user_util.create_user(password=self.password)
        user.has_enabled_2fa = True
        user.secret_2fa
        Session().add(user)
        Session().commit()
        self.app.post(
            route_path('login'),
            {'username': user.username,
             'password': self.password})
        response = self.app.get('/')
        assert response.status_code == 302
        assert response.location.endswith(route_path('check_2fa'))

    def test_2fa_recovery_codes_works_only_once(self, user_util):
        user = user_util.create_user(password=self.password)
        user.has_enabled_2fa = True
        user.secret_2fa
        recovery_cod_to_check = user.get_2fa_recovery_codes()[0]
        Session().add(user)
        Session().commit()
        self.app.post(
            route_path('login'),
            {'username': user.username,
             'password': self.password})
        response = self.app.post(route_path('check_2fa'), {'totp': recovery_cod_to_check})
        assert response.status_code == 302
        response = self.app.post(route_path('check_2fa'), {'totp': recovery_cod_to_check})
        response.mustcontain('Code is invalid. Try again!')

    def test_2fa_state_when_forced_by_admin(self, user_util):
        user = user_util.create_user(password=self.password)
        user.has_enabled_2fa = False
        SettingsModel().create_or_update_setting('auth_rhodecode_global_2fa', True)
        assert user.has_enabled_2fa