From eb682f23ff9528a26dc5d424d166e068ef67796e Mon Sep 17 00:00:00 2001
From: Tuomas Suutari <tuomas.suutari@anders.fi>
Date: Fri, 7 Jul 2017 22:55:18 +0300
Subject: [PATCH] Pass scope to OIDC_IDTOKEN_PROCESSING_HOOK

The ID token processing hook might want to add claims to the ID token
conditionally based on the scope parameter.  Therefore it would be very
useful to provide the scope parameter to the processing hook.
---
 oidc_provider/lib/utils/common.py          |  5 ++++-
 oidc_provider/lib/utils/token.py           | 13 +++++++------
 oidc_provider/tests/app/utils.py           | 12 ++++++++++--
 oidc_provider/tests/test_token_endpoint.py | 20 ++++++++++++++++++++
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/oidc_provider/lib/utils/common.py b/oidc_provider/lib/utils/common.py
index c4778bd..9ef225d 100644
--- a/oidc_provider/lib/utils/common.py
+++ b/oidc_provider/lib/utils/common.py
@@ -114,7 +114,7 @@ def default_after_end_session_hook(request, id_token=None, post_logout_redirect_
     return None
 
 
-def default_idtoken_processing_hook(id_token, user):
+def default_idtoken_processing_hook(id_token, user, scope=None):
     """
     Hook to perform some additional actions ti `id_token` dictionary just before serialization.
 
@@ -124,6 +124,9 @@ def default_idtoken_processing_hook(id_token, user):
     :param user: user for whom id_token is generated
     :type user: User
 
+    :param scope: scope for the token
+    :type scope: list[str]|None
+
     :return: custom modified dictionary of values for `id_token`
     :rtype dict
     """
diff --git a/oidc_provider/lib/utils/token.py b/oidc_provider/lib/utils/token.py
index 3f0d2e9..91bf459 100644
--- a/oidc_provider/lib/utils/token.py
+++ b/oidc_provider/lib/utils/token.py
@@ -53,13 +53,14 @@ def create_id_token(user, aud, nonce='', at_hash='', request=None, scope=[]):
     if ('email' in scope) and getattr(user, 'email', None):
         dic['email'] = user.email
 
-    processing_hook = settings.get('OIDC_IDTOKEN_PROCESSING_HOOK')
+    processing_hooks = settings.get('OIDC_IDTOKEN_PROCESSING_HOOK')
 
-    if isinstance(processing_hook, (list, tuple)):
-        for hook in processing_hook:
-            dic = settings.import_from_str(hook)(dic, user=user)
-    else:
-        dic = settings.import_from_str(processing_hook)(dic, user=user)
+    if not isinstance(processing_hooks, (list, tuple)):
+        processing_hooks = [processing_hooks]
+
+    for hook_string in processing_hooks:
+        hook = settings.import_from_str(hook_string)
+        dic = hook(dic, user=user, scope=scope)
 
     return dic
 
diff --git a/oidc_provider/tests/app/utils.py b/oidc_provider/tests/app/utils.py
index 4f824ec..31a9aca 100644
--- a/oidc_provider/tests/app/utils.py
+++ b/oidc_provider/tests/app/utils.py
@@ -107,7 +107,7 @@ def fake_sub_generator(user):
     return user.email
 
 
-def fake_idtoken_processing_hook(id_token, user):
+def fake_idtoken_processing_hook(id_token, user, scope=None):
     """
     Fake function for inserting some keys into token. Testing OIDC_IDTOKEN_PROCESSING_HOOK.
     """
@@ -116,10 +116,18 @@ def fake_idtoken_processing_hook(id_token, user):
     return id_token
 
 
-def fake_idtoken_processing_hook2(id_token, user):
+def fake_idtoken_processing_hook2(id_token, user, scope=None):
     """
     Fake function for inserting some keys into token. Testing OIDC_IDTOKEN_PROCESSING_HOOK - tuple or list as param
     """
     id_token['test_idtoken_processing_hook2'] = FAKE_RANDOM_STRING
     id_token['test_idtoken_processing_hook_user_email2'] = user.email
     return id_token
+
+
+def fake_idtoken_processing_hook3(id_token, user, scope=None):
+    """
+    Fake function for checking scope is passed to processing hook.
+    """
+    id_token['scope_passed_to_processing_hook'] = scope
+    return id_token
diff --git a/oidc_provider/tests/test_token_endpoint.py b/oidc_provider/tests/test_token_endpoint.py
index 0760555..3fedd83 100644
--- a/oidc_provider/tests/test_token_endpoint.py
+++ b/oidc_provider/tests/test_token_endpoint.py
@@ -734,6 +734,26 @@ class TokenTestCase(TestCase):
         self.assertEqual(id_token.get('test_idtoken_processing_hook2'), FAKE_RANDOM_STRING)
         self.assertEqual(id_token.get('test_idtoken_processing_hook_user_email2'), self.user.email)
 
+    @override_settings(
+        OIDC_IDTOKEN_PROCESSING_HOOK=(
+                'oidc_provider.tests.app.utils.fake_idtoken_processing_hook3'))
+    def test_additional_idtoken_processing_hook_scope_param(self):
+        """
+        Test scope parameter is passed to OIDC_IDTOKEN_PROCESSING_HOOK.
+        """
+        code = self._create_code(['openid', 'email', 'profile', 'dummy'])
+
+        post_data = self._auth_code_post_data(code=code.code)
+
+        response = self._post_request(post_data)
+
+        response_dic = json.loads(response.content.decode('utf-8'))
+        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()
+
+        self.assertEqual(
+            id_token.get('scope_passed_to_processing_hook'),
+            ['openid', 'email', 'profile', 'dummy'])
+
     def test_pkce_parameters(self):
         """
         Test Proof Key for Code Exchange by OAuth Public Clients.