Merge pull request #5 from mozilla-services/bug/03

feat: make python2/3 compatible

Author: jrconlin

.coveragerc

@@ -1,2 +1,3 @@
 [report]
 omit = *noseplugin*
+show_missing = True

CHANGELOG.md

@@ -1,3 +1,6 @@
+## 0.4.0 (2016-06-05)
+feat: make python 2.7 / 3.5 polyglot
+
 ## 0.3.4 (2016-05-17)
 bug: make header keys case insenstive
 

README.md

@@ -36,8 +36,9 @@ The data can be any serial content (string, bit array, serialized
 JSON, etc), but be sure that your receiving application is able to
 parse and understand it. (e.g. `data = "Mary had a little lamb."`)
 
-gcm_key is the API key obtained from the Google Developer Console. 
-It is only needed if endpoint is https://android.googleapis.com/gcm/send
+gcm_key is the API key obtained from the Google Developer Console.
+It is only needed if endpoint is
+https://android.googleapis.com/gcm/send
 
 `headers` is a `dict`ionary of additional HTTP header values (e.g.
 [VAPID](https://github.com/mozilla-services/vapid/tree/master/python)
@@ -50,7 +51,6 @@ WebPusher(subscription_info).send(data, headers)
 to send for Chrome:
 ```
 WebPusher(subscription_info).send(data, headers, ttl, gcm_key)
-```
 
 You can also simply encode the data to send later by calling
 

pywebpush/__init__.py

@@ -3,9 +3,10 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import base64
+import json
 import os
 
-import json
+import six
 import http_ece
 import pyelliptic
 import requests
@@ -100,22 +101,22 @@ def __init__(self, subscription_info):
         for k in ['p256dh', 'auth']:
             if keys.get(k) is None:
                 raise WebPushException("Missing keys value: %s", k)
-        receiver_raw = base64.urlsafe_b64decode(
-            self._repad(keys['p256dh'].encode('utf8')))
+            if isinstance(keys[k], six.string_types):
+                keys[k] = bytes(keys[k].encode('utf8'))
+        receiver_raw = base64.urlsafe_b64decode(self._repad(keys['p256dh']))
         if len(receiver_raw) != 65 and receiver_raw[0] != "\x04":
             raise WebPushException("Invalid p256dh key specified")
         self.receiver_key = receiver_raw
-        self.auth_key = base64.urlsafe_b64decode(
-            self._repad(keys['auth'].encode('utf8')))
+        self.auth_key = base64.urlsafe_b64decode(self._repad(keys['auth']))
 
-    def _repad(self, str):
+    def _repad(self, data):
         """Add base64 padding to the end of a string, if required"""
-        return str + "===="[:len(str) % 4]
+        return data + b"===="[:len(data) % 4]
 
     def encode(self, data):
         """Encrypt the data.
 
-        :param data: A serialized block of data (String, JSON, bit array,
+        :param data: A serialized block of byte data (String, JSON, bit array,
             etc.) Make sure that whatever you send, your client knows how
             to understand it.
 
@@ -129,6 +130,9 @@ def encode(self, data):
         # ID tag.
         server_key_id = base64.urlsafe_b64encode(server_key.get_pubkey()[1:])
 
+        if isinstance(data, six.string_types):
+            data = bytes(data.encode('utf8'))
+
         # http_ece requires that these both be set BEFORE encrypt or
         # decrypt is called if you specify the key as "dh".
         http_ece.keys[server_key_id] = server_key
@@ -143,8 +147,8 @@ def encode(self, data):
 
         return CaseInsensitiveDict({
             'crypto_key': base64.urlsafe_b64encode(
-                server_key.get_pubkey()).strip('='),
-            'salt': base64.urlsafe_b64encode(salt).strip("="),
+                server_key.get_pubkey()).strip(b'='),
+            'salt': base64.urlsafe_b64encode(salt).strip(b'='),
             'body': encrypted,
         })
 
@@ -160,6 +164,10 @@ def send(self, data, headers={}, ttl=0, gcm_key=None, reg_id=None):
         :param ttl: The Time To Live in seconds for this message if the
             recipient is not online. (Defaults to "0", which discards the
             message immediately if the recipient is unavailable.)
+        :param gcm_key: API key obtained from the Google Developer Console.
+            Needed if endpoint is https://android.googleapis.com/gcm/send
+        :param reg_id: registration id of the recipient. If not provided,
+            it will be extracted from the endpoint.
 
         """
         # Encode the data.
@@ -169,11 +177,12 @@ def send(self, data, headers={}, ttl=0, gcm_key=None, reg_id=None):
         crypto_key = headers.get("crypto-key", "")
         if crypto_key:
             crypto_key += ','
-        crypto_key += "keyid=p256dh;dh=" + encoded["crypto_key"]
+        crypto_key += "keyid=p256dh;dh=" + encoded["crypto_key"].decode('utf8')
         headers.update({
             'crypto-key': crypto_key,
             'content-encoding': 'aesgcm',
-            'encryption': "keyid=p256dh;salt=" + encoded['salt'],
+            'encryption': "keyid=p256dh;salt=" +
+            encoded['salt'].decode('utf8'),
         })
         gcm_endpoint = 'https://android.googleapis.com/gcm/send'
         if self.subscription_info['endpoint'].startswith(gcm_endpoint):

pywebpush/tests/test_webpush.py

@@ -16,9 +16,9 @@ def _gen_subscription_info(self, recv_key, endpoint="https://example.com"):
         return {
             "endpoint": endpoint,
             "keys": {
-                'auth': base64.urlsafe_b64encode(os.urandom(16)).strip('='),
+                'auth': base64.urlsafe_b64encode(os.urandom(16)).strip(b'='),
                 'p256dh': base64.urlsafe_b64encode(
-                    recv_key.get_pubkey()).strip('='),
+                    recv_key.get_pubkey()).strip(b'='),
             }
         }
 
@@ -32,6 +32,11 @@ def test_init(self):
                 u"auth": u"k8JV6sjdbhAi1n3_LDBLvA"
             }
         }
+        rk_decode = (b'\x04\xea\xe7"\xc9W\xadJ0\xd9P3(%\x00\x13\x8b'
+                     b'\x08l\xad4u\xa1\x19\n\xcc\x0eq\xff&\xdd1'
+                     b'|W\xcd\x81\xf8\xeaN\x83\x92[\x99\x82\xe0\xe3'
+                     b'\x89\x11r\xf4\x02\xd4M\xa00\x9b!\xb1F\x00'
+                     b'\xfb\xfc\xcc=\x1f')
         self.assertRaises(
             WebPushException,
             WebPusher,
@@ -56,12 +61,8 @@ def test_init(self):
 
         push = WebPusher(subscription_info)
         eq_(push.subscription_info, subscription_info)
-        eq_(push.receiver_key, ('\x04\xea\xe7"\xc9W\xadJ0\xd9P3(%\x00\x13\x8b'
-                                '\x08l\xad4u\xa1\x19\n\xcc\x0eq\xff&\xdd1'
-                                '|W\xcd\x81\xf8\xeaN\x83\x92[\x99\x82\xe0\xe3'
-                                '\x89\x11r\xf4\x02\xd4M\xa00\x9b!\xb1F\x00'
-                                '\xfb\xfc\xcc=\x1f'))
-        eq_(push.auth_key, '\x93\xc2U\xea\xc8\xddn\x10"\xd6}\xff,0K\xbc')
+        eq_(push.receiver_key, rk_decode)
+        eq_(push.auth_key, b'\x93\xc2U\xea\xc8\xddn\x10"\xd6}\xff,0K\xbc')
 
     def test_encode(self):
         recv_key = pyelliptic.ECC(curve="prime256v1")
@@ -89,7 +90,7 @@ def test_encode(self):
             authSecret=raw_auth
             )
 
-        eq_(decoded, data)
+        eq_(decoded.decode('utf8'), data)
 
     @patch("requests.post")
     def test_send(self, mock_post):

setup.py

@@ -3,7 +3,7 @@
 
 from setuptools import find_packages, setup
 
-__version__ = "0.3.4"
+__version__ = "0.4.0"
 
 
 def read_from(file):