Factor out initialization into separate class methods.

This commit splits the `Vapid` constructor into three class methods:

* `Vapid.from_file` reads a private key file in PEM or DER format, and
  creates a `Vapid` object based on the contents. This is equivalent
  to calling `Vapid(private_key_file="/path")`.
* `Vapid.from_pem` creates a `Vapid` object from a string containing a
  PEM-encoded private key.
* `Vapid.from_der` creates a `Vapid` object from a string containing a
  DER-encoded key.

Author: Kit Cambridge

python/py_vapid/__init__.py

@@ -29,36 +29,65 @@ class Vapid01(object):
     """
     _private_key = None
     _public_key = None
+    _curve = ecdsa.NIST256p
     _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
-        in PEM format, or a string containing the PEM formatted private key.
+    def __init__(self, private_key=None):
+        """Initialize VAPID with an optional private key.
+
+        :param private_key: A private key object
+        :type private_key: ecdsa.SigningKey
+
+        """
+        self.private_key = private_key
+
+    @classmethod
+    def from_pem(cls, private_key):
+        """Initialize VAPID using a private key in PEM format.
+
+        :param private_key: A private key in PEM format.
+        :type private_key: str
+
+        """
+        key = ecdsa.SigningKey.from_pem(private_key)
+        return cls(key)
+
+    @classmethod
+    def from_der(cls, private_key):
+        """Initialize VAPID using a private key in DER format.
+
+        :param private_key: A private key in DER format and Base64-encoded.
+        :type private_key: str
+
+        """
+        key = ecdsa.SigningKey.from_der(base64.b64decode(private_key))
+        return cls(key)
+
+    @classmethod
+    def from_file(cls, private_key_file=None):
+        """Initialize VAPID using a file containing a private key in PEM or
+        DER format.
 
         :param private_key_file: Name of the file containing the private key
         :type private_key_file: str
-        :param private_key: A private key in PEM format
-        :type private_key: str
 
         """
-        if private_key_file:
-            if not os.path.isfile(private_key_file):
-                self.save_key(private_key_file)
-                return
-            private_key = open(private_key_file, 'r').read()
-        if private_key:
-            try:
-                if "BEGIN EC" in private_key:
-                    self._private_key = ecdsa.SigningKey.from_pem(private_key)
-                else:
-                    self._private_key = \
-                        ecdsa.SigningKey.from_der(
-                            base64.urlsafe_b64decode(private_key))
-            except Exception as exc:
-                logging.error("Could not open private key file: %s", repr(exc))
-                raise VapidException(exc)
-            self._public_key = self._private_key.get_verifying_key()
+        if not os.path.isfile(private_key_file):
+            vapid = cls()
+            vapid.save_key(private_key_file)
+            return vapid
+        private_key = open(private_key_file, 'r').read()
+        vapid = None
+        try:
+            if "BEGIN EC" in private_key:
+                vapid = cls.from_pem(private_key)
+            else:
+                vapid = cls.from_der(private_key)
+        except Exception as exc:
+            logging.error("Could not open private key file: %s", repr(exc))
+            raise VapidException(exc)
+        return vapid
 
     @property
     def private_key(self):
@@ -77,6 +106,7 @@ def private_key(self, value):
 
         """
         self._private_key = value
+        self._public_key = None
 
     @property
     def public_key(self):
@@ -92,8 +122,7 @@ def public_key(self):
 
     def generate_keys(self):
         """Generate a valid ECDSA Key Pair."""
-        self.private_key = ecdsa.SigningKey.generate(curve=ecdsa.NIST256p)
-        self._public_key = self.private_key.get_verifying_key()
+        self.private_key = ecdsa.SigningKey.generate(curve=self._curve)
 
     def save_key(self, key_file):
         """Save the private key to a PEM file.

python/py_vapid/main.py

@@ -39,7 +39,7 @@ def main():
                     exit
         print("Generating private_key.pem")
         Vapid().save_key('private_key.pem')
-    vapid = Vapid('private_key.pem')
+    vapid = Vapid.from_file('private_key.pem')
     if args.gen or not os.path.exists('public_key.pem'):
         if not args.gen:
             print("No public_key.pem file found. You'll need this to access "

python/py_vapid/tests/test_vapid.py

@@ -10,21 +10,25 @@
 from jose import jws
 from py_vapid import Vapid01, Vapid02, VapidException
 
+# This is a private key in DER form.
 T_DER = """
 MHcCAQEEIPeN1iAipHbt8+/KZ2NIF8NeN24jqAmnMLFZEMocY8RboAoGCCqGSM49
 AwEHoUQDQgAEEJwJZq/GN8jJbo1GGpyU70hmP2hbWAUpQFKDByKB81yldJ9GTklB
 M5xqEwuPM7VuQcyiLDhvovthPIXx+gsQRQ==
 """
+
+# This is the same private key, in PEM form.
 T_PRIVATE = ("-----BEGIN EC PRIVATE KEY-----{}"
              "-----END EC PRIVATE KEY-----\n").format(T_DER)
 
+# This is a public key in PEM form.
 T_PUBLIC = """-----BEGIN PUBLIC KEY-----
 MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEJwJZq/GN8jJbo1GGpyU70hmP2hb
 WAUpQFKDByKB81yldJ9GTklBM5xqEwuPM7VuQcyiLDhvovthPIXx+gsQRQ==
 -----END PUBLIC KEY-----
 """
 
-# this is a DER RAW key ('\x04' + 2 32 octet digits)
+# this is a public key in uncompressed form ('\x04' + 2 * 32 octets)
 # Remember, this should have any padding stripped.
 T_PUBLIC_RAW = (
     "BBCcCWavxjfIyW6NRhqclO9IZj9oW1gFKUBSgwcigfNc"
@@ -33,12 +37,12 @@
 
 
 def setUp(self):
-    ff = open('/tmp/private', 'w')
-    ff.write(T_PRIVATE)
-    ff.close()
-    ff = open('/tmp/public', 'w')
-    ff.write(T_PUBLIC)
-    ff.close()
+    with open('/tmp/private', 'w') as ff:
+        ff.write(T_PRIVATE)
+    with open('/tmp/public', 'w') as ff:
+        ff.write(T_PUBLIC)
+    with open('/tmp/private.der', 'w') as ff:
+        ff.write(T_DER)
 
 
 def tearDown(self):
@@ -48,17 +52,20 @@ def tearDown(self):
 
 class VapidTestCase(unittest.TestCase):
     def test_init(self):
-        v1 = Vapid01(private_key_file="/tmp/private")
+        v1 = Vapid01.from_file("/tmp/private")
         eq_(v1.private_key.to_pem(), T_PRIVATE.encode('utf8'))
         eq_(v1.public_key.to_pem(), T_PUBLIC.encode('utf8'))
-        v2 = Vapid01(private_key=T_PRIVATE)
+        v2 = Vapid01.from_pem(T_PRIVATE)
         eq_(v2.private_key.to_pem(), T_PRIVATE.encode('utf8'))
         eq_(v2.public_key.to_pem(), T_PUBLIC.encode('utf8'))
-        v3 = Vapid01(private_key=T_DER)
+        v3 = Vapid01.from_der(T_DER)
         eq_(v3.private_key.to_pem(), T_PRIVATE.encode('utf8'))
         eq_(v3.public_key.to_pem(), T_PUBLIC.encode('utf8'))
+        v4 = Vapid01.from_file("/tmp/private.der")
+        eq_(v4.private_key.to_pem(), T_PRIVATE.encode('utf8'))
+        eq_(v4.public_key.to_pem(), T_PUBLIC.encode('utf8'))
         no_exist = '/tmp/not_exist'
-        Vapid01(private_key_file=no_exist)
+        Vapid01.from_file(no_exist)
         ok_(os.path.isfile(no_exist))
         os.unlink(no_exist)
 
@@ -68,7 +75,7 @@ def repad(self, data):
     @patch("ecdsa.SigningKey.from_pem", side_effect=Exception)
     def test_init_bad_priv(self, mm):
         self.assertRaises(Exception,
-                          Vapid01,
+                          Vapid01.from_file,
                           private_key_file="/tmp/private")
 
     def test_private(self):
@@ -98,7 +105,7 @@ def test_same_public_key(self):
         os.unlink("/tmp/p2")
 
     def test_validate(self):
-        v = Vapid01("/tmp/private")
+        v = Vapid01.from_file("/tmp/private")
         msg = "foobar".encode('utf8')
         vtoken = v.validate(msg)
         ok_(v.public_key.verify(base64.urlsafe_b64decode(vtoken),
@@ -108,7 +115,7 @@ def test_validate(self):
         ok_(v.verify_token(msg, vtoken))
 
     def test_sign_01(self):
-        v = Vapid01("/tmp/private")
+        v = Vapid01.from_file("/tmp/private")
         claims = {"aud": "example.com", "sub": "[email protected]"}
         result = v.sign(claims, "id=previous")
         eq_(result['Crypto-Key'],
@@ -123,7 +130,7 @@ def test_sign_01(self):
             'p256ecdsa=' + T_PUBLIC_RAW)
 
     def test_sign_02(self):
-        v = Vapid02("/tmp/private")
+        v = Vapid02.from_file("/tmp/private")
         claims = {"aud": "example.com",
                   "sub": "[email protected]",
                   "foo": "extra value"}
@@ -144,7 +151,7 @@ def test_sign_02(self):
             eq_(t_val[k], claims[k])
 
     def test_bad_sign(self):
-        v = Vapid01("/tmp/private")
+        v = Vapid01.from_file("/tmp/private")
         self.assertRaises(VapidException,
                           v.sign,
                           {'aud': "p.example.com"})