Merge pull request #34 from vanzi/master

Added support for custom headers in encode method

3 files changed, 18 insertions, 3 deletions

diff --git a/README.md b/README.md
index c612b8c..3756ecc 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,10 @@ Usage
     import jwt
     jwt.encode({"some": "payload"}, "secret")
 
+Additional headers may also be specified.
+
+    jwt.encode({"some": "payload"}, "secret", headers={"kid": "230498151c214b788dd97f22b85410a5"})
+
 Note the resulting JWT will not be encrypted, but verifiable with a secret key.
 
     jwt.decode("someJWTstring", "secret")
diff --git a/jwt/__init__.py b/jwt/__init__.py
index 07e77fe..fd38bd1 100644
--- a/jwt/__init__.py
+++ b/jwt/__init__.py
@@ -67,13 +67,13 @@ try:
     from Crypto.Hash import SHA384
     from Crypto.Hash import SHA512
     from Crypto.PublicKey import RSA
-    
+
     signing_methods.update({
         'RS256': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA256.new(msg)),
         'RS384': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA384.new(msg)),
         'RS512': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA512.new(msg))
     })
-    
+
     verify_methods.update({
         'RS256': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA256.new(msg), sig),
         'RS384': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA384.new(msg), sig),
@@ -139,7 +139,7 @@ def header(jwt):
         raise DecodeError("Invalid header encoding")
 
 
-def encode(payload, key, algorithm='HS256'):
+def encode(payload, key, algorithm='HS256', headers=None):
     segments = []
 
     # Check that we get a mapping
@@ -149,6 +149,8 @@ def encode(payload, key, algorithm='HS256'):
 
     # Header
     header = {"typ": "JWT", "alg": algorithm}
+    if headers:
+        header.update(headers)
     json_header = json.dumps(header, separators=(',', ':')).encode('utf-8')
     segments.append(base64url_encode(json_header))
 
diff --git a/tests/test_jwt.py b/tests/test_jwt.py
index 5dd7374..8d93440 100644
--- a/tests/test_jwt.py
+++ b/tests/test_jwt.py
@@ -123,6 +123,15 @@ class TestJWT(unittest.TestCase):
             lambda: jwt.verify_signature(decoded_payload, signing,
                                          header, signature))
 
+    def test_custom_headers(self):
+        right_secret = 'foo'
+        headers = {'foo': 'bar', 'kid': 'test'}
+        jwt_message = jwt.encode(self.payload, right_secret, headers=headers)
+        decoded_payload, signing, header, signature = jwt.load(jwt_message)
+
+        for key, value in headers.items():
+            self.assertEqual(header[key], value)
+
     def test_invalid_crypto_alg(self):
         self.assertRaises(NotImplementedError, jwt.encode, self.payload,
                           "secret", "HS1024")