rhodecode/apps/login/tests/test_2fa.py

import pytest
import mock

from rhodecode.lib.type_utils import AttributeDict
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'

    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.init_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.init_secret_2fa()
        recovery_code_to_check = user.init_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_code_to_check})
        assert response.status_code == 302
        response = self.app.post(route_path('check_2fa'), {'totp': recovery_code_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
        with mock.patch.object(
                SettingsModel, 'get_setting_by_name', lambda *a, **kw: AttributeDict(app_settings_value=True)):

            assert user.has_enabled_2fa