From dff76cd1ea1927fce40f4f3bb2367d3fcf2f7981 Mon Sep 17 00:00:00 2001
From: Ignacio Fiorentino <juanifioren@gmail.com>
Date: Tue, 22 Mar 2016 16:17:56 -0300
Subject: [PATCH] Add HS256 support for JWS.

---
 oidc_provider/lib/endpoints/authorize.py      |  2 +-
 oidc_provider/lib/endpoints/token.py          |  4 ++--
 oidc_provider/lib/utils/token.py              | 22 ++++++++++++-------
 .../migrations/0011_client_jwt_alg.py         | 20 +++++++++++++++++
 oidc_provider/models.py                       |  9 ++++++--
 5 files changed, 44 insertions(+), 13 deletions(-)
 create mode 100644 oidc_provider/migrations/0011_client_jwt_alg.py

diff --git a/oidc_provider/lib/endpoints/authorize.py b/oidc_provider/lib/endpoints/authorize.py
index 56972a4..f38c6ce 100644
--- a/oidc_provider/lib/endpoints/authorize.py
+++ b/oidc_provider/lib/endpoints/authorize.py
@@ -113,7 +113,7 @@ class AuthorizeEndpoint(object):
                         user=self.request.user,
                         aud=self.client.client_id,
                         nonce=self.params.nonce)
-                    query_fragment['id_token'] = encode_id_token(id_token_dic)
+                    query_fragment['id_token'] = encode_id_token(id_token_dic, self.client)
                 else:
                     id_token_dic = {}
 
diff --git a/oidc_provider/lib/endpoints/token.py b/oidc_provider/lib/endpoints/token.py
index a981eee..099947b 100644
--- a/oidc_provider/lib/endpoints/token.py
+++ b/oidc_provider/lib/endpoints/token.py
@@ -140,7 +140,7 @@ class TokenEndpoint(object):
             'refresh_token': token.refresh_token,
             'token_type': 'bearer',
             'expires_in': settings.get('OIDC_TOKEN_EXPIRE'),
-            'id_token': encode_id_token(id_token_dic),
+            'id_token': encode_id_token(id_token_dic, token.client),
         }
 
         return dic
@@ -173,7 +173,7 @@ class TokenEndpoint(object):
             'refresh_token': token.refresh_token,
             'token_type': 'bearer',
             'expires_in': settings.get('OIDC_TOKEN_EXPIRE'),
-            'id_token': encode_id_token(id_token_dic),
+            'id_token': encode_id_token(id_token_dic, self.token.client),
         }
 
         return dic
diff --git a/oidc_provider/lib/utils/token.py b/oidc_provider/lib/utils/token.py
index e512326..4c8ca93 100644
--- a/oidc_provider/lib/utils/token.py
+++ b/oidc_provider/lib/utils/token.py
@@ -6,6 +6,7 @@ from Crypto.PublicKey.RSA import importKey
 from django.utils import timezone
 from hashlib import md5
 from jwkest.jwk import RSAKey as jwk_RSAKey
+from jwkest.jwk import SYMKey
 from jwkest.jws import JWS
 
 from oidc_provider.lib.utils.common import get_issuer
@@ -55,21 +56,26 @@ def create_id_token(user, aud, nonce):
     return dic
 
 
-def encode_id_token(payload):
+def encode_id_token(payload, client):
     """
     Represent the ID Token as a JSON Web Token (JWT).
 
     Return a hash.
     """
-    keys = []
+    alg = client.jwt_alg
+    if alg == 'RS256':
+        keys = []
+        for rsakey in RSAKey.objects.all():
+            keys.append(jwk_RSAKey(key=importKey(rsakey.key), kid=rsakey.kid))
 
-    for rsakey in RSAKey.objects.all():
-        keys.append(jwk_RSAKey(key=importKey(rsakey.key), kid=rsakey.kid))
-
-    if not keys:
-        raise Exception('You must add at least one RSA Key.')
+        if not keys:
+            raise Exception('You must add at least one RSA Key.')
+    elif alg == 'HS256':
+        keys = [SYMKey(key=client.client_secret, alg=alg)]
+    else:
+        raise Exception('Unsupported key algorithm.')
     
-    _jws = JWS(payload, alg='RS256')
+    _jws = JWS(payload, alg=alg)
 
     return _jws.sign_compact(keys)
 
diff --git a/oidc_provider/migrations/0011_client_jwt_alg.py b/oidc_provider/migrations/0011_client_jwt_alg.py
new file mode 100644
index 0000000..d5e552b
--- /dev/null
+++ b/oidc_provider/migrations/0011_client_jwt_alg.py
@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.9 on 2016-03-22 17:42
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('oidc_provider', '0010_code_is_authentication'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='client',
+            name='jwt_alg',
+            field=models.CharField(choices=[(b'HS256', b'HS256'), (b'RS256', b'RS256')], default=b'RS256', max_length=10, verbose_name='JWT Algorithm'),
+        ),
+    ]
diff --git a/oidc_provider/models.py b/oidc_provider/models.py
index 68ace73..c2fe67f 100644
--- a/oidc_provider/models.py
+++ b/oidc_provider/models.py
@@ -16,11 +16,16 @@ class Client(models.Model):
         ('id_token token', 'id_token token (Implicit Flow)'),
     ]
 
+    JWT_ALGS = [
+        ('HS256', 'HS256'),
+        ('RS256', 'RS256'),
+    ]
+
     name = models.CharField(max_length=100, default='')
     client_id = models.CharField(max_length=255, unique=True)
     client_secret = models.CharField(max_length=255, unique=True)
-    response_type = models.CharField(max_length=30,
-                                     choices=RESPONSE_TYPE_CHOICES)
+    response_type = models.CharField(max_length=30, choices=RESPONSE_TYPE_CHOICES)
+    jwt_alg = models.CharField(max_length=10, choices=JWT_ALGS, default='RS256', verbose_name=_(u'JWT Algorithm'))
     date_created = models.DateField(auto_now_add=True)
 
     _redirect_uris = models.TextField(default='', verbose_name=_(u'Redirect URI'), help_text=_(u'Enter each URI on a new line.'))