Repository: jmwri/simplejwt

Issue:
Create a Jwt object
It would be useful to have a `Jwt` object that provides helpers, such as `expired()` etc.

Discussion:

Patch:
diff --git a/README.md b/README.md
index 2b7b764..e8506f6 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,12 @@
 
 A dead simple JWT library.
 
+# Supported algorithms
+
+* HS256
+* HS384
+* HS512
+
 # Usage
 ## Encode
 Returns a new token.
diff --git a/setup.py b/setup.py
index 18699b4..d39a525 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 github = 'https://github.com/jmwri/simplejwt'
-version = '0.3.0'
+version = '0.3.1'
 
 setup(
     name='simplejwt',
diff --git a/simplejwt/jwt.py b/simplejwt/jwt.py
index c776ec4..4ff7097 100644
--- a/simplejwt/jwt.py
+++ b/simplejwt/jwt.py
@@ -2,6 +2,7 @@ from typing import Union
 import json
 import hmac
 import hashlib
+from datetime import datetime
 
 from simplejwt import util
 from simplejwt.exception import InvalidSignatureError
@@ -12,6 +13,18 @@ algorithms = {
     'HS512': hashlib.sha512,
 }
 
+default_alg = 'HS256'
+
+registered_claims = {
+    'issuer': 'iss',
+    'subject': 'sub',
+    'audience': 'aud',
+    'valid_to': 'exp',
+    'valid_from': 'nbf',
+    'issued_at': 'iat',
+    'id': 'jti',
+}
+
 
 def get_algorithm(alg: str):
     if alg not in algorithms:
@@ -26,35 +39,136 @@ def _hash(secret: bytes, data: bytes, alg: str) -> bytes:
         .digest()
 
 
-def make(secret: Union[str, bytes], payload: dict, alg='HS256',
-         issuer: str = None, subject: str = None, audience: str = None,
-         valid_to: int = None, valid_from: int = None, issued_at: int = None,
-         id: str = None):
-    new_payload = payload.copy()
-    if issuer and 'iss' not in new_payload:
-        new_payload['iss'] = issuer
-    if subject and 'sub' not in new_payload:
-        new_payload['sub'] = subject
-    if audience and 'aud' not in new_payload:
-        new_payload['aud'] = audience
-    if valid_to and 'exp' not in new_payload:
-        new_payload['exp'] = valid_to
-    if valid_from and 'nbf' not in new_payload:
-        new_payload['nbf'] = valid_from
-    if issued_at and 'iat' not in new_payload:
-        new_payload['iat'] = issued_at
-    if id and 'jti' not in new_payload:
-        new_payload['jti'] = id
-    return encode(secret, new_payload, alg)
-
-
-def encode(secret: Union[str, bytes], payload: dict, alg='HS256'):
+class Jwt:
+    def __init__(self, secret: Union[str, bytes], payload: dict = None,
+                 alg: str = default_alg, header: dict = None,
+                 issuer: str = None, subject: str = None, audience: str = None,
+                 valid_to: int = None, valid_from: int = None,
+                 issued_at: int = None, id: str = None):
+        self.secret = secret
+        self.payload = payload or {}
+        self.alg = alg
+        self.header = header or {}
+        self.registered_claims = {}
+        if issuer:
+            self.issuer = issuer
+        if subject:
+            self.subject = subject
+        if audience:
+            self.audience = audience
+        if valid_to:
+            self.valid_to = valid_to
+        if valid_from:
+            self.valid_from = valid_from
+        if issued_at:
+            self.issued_at = issued_at
+        if id:
+            self.id = id
+        self._pop_claims_from_payload()
+
+    @property
+    def issuer(self):
+        return self.registered_claims.get('iss')
+
+    @issuer.setter
+    def issuer(self, issuer: str):
+        self.registered_claims['iss'] = issuer
+
+    @property
+    def subject(self):
+        return self.registered_claims.get('sub')
+
+    @subject.setter
+    def subject(self, subject: str):
+        self.registered_claims['sub'] = subject
+
+    @property
+    def audience(self):
+        return self.registered_claims.get('aud')
+
+    @audience.setter
+    def audience(self, audience: str):
+        self.registered_claims['aud'] = audience
+
+    @property
+    def valid_to(self):
+        return self.registered_claims.get('exp')
+
+    @valid_to.setter
+    def valid_to(self, valid_to: int):
+        self.registered_claims['exp'] = valid_to
+
+    @property
+    def valid_from(self):
+        return self.registered_claims.get('nbf')
+
+    @valid_from.setter
+    def valid_from(self, valid_from: int):
+        self.registered_claims['nbf'] = valid_from
+
+    @property
+    def issued_at(self):
+        return self.registered_claims.get('iat')
+
+    @issued_at.setter
+    def issued_at(self, issued_at: int):
+        self.registered_claims['iat'] = issued_at
+
+    @property
+    def id(self):
+        return self.registered_claims.get('jti')
+
+    @id.setter
+    def id(self, id: str):
+        self.registered_claims['jti'] = id
+
+    def valid(self, time: int = None):
+        time = time or int(datetime.utcnow().timestamp())
+        if time < self.valid_from:
+            return False
+        if time > self.valid_to:
+            return False
+        return True
+
+    def _pop_claims_from_payload(self):
+        claims_in_payload = [k for k in self.payload.keys() if
+                             k in registered_claims.values()]
+        for name in claims_in_payload:
+            self.registered_claims[name] = self.payload.pop(name)
+
+    def encode(self):
+        payload = {}
+        payload.update(self.registered_claims)
+        payload.update(self.payload)
+        return encode(self.secret, payload, self.alg, self.header)
+
+    @staticmethod
+    def decode(secret: Union[str, bytes], token: Union[str, bytes],
+               alg: str = default_alg):
+        header, payload = _decode(secret, token, alg)
+        return Jwt(secret, payload, alg, header)
+
+
+def make(secret: Union[str, bytes], payload: dict, alg: str = default_alg,
+         **kwargs):
+    jwt = Jwt(secret, payload, alg, **kwargs)
+    return jwt.encode()
+
+
+def encode(secret: Union[str, bytes], payload: dict = None,
+           alg: str = default_alg, header: dict = None):
     secret = util.to_bytes(secret)
 
-    header = {
-        'type': 'JWT',
-        'alg': alg
-    }
+    payload = payload or {}
+    header = header or {}
+
+    if isinstance(header, dict):
+        header = header.copy()
+        header.update({
+            'type': 'JWT',
+            'alg': alg
+        })
+
     header_json = util.to_bytes(json.dumps(header))
     header_b64 = util.b64_encode(header_json)
     payload_json = util.to_bytes(json.dumps(payload))
@@ -68,14 +182,19 @@ def encode(secret: Union[str, bytes], payload: dict, alg='HS256'):
     return util.from_bytes(token)
 
 
-def decode(secret: Union[str, bytes], token: Union[str, bytes], alg='HS256'):
+def _decode(secret: Union[str, bytes], token: Union[str, bytes],
+            alg: str = default_alg):
     secret = util.to_bytes(secret)
     token = util.to_bytes(token)
     pre_signature, signature_segment = token.rsplit(b'.', 1)
-    payload_b64 = pre_signature.split(b'.', 1)[1]
+    header_b64, payload_b64 = pre_signature.split(b'.')
+    header_json = util.b64_decode(header_b64)
+    header = json.loads(util.from_bytes(header_json))
     payload_json = util.b64_decode(payload_b64)
     payload = json.loads(util.from_bytes(payload_json))
 
+    if not isinstance(header, dict):
+        raise RuntimeError('Invalid header: {}'.format(header))
     if not isinstance(payload, dict):
         raise RuntimeError('Invalid payload: {}'.format(payload))
 
@@ -84,4 +203,10 @@ def decode(secret: Union[str, bytes], token: Union[str, bytes], alg='HS256'):
 
     if not hmac.compare_digest(signature, calculated_signature):
         raise InvalidSignatureError('Invalid signature')
+    return header, payload
+
+
+def decode(secret: Union[str, bytes], token: Union[str, bytes],
+           alg: str = default_alg):
+    _, payload = _decode(secret, token, alg)
     return payload


Test patch:
diff --git a/tests/test_jwt.py b/tests/test_jwt.py
index cacda88..0f8665a 100644
--- a/tests/test_jwt.py
+++ b/tests/test_jwt.py
@@ -1,5 +1,6 @@
 import pytest
 import hashlib
+from datetime import datetime
 
 from simplejwt import jwt, util
 from simplejwt.exception import InvalidSignatureError
@@ -25,16 +26,6 @@ test_token_data = {
     }
 }
 
-registered_claims = {
-    'issuer': 'iss',
-    'subject': 'sub',
-    'audience': 'aud',
-    'valid_to': 'exp',
-    'valid_from': 'nbf',
-    'issued_at': 'iat',
-    'id': 'jti',
-}
-
 test_registered_claims = {
     'issuer': 'test_issuer',
     'subject': 'test_subject',
@@ -84,7 +75,7 @@ def test_encode():
 
 
 def test_make_claims():
-    for name, abb in registered_claims.items():
+    for name, abb in jwt.registered_claims.items():
         args = {
             'secret': test_token_data['secret'],
             'payload': test_token_data['payload'],
@@ -95,6 +86,69 @@ def test_make_claims():
         assert payload[abb] == test_registered_claims[name]
 
 
+def test_jwt_registered_claims_constructor():
+    for name, abb in jwt.registered_claims.items():
+        args = {
+            'secret': test_token_data['secret'],
+            'payload': test_token_data['payload'],
+            name: test_registered_claims[name]
+        }
+        obj = jwt.Jwt(**args)
+        assert getattr(obj, name) == test_registered_claims[name]
+
+
+def test_jwt_registered_claims():
+    for name, abb in jwt.registered_claims.items():
+        args = {
+            'secret': test_token_data['secret'],
+            'payload': test_token_data['payload'],
+        }
+        obj = jwt.Jwt(**args)
+        setattr(obj, name, test_registered_claims[name])
+        token = obj.encode()
+        payload = jwt.decode(test_token_data['secret'], token)
+        assert getattr(obj, name) == test_registered_claims[name]
+        assert payload[abb] == test_registered_claims[name]
+
+
+def test_jwt_precedence():
+    obj = jwt.Jwt('secret', {'iss': 'usr_defined_iss'}, issuer='my_iss')
+    assert obj.registered_claims['iss'] == 'usr_defined_iss'
+
+
+def test_jwt_decode():
+    for alg, token in test_tokens.items():
+        obj = jwt.Jwt.decode(
+            test_token_data['secret'],
+            token,
+            alg
+        )
+        assert obj.secret == test_token_data['secret']
+        assert obj.alg == alg
+        assert obj.header == {
+            'type': 'JWT',
+            'alg': alg
+        }
+        assert obj.payload == test_token_data['payload']
+
+
+def test_jwt_valid():
+    obj = jwt.Jwt('secret', {}, valid_from=2, valid_to=4)
+    assert not obj.valid(1)
+    assert obj.valid(2)
+    assert obj.valid(3)
+    assert obj.valid(4)
+    assert not obj.valid(5)
+
+
+def test_jwt_valid_current_time():
+    now = int(datetime.utcnow().timestamp())
+    obj = jwt.Jwt('secret', {}, valid_from=now, valid_to=now)
+    assert obj.valid()
+    obj = jwt.Jwt('secret', {}, valid_from=now+1, valid_to=now+1)
+    assert not obj.valid()
+
+
 def test_make_precedence():
     token = jwt.make(test_token_data['secret'], {'iss': 'usr_defined_iss'},
                      issuer='my_iss')
@@ -131,6 +185,12 @@ def test_decode_invalid_signature():
             )
 
 
+def test_decode_invalid_header():
+    token = jwt.encode(test_token_data['secret'], header='should be dict')
+    with pytest.raises(RuntimeError):
+        jwt.decode(test_token_data['secret'], token)
+
+
 def test_decode_invalid_payload():
     token = jwt.encode(test_token_data['secret'], 'should be dict')
     with pytest.raises(RuntimeError):