66from mock import patch , Mock
77from nose .tools import eq_ , ok_ , assert_raises
88import http_ece
9- import pyelliptic
9+ from cryptography .hazmat .primitives .asymmetric import ec
10+ from cryptography .hazmat .backends import default_backend
1011
1112from pywebpush import WebPusher , WebPushException , CaseInsensitiveDict , webpush
1213
@@ -21,17 +22,24 @@ class WebpushTestCase(unittest.TestCase):
2122 "M5xqEwuPM7VuQcyiLDhvovthPIXx+gsQRQ=="
2223 )
2324
24- def _gen_subscription_info (self , recv_key ,
25+ def _gen_subscription_info (self ,
26+ recv_key = None ,
2527 endpoint = "https://example.com/" ):
28+ if not recv_key :
29+ recv_key = ec .generate_private_key (ec .SECP256R1 , default_backend ())
2630 return {
2731 "endpoint" : endpoint ,
2832 "keys" : {
2933 'auth' : base64 .urlsafe_b64encode (os .urandom (16 )).strip (b'=' ),
30- 'p256dh' : base64 .urlsafe_b64encode (
31- recv_key .get_pubkey ()).strip (b'=' ),
34+ 'p256dh' : self ._get_pubkey_str (recv_key ),
3235 }
3336 }
3437
38+ def _get_pubkey_str (self , priv_key ):
39+ return base64 .urlsafe_b64encode (
40+ priv_key .public_key ().public_numbers ().encode_point ()
41+ ).strip (b'=' )
42+
3543 def test_init (self ):
3644 # use static values so we know what to look for in the reply
3745 subscription_info = {
@@ -72,14 +80,17 @@ def test_init(self):
7280
7381 def test_encode (self ):
7482 for content_encoding in ["aesgcm" , "aes128gcm" ]:
75- recv_key = pyelliptic .ECC (curve = "prime256v1" )
83+ recv_key = ec .generate_private_key (
84+ ec .SECP256R1 , default_backend ())
7685 subscription_info = self ._gen_subscription_info (recv_key )
7786 data = "Mary had a little lamb, with some nice mint jelly"
7887 push = WebPusher (subscription_info )
7988 encoded = push .encode (data , content_encoding = content_encoding )
80- keyid = base64 .urlsafe_b64encode (recv_key .get_pubkey ()[1 :])
81- http_ece .keys [keyid ] = recv_key
82- http_ece .labels [keyid ] = 'P-256'
89+ """
90+ crypto_key = base64.urlsafe_b64encode(
91+ self._get_pubkey_str(recv_key)
92+ ).strip(b'=')
93+ """
8394 # Convert these b64 strings into their raw, binary form.
8495 raw_salt = None
8596 if 'salt' in encoded :
@@ -94,15 +105,14 @@ def test_encode(self):
94105 encoded ['body' ],
95106 salt = raw_salt ,
96107 dh = raw_dh ,
97- keyid = keyid ,
98- authSecret = raw_auth ,
108+ private_key = recv_key ,
109+ auth_secret = raw_auth ,
99110 version = content_encoding
100111 )
101112 eq_ (decoded .decode ('utf8' ), data )
102113
103114 def test_bad_content_encoding (self ):
104- recv_key = pyelliptic .ECC (curve = "prime256v1" )
105- subscription_info = self ._gen_subscription_info (recv_key )
115+ subscription_info = self ._gen_subscription_info ()
106116 data = "Mary had a little lamb, with some nice mint jelly"
107117 push = WebPusher (subscription_info )
108118 self .assertRaises (WebPushException ,
@@ -112,8 +122,7 @@ def test_bad_content_encoding(self):
112122
113123 @patch ("requests.post" )
114124 def test_send (self , mock_post ):
115- recv_key = pyelliptic .ECC (curve = "prime256v1" )
116- subscription_info = self ._gen_subscription_info (recv_key )
125+ subscription_info = self ._gen_subscription_info ()
117126 headers = {"Crypto-Key" : "pre-existing" ,
118127 "Authentication" : "bearer vapid" }
119128 data = "Mary had a little lamb"
@@ -131,9 +140,7 @@ def test_send(self, mock_post):
131140 def test_send_vapid (self , mock_post ):
132141 mock_post .return_value = Mock ()
133142 mock_post .return_value .status_code = 200
134- recv_key = pyelliptic .ECC (curve = "prime256v1" )
135-
136- subscription_info = self ._gen_subscription_info (recv_key )
143+ subscription_info = self ._gen_subscription_info ()
137144 data = "Mary had a little lamb"
138145 webpush (
139146 subscription_info = subscription_info ,
@@ -165,43 +172,40 @@ def repad(str):
165172 def test_send_bad_vapid_no_key (self , mock_post ):
166173 mock_post .return_value = Mock ()
167174 mock_post .return_value .status_code = 200
168- recv_key = pyelliptic .ECC (curve = "prime256v1" )
169175
170- subscription_info = self ._gen_subscription_info (recv_key )
176+ subscription_info = self ._gen_subscription_info ()
171177 data = "Mary had a little lamb"
172178 assert_raises (WebPushException ,
173179 webpush ,
174180 subscription_info = subscription_info ,
175181 data = data ,
176182 vapid_claims = {
177- "aud" : "https://example.com" ,
178- 179- }
183+ "aud" : "https://example.com" ,
184+ 185+ }
180186 )
181187
182188 @patch ("requests.post" )
183189 def test_send_bad_vapid_bad_return (self , mock_post ):
184190 mock_post .return_value = Mock ()
185191 mock_post .return_value .status_code = 410
186- recv_key = pyelliptic .ECC (curve = "prime256v1" )
187192
188- subscription_info = self ._gen_subscription_info (recv_key )
193+ subscription_info = self ._gen_subscription_info ()
189194 data = "Mary had a little lamb"
190195 assert_raises (WebPushException ,
191196 webpush ,
192197 subscription_info = subscription_info ,
193198 data = data ,
194199 vapid_claims = {
195- "aud" : "https://example.com" ,
196- 197- },
200+ "aud" : "https://example.com" ,
201+ 202+ },
198203 vapid_private_key = self .vapid_key
199204 )
200205
201206 @patch ("requests.post" )
202207 def test_send_empty (self , mock_post ):
203- recv_key = pyelliptic .ECC (curve = "prime256v1" )
204- subscription_info = self ._gen_subscription_info (recv_key )
208+ subscription_info = self ._gen_subscription_info ()
205209 headers = {"Crypto-Key" : "pre-existing" ,
206210 "Authentication" : "bearer vapid" }
207211 WebPusher (subscription_info ).send ('' , headers )
@@ -214,16 +218,14 @@ def test_send_empty(self, mock_post):
214218 ok_ ('pre-existing' in ckey )
215219
216220 def test_encode_empty (self ):
217- recv_key = pyelliptic .ECC (curve = "prime256v1" )
218- subscription_info = self ._gen_subscription_info (recv_key )
221+ subscription_info = self ._gen_subscription_info ()
219222 headers = {"Crypto-Key" : "pre-existing" ,
220223 "Authentication" : "bearer vapid" }
221224 encoded = WebPusher (subscription_info ).encode ('' , headers )
222225 eq_ (encoded , None )
223226
224227 def test_encode_no_crypto (self ):
225- recv_key = pyelliptic .ECC (curve = "prime256v1" )
226- subscription_info = self ._gen_subscription_info (recv_key )
228+ subscription_info = self ._gen_subscription_info ()
227229 del (subscription_info ['keys' ])
228230 headers = {"Crypto-Key" : "pre-existing" ,
229231 "Authentication" : "bearer vapid" }
@@ -236,8 +238,7 @@ def test_encode_no_crypto(self):
236238
237239 @patch ("requests.post" )
238240 def test_send_no_headers (self , mock_post ):
239- recv_key = pyelliptic .ECC (curve = "prime256v1" )
240- subscription_info = self ._gen_subscription_info (recv_key )
241+ subscription_info = self ._gen_subscription_info ()
241242 data = "Mary had a little lamb"
242243 WebPusher (subscription_info ).send (data )
243244 eq_ (subscription_info .get ('endpoint' ), mock_post .call_args [0 ][0 ])
@@ -248,15 +249,14 @@ def test_send_no_headers(self, mock_post):
248249
249250 @patch ("pywebpush.open" )
250251 def test_as_curl (self , opener ):
251- recv_key = pyelliptic .ECC (curve = "prime256v1" )
252- subscription_info = self ._gen_subscription_info (recv_key )
252+ subscription_info = self ._gen_subscription_info ()
253253 result = webpush (
254254 subscription_info ,
255255 data = "Mary had a little lamb" ,
256256 vapid_claims = {
257- "aud" : "https://example.com" ,
258- 259- },
257+ "aud" : "https://example.com" ,
258+ 259+ },
260260 vapid_private_key = self .vapid_key ,
261261 curl = True
262262 )
@@ -281,9 +281,8 @@ def test_ci_dict(self):
281281
282282 @patch ("requests.post" )
283283 def test_gcm (self , mock_post ):
284- recv_key = pyelliptic .ECC (curve = "prime256v1" )
285284 subscription_info = self ._gen_subscription_info (
286- recv_key ,
285+ None ,
287286 endpoint = "https://android.googleapis.com/gcm/send/regid123" )
288287 headers = {"Crypto-Key" : "pre-existing" ,
289288 "Authentication" : "bearer vapid" }
0 commit comments