diff --git a/ghostwriter/home/templatetags/custom_tags.py b/ghostwriter/home/templatetags/custom_tags.py index 12e179bde..618655a6d 100644 --- a/ghostwriter/home/templatetags/custom_tags.py +++ b/ghostwriter/home/templatetags/custom_tags.py @@ -8,6 +8,7 @@ # 3rd Party Libraries from bs4 import BeautifulSoup +from allauth_2fa.utils import user_has_valid_totp_device # Ghostwriter Libraries from ghostwriter.api.utils import verify_access, verify_finding_access, verify_user_is_privileged @@ -124,3 +125,9 @@ def can_create_finding(user): def is_privileged(user): """Check if the user has the permission to create a finding.""" return verify_user_is_privileged(user) + + +@register.filter +def has_2fa(user): + """Check if the user has a valid TOTP method configured.""" + return user_has_valid_totp_device(user) diff --git a/ghostwriter/home/tests/test_views.py b/ghostwriter/home/tests/test_views.py index 31aa2bc9c..913ab4287 100644 --- a/ghostwriter/home/tests/test_views.py +++ b/ghostwriter/home/tests/test_views.py @@ -11,6 +11,9 @@ from django.test.utils import override_settings from django.urls import reverse +# 3rd Party Libraries +from django_otp.plugins.otp_static.models import StaticToken + # Ghostwriter Libraries from ghostwriter.factories import ( GroupFactory, @@ -138,6 +141,12 @@ def test_tags(self): self.user.save() self.assertTrue(custom_tags.can_create_finding(self.user)) + self.assertFalse(custom_tags.has_2fa(self.user)) + self.user.totpdevice_set.create() + static_model = self.user.staticdevice_set.create() + static_model.token_set.create(token=StaticToken.random_token()) + self.assertTrue(custom_tags.has_2fa(self.user)) + class DashboardTests(TestCase): """Collection of tests for :view:`home.dashboard`."""