feat: Add draft 02 support for python

https://tools.ietf.org/html/draft-ietf-webpush-vapid-02

fixes: #25

Author: jrconlin

python/py_vapid/__init__.py

@@ -11,17 +11,25 @@
 import ecdsa
 from jose import jws
 
+# Show compliance version. For earlier versions see previously tagged releases.
+VERSION = "VAPID-DRAFT-02/ECE-DRAFT-07"
+
 
 class VapidException(Exception):
     """An exception wrapper for Vapid."""
     pass
 
 
-class Vapid(object):
-    """Minimal VAPID signature generation library. """
+class Vapid01(object):
+    """Minimal VAPID Draft 01 signature generation library.
+
+    https://tools.ietf.org/html/draft-ietf-webpush-vapid-01
+
+    """
     _private_key = None
     _public_key = None
     _hasher = hashlib.sha256
+    _schema = "WebPush"
 
     def __init__(self, private_key_file=None, private_key=None):
         """Initialize VAPID using an optional file containing a private key
@@ -137,6 +145,15 @@ def verify_token(self, validation_token, verification_token):
         return self.public_key.verify(hsig, validation_token,
                                       hashfunc=self._hasher)
 
+    def _base_sign(self, claims):
+        if not claims.get('exp'):
+            claims['exp'] = int(time.time()) + 86400
+        if not claims.get('sub'):
+            raise VapidException(
+                "Missing 'sub' from claims. "
+                "'sub' is your admin email as a mailto: link.")
+        return claims
+
     def sign(self, claims, crypto_key=None):
         """Sign a set of claims.
         :param claims: JSON object containing the JWT claims to use.
@@ -149,12 +166,7 @@ def sign(self, claims, crypto_key=None):
         :rtype: dict
 
         """
-        if not claims.get('exp'):
-            claims['exp'] = int(time.time()) + 86400
-        if not claims.get('sub'):
-            raise VapidException(
-                "Missing 'sub' from claims. "
-                "'sub' is your admin email as a mailto: link.")
+        claims = self._base_sign(claims)
         sig = jws.sign(claims, self.private_key, algorithm="ES256")
         pkey = 'p256ecdsa='
         pkey += base64.urlsafe_b64encode(self.public_key.to_string())
@@ -163,5 +175,32 @@ def sign(self, claims, crypto_key=None):
         else:
             crypto_key = pkey
 
-        return {"Authorization": "WebPush " + sig.strip('='),
+        return {"Authorization": "{} {}".format(self._schema, sig.strip('=')),
                 "Crypto-Key": crypto_key}
+
+
+class Vapid02(Vapid01):
+    """Minimal Vapid 02 signature generation library
+
+    https://tools.ietf.org/html/draft-ietf-webpush-vapid-02
+
+    """
+    _schema = "vapid"
+
+    def sign(self, claims, crypto_key=None):
+        claims = self._base_sign(claims)
+        sig = jws.sign(claims, self.private_key, algorithm="ES256")
+        pkey = self.public_key.to_string()
+        # Make sure that the key is properly prefixed.
+        if len(pkey) == 64:
+            pkey = '\04' + pkey
+        return{
+            "Authorization": "{schema} t={t},k={k}".format(
+                schema=self._schema,
+                t=sig,
+                k=base64.urlsafe_b64encode(pkey).strip('=')
+            )
+        }
+
+
+Vapid = Vapid01

python/py_vapid/main.py

@@ -6,14 +6,21 @@
 import os
 import json
 
-from py_vapid import Vapid
+from py_vapid import Vapid01, Vapid02
 
 
 def main():
     parser = argparse.ArgumentParser(description="VAPID tool")
     parser.add_argument('--sign', '-s', help='claims file to sign')
     parser.add_argument('--validate', '-v', help='dashboard token to validate')
+    parser.add_argument('--version2', '-2', help="use VAPID spec Draft-02",
+                        default=False, action="store_true")
+    parser.add_argument('--version1', '-1', help="use VAPID spec Draft-01",
+                        default=True, action="store_true")
     args = parser.parse_args()
+    Vapid = Vapid01
+    if args.version2:
+        Vapid = Vapid02
     if not os.path.exists('private_key.pem'):
         print "No private_key.pem file found."
         answer = None

python/py_vapid/tests/test_vapid.py

@@ -7,7 +7,7 @@
 from mock import patch
 
 from jose import jws
-from py_vapid import Vapid, VapidException
+from py_vapid import Vapid01, Vapid02, VapidException
 
 T_DER = """
 MHcCAQEEIPeN1iAipHbt8+/KZ2NIF8NeN24jqAmnMLFZEMocY8RboAoGCCqGSM49
@@ -43,54 +43,57 @@ def tearDown(self):
 
 class VapidTestCase(unittest.TestCase):
     def test_init(self):
-        v1 = Vapid(private_key_file="/tmp/private")
+        v1 = Vapid01(private_key_file="/tmp/private")
         eq_(v1.private_key.to_pem(), T_PRIVATE)
         eq_(v1.public_key.to_pem(), T_PUBLIC)
-        v2 = Vapid(private_key=T_PRIVATE)
+        v2 = Vapid01(private_key=T_PRIVATE)
         eq_(v2.private_key.to_pem(), T_PRIVATE)
         eq_(v2.public_key.to_pem(), T_PUBLIC)
-        v3 = Vapid(private_key=T_DER)
+        v3 = Vapid01(private_key=T_DER)
         eq_(v3.private_key.to_pem(), T_PRIVATE)
         eq_(v3.public_key.to_pem(), T_PUBLIC)
         no_exist = '/tmp/not_exist'
-        Vapid(private_key_file=no_exist)
+        Vapid01(private_key_file=no_exist)
         ok_(os.path.isfile(no_exist))
         os.unlink(no_exist)
 
+    def repad(self, data):
+        return data + b"===="[:len(data) % 4]
+
     @patch("ecdsa.SigningKey.from_pem", side_effect=Exception)
     def test_init_bad_priv(self, mm):
         self.assertRaises(Exception,
-                          Vapid,
+                          Vapid01,
                           private_key_file="/tmp/private")
 
     def test_private(self):
-        v = Vapid()
+        v = Vapid01()
         self.assertRaises(VapidException, lambda x=None: v.private_key)
 
     def test_public(self):
-        v = Vapid()
+        v = Vapid01()
 
         self.assertRaises(VapidException, lambda x=None: v.public_key)
 
     def test_gen_key(self):
-        v = Vapid()
+        v = Vapid01()
         v.generate_keys()
         ok_(v.public_key)
         ok_(v.private_key)
 
     def test_save_key(self):
-        v = Vapid()
+        v = Vapid01()
         v.save_key("/tmp/p2")
         os.unlink("/tmp/p2")
 
     def test_save_public_key(self):
-        v = Vapid()
+        v = Vapid01()
         v.generate_keys()
         v.save_public_key("/tmp/p2")
         os.unlink("/tmp/p2")
 
     def test_validate(self):
-        v = Vapid("/tmp/private")
+        v = Vapid01("/tmp/private")
         msg = "foobar"
         vtoken = v.validate(msg)
         ok_(v.public_key.verify(base64.urlsafe_b64decode(vtoken),
@@ -99,8 +102,8 @@ def test_validate(self):
         # test verify
         ok_(v.verify_token(msg, vtoken))
 
-    def test_sign(self):
-        v = Vapid("/tmp/private")
+    def test_sign_01(self):
+        v = Vapid01("/tmp/private")
         claims = {"aud": "example.com", "sub": "[email protected]"}
         result = v.sign(claims, "id=previous")
         eq_(result['Crypto-Key'],
@@ -114,8 +117,29 @@ def test_sign(self):
         eq_(result['Crypto-Key'],
             'p256ecdsa=' + T_PUBLIC_RAW)
 
+    def test_sign_02(self):
+        v = Vapid02("/tmp/private")
+        claims = {"aud": "example.com",
+                  "sub": "[email protected]",
+                  "foo": "extra value"}
+        result = v.sign(claims, "id=previous")
+        auth = result['Authorization']
+        eq_(auth[:6], 'vapid ')
+        ok_(' t=' in auth)
+        ok_(',k=' in auth)
+        parts = auth[6:].split(',')
+        eq_(len(parts), 2)
+        t_val = json.loads(base64.urlsafe_b64decode(
+            self.repad(parts[0][2:].split('.')[1])
+        ))
+        k_val = base64.urlsafe_b64decode(self.repad(parts[1][2:]))
+        eq_(k_val[0], "\04")
+        eq_(len(k_val), 65)
+        for k in claims:
+            eq_(t_val[k], claims[k])
+
     def test_bad_sign(self):
-        v = Vapid("/tmp/private")
+        v = Vapid01("/tmp/private")
         self.assertRaises(VapidException,
                           v.sign,
                           {'aud': "p.example.com"})