From 7a357001b6d06fda6a8d6aa14ea213c573445848 Mon Sep 17 00:00:00 2001
From: Wojciech Bartosiak <wojtek@flipsports.com>
Date: Fri, 12 Feb 2016 16:02:35 +0000
Subject: [PATCH] Added OIDC_ID_TOKEN_PROCESSING_HOOK functionality

---
 .gitignore                                 |  2 ++
 oidc_provider/lib/utils/common.py          | 13 ++++++++++++-
 oidc_provider/lib/utils/token.py           |  2 ++
 oidc_provider/settings.py                  |  8 ++++++++
 oidc_provider/tests/app/utils.py           | 13 ++++++++++++-
 oidc_provider/tests/test_token_endpoint.py | 20 ++++++++++++++++----
 tox.ini                                    |  6 +++---
 7 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9046185..21e397b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,3 +7,5 @@ dist/
 .tox
 .coverage
 src/
+.venv
+.idea
diff --git a/oidc_provider/lib/utils/common.py b/oidc_provider/lib/utils/common.py
index 2a529d9..3743684 100644
--- a/oidc_provider/lib/utils/common.py
+++ b/oidc_provider/lib/utils/common.py
@@ -1,4 +1,3 @@
-from django.conf import settings as django_settings
 from django.core.urlresolvers import reverse
 from django.http import HttpResponse
 
@@ -49,3 +48,15 @@ def default_after_userlogin_hook(request, user, client):
     Default function for setting OIDC_AFTER_USERLOGIN_HOOK.
     """
     return None
+
+def additional_id_token_processing_hook(id_token):
+    """
+    Hook to perform some additional actions ti `id_token` dictionary just before serialization.
+
+    :param id_token: dictionary contains values that going to be serialized into `id_token`
+    :type id_token: dict
+    :return: custom modified dictionary of values for `id_token`
+    :rtype dict
+    """
+    return id_token
+
diff --git a/oidc_provider/lib/utils/token.py b/oidc_provider/lib/utils/token.py
index 9ad6f86..05eeaa6 100644
--- a/oidc_provider/lib/utils/token.py
+++ b/oidc_provider/lib/utils/token.py
@@ -44,6 +44,8 @@ def create_id_token(user, aud, nonce):
     if nonce:
         dic['nonce'] = str(nonce)
 
+    dic = settings.get('OIDC_ID_TOKEN_PROCESSING_HOOK', import_str=True)(dic)
+
     return dic
 
 
diff --git a/oidc_provider/settings.py b/oidc_provider/settings.py
index 9b5f9f7..0bae84b 100644
--- a/oidc_provider/settings.py
+++ b/oidc_provider/settings.py
@@ -97,6 +97,14 @@ class DefaultSettings(object):
         """
         return 'oidc_provider.lib.utils.common.DefaultUserInfo'
 
+    @property
+    def OIDC_ID_TOKEN_PROCESSING_HOOK(self):
+        """
+        OPTIONAL. A string with the location of your hook.
+        Used to add extra dictionary values specific for your app into id_token.
+        """
+        return 'oidc_provider.lib.utils.common.additional_id_token_processing_hook'
+
 default_settings = DefaultSettings()
 
 
diff --git a/oidc_provider/tests/app/utils.py b/oidc_provider/tests/app/utils.py
index 4c707ba..0b78eba 100644
--- a/oidc_provider/tests/app/utils.py
+++ b/oidc_provider/tests/app/utils.py
@@ -1,4 +1,6 @@
 import os
+import random
+import string
 try:
     from urlparse import parse_qs, urlsplit
 except ImportError:
@@ -9,7 +11,8 @@ from django.contrib.auth.models import User
 from oidc_provider.models import *
 
 
-FAKE_NONCE = 'cb584e44c43ed6bd0bc2d9c7e242837d'  
+FAKE_NONCE = 'cb584e44c43ed6bd0bc2d9c7e242837d'
+FAKE_RANDOM_STRING = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(32))
 
 
 def create_fake_user():
@@ -103,3 +106,11 @@ def fake_sub_generator(user):
     Fake function for setting OIDC_IDTOKEN_SUB_GENERATOR.
     """
     return user.email
+
+
+def fake_id_token_processing_hook(id_token):
+    """
+    Fake function for inserting some keys into token. Testing OIDC_ID_TOKEN_PROCESSING_HOOK
+    """
+    id_token['test_id_token_processing_hook'] = FAKE_RANDOM_STRING
+    return id_token
\ No newline at end of file
diff --git a/oidc_provider/tests/test_token_endpoint.py b/oidc_provider/tests/test_token_endpoint.py
index 804499c..8a5047d 100644
--- a/oidc_provider/tests/test_token_endpoint.py
+++ b/oidc_provider/tests/test_token_endpoint.py
@@ -1,16 +1,12 @@
 from base64 import b64encode
-import json
 try:
     from urllib.parse import urlencode
 except ImportError:
     from urllib import urlencode
-import uuid
 
-from django.core.urlresolvers import reverse
 from django.test import RequestFactory, override_settings
 from django.test import TestCase
 from jwkest.jwk import KEYS
-from jwkest.jws import JWS
 from jwkest.jwt import JWT
 from mock import patch
 
@@ -339,3 +335,19 @@ class TokenTestCase(TestCase):
         id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()
 
         self.assertEqual(id_token.get('sub'), self.user.email)
+
+    @override_settings(OIDC_ID_TOKEN_PROCESSING_HOOK='oidc_provider.tests.app.utils.fake_id_token_processing_hook')
+    def test_additional_id_token_processing_hook(self):
+        """
+        Test custom function for setting OIDC_ID_TOKEN_PROCESSING_HOOK.
+        """
+        code = self._create_code()
+
+        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('test_id_token_processing_hook'), FAKE_RANDOM_STRING)
diff --git a/tox.ini b/tox.ini
index fbd8037..92540ff 100644
--- a/tox.ini
+++ b/tox.ini
@@ -6,9 +6,9 @@ envlist=
 [testenv]
 
 deps =
-    django17: django==1.7
-    django18: django==1.8
-    django19: django==1.9
+    django17: django>=1.7,<1.8
+    django18: django>=1.8,<1.9
+    django19: django>=1.9,<2.0
     coverage
     mock