From a5d62a40d530bf78e513ee578f1cac20d83caaa0 Mon Sep 17 00:00:00 2001
From: Valentin Deniaud <vdeniaud@entrouvert.com>
Date: Tue, 20 Sep 2022 15:38:41 +0200
Subject: [PATCH 10/10] auth_oidc: allow adding roles on login (#53442)

---
 src/authentic2_auth_oidc/backends.py |  6 ++++++
 src/authentic2_auth_oidc/models.py   |  7 ++++++-
 tests/test_auth_oidc.py              | 22 ++++++++++++++++++++++
 tests/test_manager_authenticators.py | 10 ++++++++++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/src/authentic2_auth_oidc/backends.py b/src/authentic2_auth_oidc/backends.py
index 5d816ca22..784e17bb6 100644
--- a/src/authentic2_auth_oidc/backends.py
+++ b/src/authentic2_auth_oidc/backends.py
@@ -382,6 +382,12 @@ class OIDCBackend(ModelBackend):
                 setattr(user.verified_attributes, attribute, value)
             else:
                 setattr(user.attributes, attribute, value)
+
+        for action in provider.add_role_actions.all():
+            if action.role not in user.roles.all():
+                logger.info('auth_oidc: adding role "%s" to user %s', action.role, user)
+                user.roles.add(action.role)
+
         return user
 
     def get_saml2_authn_context(self):
diff --git a/src/authentic2_auth_oidc/models.py b/src/authentic2_auth_oidc/models.py
index 3bd265925..577430e68 100644
--- a/src/authentic2_auth_oidc/models.py
+++ b/src/authentic2_auth_oidc/models.py
@@ -26,7 +26,11 @@ from django.utils.translation import gettext_lazy as _
 from jwcrypto.jwk import InvalidJWKValue, JWKSet
 
 from authentic2.a2_rbac.utils import get_default_ou
-from authentic2.apps.authenticators.models import AuthenticatorRelatedObjectBase, BaseAuthenticator
+from authentic2.apps.authenticators.models import (
+    AddRoleAction,
+    AuthenticatorRelatedObjectBase,
+    BaseAuthenticator,
+)
 from authentic2.utils.misc import make_url, redirect_to_login
 from authentic2.utils.template import validate_template
 
@@ -136,6 +140,7 @@ class OIDCProvider(BaseAuthenticator):
     def related_models(self):
         return {
             OIDCClaimMapping: self.claim_mappings.all(),
+            AddRoleAction: self.add_role_actions.all(),
         }
 
     @property
diff --git a/tests/test_auth_oidc.py b/tests/test_auth_oidc.py
index fe7f2795b..714cc5a37 100644
--- a/tests/test_auth_oidc.py
+++ b/tests/test_auth_oidc.py
@@ -1315,3 +1315,25 @@ def test_only_idtoken_claims(app, caplog, code, oidc_provider, oidc_provider_jwk
         ):
             response = app.get(login_callback_url(oidc_provider), params={'code': code, 'state': state})
         assert User.objects.count() == 1
+
+
+def test_oidc_add_role(app, code, oidc_provider, oidc_provider_jwkset, simple_role):
+    oidc_provider.add_role_actions.create(role=simple_role)
+
+    response = app.get('/').maybe_follow()
+    response = response.click(oidc_provider.name)
+    location = urllib.parse.urlparse(response.location)
+    query = QueryDict(location.query)
+    state = query['state']
+    nonce = query['nonce']
+
+    with oidc_provider_mock(
+        oidc_provider,
+        oidc_provider_jwkset,
+        code,
+        nonce=nonce,
+    ):
+        response = app.get(login_callback_url(oidc_provider), params={'code': code, 'state': state})
+
+    user = User.objects.get()
+    assert simple_role in user.roles.all()
diff --git a/tests/test_manager_authenticators.py b/tests/test_manager_authenticators.py
index 6c2dff3b6..40e4c90ff 100644
--- a/tests/test_manager_authenticators.py
+++ b/tests/test_manager_authenticators.py
@@ -223,6 +223,16 @@ def test_authenticators_oidc_claims(app, superuser):
     assert_event('authenticator.related_object.deletion', user=superuser, session=app.session)
 
 
+def test_authenticators_oidc_add_role(app, superuser, role_ou1):
+    authenticator = OIDCProvider.objects.create(slug='idp1')
+    resp = login(app, superuser, path=authenticator.get_absolute_url())
+
+    resp = resp.click('Add', href='role')
+    resp.form['role'] = role_ou1.pk
+    resp = resp.form.submit().follow()
+    assert 'role_ou1' in resp.text
+
+
 def test_authenticators_fc(app, superuser):
     resp = login(app, superuser, path='/manage/authenticators/')
 
-- 
2.30.2