Enhancement: skip local 2FA for trusted OIDC AMR (#14403)

src/paperless/tests/test_adapter.py

@@ -2,19 +2,26 @@ import logging
 
 import pytest
 from allauth.account.adapter import get_adapter
+from allauth.account.models import Login
 from allauth.core import context
+from allauth.mfa.totp.internal.auth import TOTP
+from allauth.mfa.totp.internal.auth import generate_totp_secret
 from allauth.socialaccount.adapter import get_adapter as get_social_adapter
+from allauth.socialaccount.models import SocialAccount
+from allauth.socialaccount.models import SocialLogin
 from django.contrib.auth.models import AnonymousUser
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import User
 from django.forms import ValidationError
 from django.http import HttpRequest
+from django.test import RequestFactory
 from django.urls import reverse
 from pytest_django.fixtures import Settings
 from pytest_mock import MockerFixture
 from rest_framework.authtoken.models import Token
 
 from paperless.adapter import DrfTokenStrategy
+from paperless.adapter import SocialAccountAuthenticateStage
 from paperless_testing.factories import UserFactory
 
 
@@ -129,6 +136,120 @@ class TestCustomAccountAdapter:
         assert not user2.is_superuser
 
 
+@pytest.mark.django_db
+class TestSocialAccountAuthenticateStage:
+    @staticmethod
+    def _requires_mfa(user: User, extra_data: dict | None) -> bool:
+        signal_kwargs = None
+        if extra_data is not None:
+            account = SocialAccount(
+                provider="openid_connect",
+                uid="uid",
+                extra_data=extra_data,
+            )
+            signal_kwargs = {"sociallogin": SocialLogin(user=user, account=account)}
+        request = RequestFactory().get("/")
+        request.session = {}
+        stage = SocialAccountAuthenticateStage(
+            None,
+            request,
+            Login(user=user, signal_kwargs=signal_kwargs),
+        )
+        response, cont = stage.handle()
+        assert cont is True
+        requires_mfa = response is not None
+        assert stage.state.get("authentication_required", False) is requires_mfa
+        if requires_mfa:
+            assert response.status_code == 302
+            assert response.url == reverse("mfa_authenticate")
+        return requires_mfa
+
+    def test_stage_replaces_allauth_stage(self) -> None:
+        stages = get_adapter().get_login_stages()
+
+        assert "paperless.adapter.SocialAccountAuthenticateStage" in stages
+        assert "allauth.mfa.stages.AuthenticateStage" not in stages
+
+    @pytest.mark.parametrize(
+        ("trusted_amr", "extra_data", "expected"),
+        [
+            pytest.param(
+                ["mfa"],
+                {"id_token": {"amr": ["pwd", "mfa"]}},
+                False,
+                id="trusted-amr",
+            ),
+            pytest.param(
+                ["phr"],
+                {"id_token": {"amr": ["otp"]}},
+                True,
+                id="untrusted-amr",
+            ),
+            pytest.param([], {"id_token": {"amr": ["mfa"]}}, True, id="setting-unset"),
+            pytest.param(["mfa"], {"id_token": {}}, True, id="no-amr-claim"),
+            pytest.param(
+                ["mfa"],
+                {"userinfo": {"amr": ["mfa"]}},
+                True,
+                id="userinfo-only",
+            ),
+            pytest.param(["m"], {"id_token": {"amr": "mfa"}}, True, id="string-amr"),
+            pytest.param(
+                ["mfa"],
+                {"id_token": {"amr": {"mfa": 1}}},
+                True,
+                id="dict-amr",
+            ),
+            pytest.param(
+                ["mfa"],
+                {"id_token": {"amr": [["mfa"]]}},
+                True,
+                id="nested-amr",
+            ),
+            pytest.param(
+                ["mfa"],
+                {"id_token": {"amr": ["mfa", 1]}},
+                True,
+                id="mixed-amr",
+            ),
+            pytest.param(["mfa"], {"id_token": {"amr": None}}, True, id="null-amr"),
+            pytest.param(["mfa"], {"id_token": {"amr": []}}, True, id="empty-amr"),
+            pytest.param(["mfa"], {"id_token": "mfa"}, True, id="string-id-token"),
+            pytest.param(["mfa"], {"id_token": ["mfa"]}, True, id="list-id-token"),
+            pytest.param(["mfa"], {"id_token": None}, True, id="null-id-token"),
+            pytest.param(["mfa"], None, True, id="regular-login"),
+        ],
+    )
+    def test_mfa_skipped_only_for_trusted_amr(
+        self,
+        settings: Settings,
+        trusted_amr: list[str],
+        extra_data: dict | None,
+        *,
+        expected: bool,
+    ) -> None:
+        """
+        GIVEN:
+            - A user with TOTP enabled
+        WHEN:
+            - The user logs in, via a social account or not
+        THEN:
+            - The 2FA prompt is only skipped if the ID token's amr claim contains
+              one of SOCIAL_ACCOUNT_MFA_TRUSTED_AMR
+        """
+        settings.SOCIAL_ACCOUNT_MFA_TRUSTED_AMR = trusted_amr
+        user = UserFactory(username="testuser")
+        TOTP.activate(user, generate_totp_secret())
+
+        assert self._requires_mfa(user, extra_data) is expected
+
+    def test_no_mfa_for_user_without_authenticator(self, settings: Settings) -> None:
+        settings.SOCIAL_ACCOUNT_MFA_TRUSTED_AMR = ["mfa"]
+        user = UserFactory(username="testuser")
+
+        assert not self._requires_mfa(user, {"id_token": {"amr": ["pwd"]}})
+
+
 class TestCustomSocialAccountAdapter:
     @pytest.mark.django_db
     def test_is_open_for_signup(self, settings: Settings) -> None: