diff --git a/tests/test_utils.py b/tests/test_utils.py index 64762e98e..910f04601 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -7,10 +7,13 @@ from phonenumber_field.phonenumber import PhoneNumber from two_factor.plugins.email.utils import mask_email +from two_factor.plugins.phonenumber.method import PhoneCallMethod, SMSMethod from two_factor.plugins.phonenumber.models import PhoneDevice from two_factor.plugins.phonenumber.utils import ( - backup_phones, format_phone_number, mask_phone_number, + backup_phones, format_phone_number, get_available_phone_methods, + mask_phone_number, ) +from two_factor.plugins.registry import GeneratorMethod, MethodRegistry from two_factor.utils import ( USER_DEFAULT_DEVICE_ATTR_NAME, default_device, get_otpauth_url, totp_digits, @@ -137,11 +140,28 @@ def test_wrong_device_hash(self): class PhoneUtilsTests(UserMixin, TestCase): + def test_get_available_phone_methods(self): + parameters = [ + # registered_methods, expected_codes + ([GeneratorMethod()], set()), + ([GeneratorMethod(), PhoneCallMethod()], {'call'}), + ([GeneratorMethod(), PhoneCallMethod(), SMSMethod()], {'call', 'sms'}), + ] + with mock.patch('two_factor.plugins.phonenumber.utils.registry', new_callable=MethodRegistry) as test_registry: + for registered_methods, expected_codes in parameters: + with self.subTest( + registered_methods=registered_methods, + expected_codes=expected_codes, + ): + test_registry._methods = registered_methods + codes = {method.code for method in get_available_phone_methods()} + self.assertEqual(codes, expected_codes) + def test_backup_phones(self): gateway = 'two_factor.gateways.fake.Fake' user = self.create_user() - user.phonedevice_set.create(name='default', number='+12024561111') - backup = user.phonedevice_set.create(name='backup', number='+12024561111') + user.phonedevice_set.create(name='default', number='+12024561111', method='call') + backup = user.phonedevice_set.create(name='backup', number='+12024561111', method='call') parameters = [ # with_gateway, with_user, expected_output diff --git a/tests/test_views_phone.py b/tests/test_views_phone.py index 6b527d73a..fdf4486ad 100644 --- a/tests/test_views_phone.py +++ b/tests/test_views_phone.py @@ -161,11 +161,13 @@ def setUp(self): self.default = self.user.phonedevice_set.create(name='default', method='call', number='+12024561111') self.login_user() + @override_settings(TWO_FACTOR_SMS_GATEWAY='two_factor.gateways.fake.Fake') def test_delete(self): + self.assertEqual(len(backup_phones(self.user)), 1) response = self.client.post(reverse('two_factor:phone_delete', args=[self.backup.pk])) self.assertRedirects(response, resolve_url(settings.LOGIN_REDIRECT_URL)) - self.assertEqual(list(backup_phones(self.user)), []) + self.assertEqual(backup_phones(self.user), []) def test_cannot_delete_default(self): response = self.client.post(reverse('two_factor:phone_delete', diff --git a/tests/test_views_profile.py b/tests/test_views_profile.py index fc1787a8c..c8d664bcc 100644 --- a/tests/test_views_profile.py +++ b/tests/test_views_profile.py @@ -2,60 +2,42 @@ from django.test import TestCase, override_settings from django.urls import reverse +from two_factor.plugins.registry import registry + from .utils import UserMixin +@override_settings( + TWO_FACTOR_SMS_GATEWAY='two_factor.gateways.fake.Fake', + TWO_FACTOR_CALL_GATEWAY='two_factor.gateways.fake.Fake', +) class ProfileTest(UserMixin, TestCase): - PHONENUMBER_PLUGIN_NAME = 'two_factor.plugins.phonenumber' - EXPECTED_BASE_CONTEXT_KEYS = { - 'default_device', - 'default_device_type', - 'backup_tokens', - } - EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS = { - 'backup_phones', - 'available_phone_methods', - } - def setUp(self): super().setUp() self.user = self.create_user() self.enable_otp() self.login_user() - @classmethod - def get_installed_apps_list(cls, with_phone_number_plugin=True): - apps = set(settings.INSTALLED_APPS) - if with_phone_number_plugin: - apps.add(cls.PHONENUMBER_PLUGIN_NAME) - else: - apps.remove(cls.PHONENUMBER_PLUGIN_NAME) - return list(apps) - def get_profile(self): url = reverse('two_factor:profile') return self.client.get(url) - def test_get_profile_without_phonenumer_plugin_enabled(self): - apps_list = self.get_installed_apps_list(with_phone_number_plugin=False) - with override_settings(INSTALLED_APPS=apps_list): + def test_get_profile_without_phonenumber_plugin_enabled(self): + without_phonenumber_plugin = [ + app for app in settings.INSTALLED_APPS if app != 'two_factor.plugins.phonenumber'] + + with override_settings(INSTALLED_APPS=without_phonenumber_plugin): + self.assertFalse(registry.get_method('call')) + self.assertFalse(registry.get_method('sms')) + response = self.get_profile() - context_keys = set(response.context.keys()) - self.assertTrue(self.EXPECTED_BASE_CONTEXT_KEYS.issubset(context_keys)) - # None of the phonenumber related keys are present - self.assertTrue( - self.EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS.isdisjoint( - context_keys - ) - ) + + self.assertTrue(response.context['available_phone_methods'] == []) def test_get_profile_with_phonenumer_plugin_enabled(self): - apps_list = self.get_installed_apps_list(with_phone_number_plugin=True) - with override_settings(INSTALLED_APPS=apps_list): - response = self.get_profile() - context_keys = set(response.context.keys()) - expected_keys = ( - self.EXPECTED_BASE_CONTEXT_KEYS - | self.EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS - ) - self.assertTrue(expected_keys.issubset(context_keys)) + self.assertTrue(registry.get_method('call')) + self.assertTrue(registry.get_method('sms')) + + response = self.get_profile() + available_phone_method_codes = {method.code for method in response.context['available_phone_methods']} + self.assertTrue(available_phone_method_codes == {'call', 'sms'}) diff --git a/tests/test_views_setup.py b/tests/test_views_setup.py index 497fe371b..b23fbe51e 100644 --- a/tests/test_views_setup.py +++ b/tests/test_views_setup.py @@ -2,8 +2,9 @@ from binascii import unhexlify from unittest import mock -from django.test import TestCase +from django.conf import settings from django.test.utils import override_settings +from django.test import TestCase from django.urls import reverse from django_otp import DEVICE_ID_SESSION_KEY from django_otp.oath import totp diff --git a/two_factor/plugins/phonenumber/apps.py b/two_factor/plugins/phonenumber/apps.py index 9992ad31d..c7b2aa9e3 100644 --- a/two_factor/plugins/phonenumber/apps.py +++ b/two_factor/plugins/phonenumber/apps.py @@ -12,11 +12,11 @@ class TwoFactorPhoneNumberConfig(AppConfig): url_prefix = 'phone' def ready(self): - updated_registered_methods(self, None, None) - setting_changed.connect(updated_registered_methods) + update_registered_methods(self, None, None) + setting_changed.connect(update_registered_methods) -def updated_registered_methods(sender, setting, value, **kwargs): +def update_registered_methods(sender, setting, value, **kwargs): # This allows for dynamic registration, typically when testing. from .method import PhoneCallMethod, SMSMethod diff --git a/two_factor/plugins/phonenumber/forms.py b/two_factor/plugins/phonenumber/forms.py index 3ce780fc9..9162f9451 100644 --- a/two_factor/plugins/phonenumber/forms.py +++ b/two_factor/plugins/phonenumber/forms.py @@ -15,9 +15,16 @@ class Meta: model = PhoneDevice fields = ['number', 'method'] + @staticmethod + def get_available_choices(): + choices = [] + for method in get_available_phone_methods(): + choices.append((method.code, method.verbose_name)) + return choices + def __init__(self, **kwargs): super().__init__(**kwargs) - self.fields['method'].choices = get_available_phone_methods() + self.fields['method'].choices = self.get_available_choices() class PhoneNumberForm(forms.ModelForm): diff --git a/two_factor/plugins/phonenumber/utils.py b/two_factor/plugins/phonenumber/utils.py index 9a8ea164e..6b9f4c8c7 100644 --- a/two_factor/plugins/phonenumber/utils.py +++ b/two_factor/plugins/phonenumber/utils.py @@ -1,33 +1,32 @@ import re import phonenumbers -from django.conf import settings -from django.utils.translation import gettext_lazy as _ -phone_mask = re.compile(r'(?<=.{3})[0-9](?=.{2})') - - -def backup_phones(user): - no_gateways = ( - getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None) is None - and getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None) is None) - no_user = not user or user.is_anonymous +from two_factor.plugins.registry import registry - if no_gateways or no_user: - from .models import PhoneDevice - return PhoneDevice.objects.none() - return user.phonedevice_set.filter(name='backup') +phone_mask = re.compile(r'(?<=.{3})[0-9](?=.{2})') def get_available_phone_methods(): methods = [] - if getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None): - methods.append(('call', _('Phone call'))) - if getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None): - methods.append(('sms', _('Text message'))) + for code in ['sms', 'call']: + if method := registry.get_method(code): + methods.append(method) + return methods +def backup_phones(user): + if not user or user.is_anonymous: + return [] + + phones = [] + for method in get_available_phone_methods(): + phones += list(method.get_devices(user)) + + return [phone for phone in phones if phone.name == 'backup'] + + def mask_phone_number(number): """ Masks a phone number, only first 3 and last 2 digits visible. diff --git a/two_factor/views/profile.py b/two_factor/views/profile.py index 1cd5cba8f..e9c45b21a 100644 --- a/two_factor/views/profile.py +++ b/two_factor/views/profile.py @@ -1,4 +1,3 @@ -from django.apps.registry import apps from django.conf import settings from django.contrib.auth.decorators import login_required from django.shortcuts import redirect, resolve_url @@ -30,24 +29,22 @@ class ProfileView(TemplateView): template_name = 'two_factor/profile/profile.html' def get_context_data(self, **kwargs): + user = self.request.user + try: - backup_tokens = self.request.user.staticdevice_set.all()[0].token_set.count() + backup_tokens = user.staticdevice_set.all()[0].token_set.count() except Exception: backup_tokens = 0 context = { - 'default_device': default_device(self.request.user), - 'default_device_type': default_device(self.request.user).__class__.__name__, + 'default_device': default_device(user), + 'default_device_type': default_device(user).__class__.__name__, 'backup_tokens': backup_tokens, + 'backup_phones': backup_phones(user), + 'available_phone_methods': get_available_phone_methods(), } - if (apps.is_installed('two_factor.plugins.phonenumber')): - context.update({ - 'backup_phones': backup_phones(self.request.user), - 'available_phone_methods': get_available_phone_methods(), - }) - return context