Skip to content

Commit 349c12f

Browse files
jrconlinjrconlin
authored
Merge pull request #36 from kitcambridge/kit/init-from-raw
Refactor initialization; add `from_raw`
2 parents d169f22 + 8c92573 commit 349c12f

File tree

4 files changed

+108
-42
lines changed

4 files changed

+108
-42
lines changed

‎python/py_vapid/__init__.py‎

Lines changed: 72 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@
1616
VERSION = "VAPID-DRAFT-02/ECE-DRAFT-07"
1717

1818

19+
def b64urldecode(data):
20+
"""Decodes an unpadded Base64url-encoded string."""
21+
return base64.urlsafe_b64decode(data + "===="[:len(data) % 4])
22+
23+
1924
class VapidException(Exception):
2025
"""An exception wrapper for Vapid."""
2126
pass
@@ -29,36 +34,79 @@ class Vapid01(object):
2934
"""
3035
_private_key = None
3136
_public_key = None
37+
_curve = ecdsa.NIST256p
3238
_hasher = hashlib.sha256
3339
_schema = "WebPush"
3440

35-
def __init__(self, private_key_file=None, private_key=None):
36-
"""Initialize VAPID using an optional file containing a private key
37-
in PEM format, or a string containing the PEM formatted private key.
41+
def __init__(self, private_key=None):
42+
"""Initialize VAPID with an optional private key.
43+
44+
:param private_key: A private key object
45+
:type private_key: ecdsa.SigningKey
46+
47+
"""
48+
self.private_key = private_key
49+
50+
@classmethod
51+
def from_raw(cls, private_key):
52+
"""Initialize VAPID using a private key point in "raw" or
53+
"uncompressed" form.
54+
55+
:param private_key: A private key point in uncompressed form.
56+
:type private_key: str
57+
58+
"""
59+
key = ecdsa.SigningKey.from_string(b64urldecode(private_key),
60+
curve=cls._curve,
61+
hashfunc=cls._hasher)
62+
return cls(key)
63+
64+
@classmethod
65+
def from_pem(cls, private_key):
66+
"""Initialize VAPID using a private key in PEM format.
67+
68+
:param private_key: A private key in PEM format.
69+
:type private_key: str
70+
71+
"""
72+
key = ecdsa.SigningKey.from_pem(private_key)
73+
return cls(key)
74+
75+
@classmethod
76+
def from_der(cls, private_key):
77+
"""Initialize VAPID using a private key in DER format.
78+
79+
:param private_key: A private key in DER format and Base64-encoded.
80+
:type private_key: str
81+
82+
"""
83+
key = ecdsa.SigningKey.from_der(base64.b64decode(private_key))
84+
return cls(key)
85+
86+
@classmethod
87+
def from_file(cls, private_key_file=None):
88+
"""Initialize VAPID using a file containing a private key in PEM or
89+
DER format.
3890
3991
:param private_key_file: Name of the file containing the private key
4092
:type private_key_file: str
41-
:param private_key: A private key in PEM format
42-
:type private_key: str
4393
4494
"""
45-
if private_key_file:
46-
if not os.path.isfile(private_key_file):
47-
self.save_key(private_key_file)
48-
return
49-
private_key = open(private_key_file, 'r').read()
50-
if private_key:
51-
try:
52-
if "BEGIN EC" in private_key:
53-
self._private_key = ecdsa.SigningKey.from_pem(private_key)
54-
else:
55-
self._private_key = \
56-
ecdsa.SigningKey.from_der(
57-
base64.urlsafe_b64decode(private_key))
58-
except Exception as exc:
59-
logging.error("Could not open private key file: %s", repr(exc))
60-
raise VapidException(exc)
61-
self._public_key = self._private_key.get_verifying_key()
95+
if not os.path.isfile(private_key_file):
96+
vapid = cls()
97+
vapid.save_key(private_key_file)
98+
return vapid
99+
private_key = open(private_key_file, 'r').read()
100+
vapid = None
101+
try:
102+
if "BEGIN EC" in private_key:
103+
vapid = cls.from_pem(private_key)
104+
else:
105+
vapid = cls.from_der(private_key)
106+
except Exception as exc:
107+
logging.error("Could not open private key file: %s", repr(exc))
108+
raise VapidException(exc)
109+
return vapid
62110

63111
@property
64112
def private_key(self):
@@ -77,6 +125,7 @@ def private_key(self, value):
77125
78126
"""
79127
self._private_key = value
128+
self._public_key = None
80129

81130
@property
82131
def public_key(self):
@@ -92,8 +141,7 @@ def public_key(self):
92141

93142
def generate_keys(self):
94143
"""Generate a valid ECDSA Key Pair."""
95-
self.private_key = ecdsa.SigningKey.generate(curve=ecdsa.NIST256p)
96-
self._public_key = self.private_key.get_verifying_key()
144+
self.private_key = ecdsa.SigningKey.generate(curve=self._curve)
97145

98146
def save_key(self, key_file):
99147
"""Save the private key to a PEM file.

‎python/py_vapid/main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def main():
3939
exit
4040
print("Generating private_key.pem")
4141
Vapid().save_key('private_key.pem')
42-
vapid = Vapid('private_key.pem')
42+
vapid = Vapid.from_file('private_key.pem')
4343
if args.gen or not os.path.exists('public_key.pem'):
4444
if not args.gen:
4545
print("No public_key.pem file found. You'll need this to access "

‎python/py_vapid/tests/test_vapid.py‎

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,21 +10,31 @@
1010
from jose import jws
1111
from py_vapid import Vapid01, Vapid02, VapidException
1212

13+
# This is a private key in DER form.
1314
T_DER = """
1415
MHcCAQEEIPeN1iAipHbt8+/KZ2NIF8NeN24jqAmnMLFZEMocY8RboAoGCCqGSM49
1516
AwEHoUQDQgAEEJwJZq/GN8jJbo1GGpyU70hmP2hbWAUpQFKDByKB81yldJ9GTklB
1617
M5xqEwuPM7VuQcyiLDhvovthPIXx+gsQRQ==
1718
"""
19+
20+
# This is the same private key, in PEM form.
1821
T_PRIVATE = ("-----BEGIN EC PRIVATE KEY-----{}"
1922
"-----END EC PRIVATE KEY-----\n").format(T_DER)
2023

24+
# This is the same private key, as a point in uncompressed form. This should
25+
# be Base64url-encoded without padding.
26+
T_RAW = """
27+
943WICKkdu3z78pnY0gXw143biOoCacwsVkQyhxjxFs
28+
"""
29+
30+
# This is a public key in PEM form.
2131
T_PUBLIC = """-----BEGIN PUBLIC KEY-----
2232
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEJwJZq/GN8jJbo1GGpyU70hmP2hb
2333
WAUpQFKDByKB81yldJ9GTklBM5xqEwuPM7VuQcyiLDhvovthPIXx+gsQRQ==
2434
-----END PUBLIC KEY-----
2535
"""
2636

27-
# this is a DER RAW key ('\x04' + 2 32 octet digits)
37+
# this is a public key in uncompressed form ('\x04' + 2 * 32 octets)
2838
# Remember, this should have any padding stripped.
2939
T_PUBLIC_RAW = (
3040
"BBCcCWavxjfIyW6NRhqclO9IZj9oW1gFKUBSgwcigfNc"
@@ -33,12 +43,12 @@
3343

3444

3545
def setUp(self):
36-
ff = open('/tmp/private', 'w')
37-
ff.write(T_PRIVATE)
38-
ff.close()
39-
ff = open('/tmp/public', 'w')
40-
ff.write(T_PUBLIC)
41-
ff.close()
46+
with open('/tmp/private', 'w') as ff:
47+
ff.write(T_PRIVATE)
48+
with open('/tmp/public', 'w') as ff:
49+
ff.write(T_PUBLIC)
50+
with open('/tmp/private.der', 'w') as ff:
51+
ff.write(T_DER)
4252

4353

4454
def tearDown(self):
@@ -48,17 +58,20 @@ def tearDown(self):
4858

4959
class VapidTestCase(unittest.TestCase):
5060
def test_init(self):
51-
v1 = Vapid01(private_key_file="/tmp/private")
61+
v1 = Vapid01.from_file("/tmp/private")
5262
eq_(v1.private_key.to_pem(), T_PRIVATE.encode('utf8'))
5363
eq_(v1.public_key.to_pem(), T_PUBLIC.encode('utf8'))
54-
v2 = Vapid01(private_key=T_PRIVATE)
64+
v2 = Vapid01.from_pem(T_PRIVATE)
5565
eq_(v2.private_key.to_pem(), T_PRIVATE.encode('utf8'))
5666
eq_(v2.public_key.to_pem(), T_PUBLIC.encode('utf8'))
57-
v3 = Vapid01(private_key=T_DER)
67+
v3 = Vapid01.from_der(T_DER)
5868
eq_(v3.private_key.to_pem(), T_PRIVATE.encode('utf8'))
5969
eq_(v3.public_key.to_pem(), T_PUBLIC.encode('utf8'))
70+
v4 = Vapid01.from_file("/tmp/private.der")
71+
eq_(v4.private_key.to_pem(), T_PRIVATE.encode('utf8'))
72+
eq_(v4.public_key.to_pem(), T_PUBLIC.encode('utf8'))
6073
no_exist = '/tmp/not_exist'
61-
Vapid01(private_key_file=no_exist)
74+
Vapid01.from_file(no_exist)
6275
ok_(os.path.isfile(no_exist))
6376
os.unlink(no_exist)
6477

@@ -68,7 +81,7 @@ def repad(self, data):
6881
@patch("ecdsa.SigningKey.from_pem", side_effect=Exception)
6982
def test_init_bad_priv(self, mm):
7083
self.assertRaises(Exception,
71-
Vapid01,
84+
Vapid01.from_file,
7285
private_key_file="/tmp/private")
7386

7487
def test_private(self):
@@ -97,8 +110,13 @@ def test_same_public_key(self):
97110
v.save_public_key("/tmp/p2")
98111
os.unlink("/tmp/p2")
99112

113+
def test_from_raw(self):
114+
v = Vapid01.from_raw(T_RAW)
115+
eq_(v.private_key.to_pem(), T_PRIVATE.encode('utf8'))
116+
eq_(v.public_key.to_pem(), T_PUBLIC.encode('utf8'))
117+
100118
def test_validate(self):
101-
v = Vapid01("/tmp/private")
119+
v = Vapid01.from_file("/tmp/private")
102120
msg = "foobar".encode('utf8')
103121
vtoken = v.validate(msg)
104122
ok_(v.public_key.verify(base64.urlsafe_b64decode(vtoken),
@@ -108,7 +126,7 @@ def test_validate(self):
108126
ok_(v.verify_token(msg, vtoken))
109127

110128
def test_sign_01(self):
111-
v = Vapid01("/tmp/private")
129+
v = Vapid01.from_file("/tmp/private")
112130
claims = {"aud": "example.com", "sub": "[email protected]"}
113131
result = v.sign(claims, "id=previous")
114132
eq_(result['Crypto-Key'],
@@ -123,7 +141,7 @@ def test_sign_01(self):
123141
'p256ecdsa=' + T_PUBLIC_RAW)
124142

125143
def test_sign_02(self):
126-
v = Vapid02("/tmp/private")
144+
v = Vapid02.from_file("/tmp/private")
127145
claims = {"aud": "example.com",
128146
"sub": "[email protected]",
129147
"foo": "extra value"}
@@ -144,7 +162,7 @@ def test_sign_02(self):
144162
eq_(t_val[k], claims[k])
145163

146164
def test_bad_sign(self):
147-
v = Vapid01("/tmp/private")
165+
v = Vapid01.from_file("/tmp/private")
148166
self.assertRaises(VapidException,
149167
v.sign,
150168
{'aud': "p.example.com"})

‎python/setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from setuptools import setup, find_packages
55

6-
__version__ = "0.8.1"
6+
__version__ = "1.0.0"
77

88

99
def read_from(file):

0 commit comments

Comments
 (0)