Proper use of jwcryptoHEADmaster

1 files changed, 26 insertions, 42 deletions

diff --git a/authserver.py b/authserver.py
index 6197082..8932066 100755
--- a/authserver.py
+++ b/authserver.py
@@ -21,8 +21,10 @@ import json
 import struct
 import pprint
 
+from jwcrypto.common import base64url_encode
 from jwcrypto.jwk import JWK
-from jwcrypto.jws import JWS, JWSCore
+from jwcrypto.jws import JWS
+from jwcrypto.jwt import JWT
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.x509 import load_pem_x509_certificate
@@ -39,19 +41,6 @@ def n2b64(n, endian='big'):
     b = n.to_bytes((n.bit_length() + 7) // 8, endian)
     return base64.urlsafe_b64encode(b).rstrip(b'=').decode('ascii')
 
-def dumps(obj):
-    """JSON dumps without white spaces as bytes
-    """
-    return json.dumps(
-        obj,
-        separators=(',', ':'),
-        sort_keys=True).encode('utf-8')
-
-def b64(obj):
-    """base64 encoding without =
-    """
-    return base64.urlsafe_b64encode(obj).rstrip(b'=')
-
 def private_key(filename, kid, password=None):
     """Get JWK from private key PEM file
     """
@@ -102,39 +91,37 @@ def get_fingerprint(pemfile):
     return ''.join(result)
 
 
-def bearer_token(claim):
+def bearer_token(claims):
     kid = get_fingerprint('ssl/server.pem')
     jwk = private_key('ssl/server.key', kid)
-    alg = 'RS256'
-    jose = {
+
+    header = {
         "typ": "JWT",
-        "alg": alg,
+        "alg": 'RS256',
         "kid": kid
     }
 
-    hdr = b64(dumps(jose))
-    clm = b64(claim.serialize())
-    payload = b'.'.join((hdr, clm))
-
-    # raw signature of custom data
-    jws = JWSCore(alg, jwk, None, '')
-    signature = jws.engine.sign(jws.key, payload)
-
-    return b'.'.join((hdr, clm, b64(signature)))
+    token = JWT(header, claims)
+    token.make_signed_token(key)
+    return token.serialize()
 
 
 class ClaimSet:
     def __init__(self, issuer, subject, audience):
-        self.iss = issuer
-        self.sub = subject
-        self.aud = audience
-        self.jti = b64(os.urandom(15)).decode('ascii')
-        self.nbf = self.iat = int(time.time())
-        self.exp = self.nbf + 600
-        self.access = []
+        now = int(time.time())
+        self.claims = {
+            'access': []
+            'aud': audience,
+            'exp': now + 600
+            'jti': base64url_encode(os.urandom(15))
+            'iat': now,
+            'iss': issuer,
+            'nbf': now,
+            'sub': subject,
+        }
 
     def add_access(self, typ, name, actions):
-        self.access.append(dict(
+        self.claims['access'].append(dict(
             type=typ,
             name=name,
             actions=actions,
@@ -142,10 +129,7 @@ class ClaimSet:
 
     def __repr__(self):
         return "<%s %r>" % (type(self).__name__,
-                            self.serialize())
-
-    def serialize(self):
-        return dumps(self.__dict__)
+                            json.dumps(self.claims)
 
 
 class DockerAuthHandler(BaseHTTPRequestHandler):
@@ -158,13 +142,13 @@ class DockerAuthHandler(BaseHTTPRequestHandler):
         subject = qs.get('account', [''])[0]
         audience = qs.get('service', [''])[0]
         scopes = qs.get('scope', [])
-        claim = ClaimSet(ISSUER, subject, audience)
+        claimset = ClaimSet(ISSUER, subject, audience)
         for scope in scopes:
             typ, name, actions = scope.split(':')
             actions = list(a.strip() for a in actions.split(','))
-            claim.add_access(typ, name, actions)
+            claimset.add_access(typ, name, actions)
         print(qs, claim)
-        token = bearer_token(claim)
+        token = bearer_token(claimset.claims)
         token = b'{"token":"' + token + b'"}'
 
         self.send_response(200)