Webpush headers (#123)

* Support adding arbitrary HTTP headers when using webpush()

Useful for adding headers like `Topic` an `Urgency` that don't have
dedicated parameters.

* Send headers provided in the `--head` CLI argument

Author: braedon

pywebpush/__init__.py

@@ -382,7 +382,8 @@ def webpush(subscription_info,
             curl=False,
             timeout=None,
             ttl=0,
-            verbose=False):
+            verbose=False,
+            headers=None):
     """
         One call solution to endcode and send `data` to the endpoint
         contained in `subscription_info` using optional VAPID auth headers.
@@ -427,8 +428,16 @@ def webpush(subscription_info,
     :param verbose: Provide verbose feedback
     :type verbose: bool
     :return requests.Response or string
+    :param headers: Dictionary of extra HTTP headers to include
+    :type headers: dict
 
     """
+    if headers is None:
+        headers = dict()
+    else:
+        # Ensure we don't leak VAPID headers by mutating the passed in dict.
+        headers = headers.copy()
+
     vapid_headers = None
     if vapid_claims:
         if verbose:
@@ -462,10 +471,11 @@ def webpush(subscription_info,
         vapid_headers = vv.sign(vapid_claims)
         if verbose:
             print("\t headers: {}".format(vapid_headers))
+        headers.update(vapid_headers)
 
     response = WebPusher(subscription_info, verbose=verbose).send(
         data,
-        vapid_headers,
+        headers,
         ttl=ttl,
         content_encoding=content_encoding,
         curl=curl,

pywebpush/__main__.py

@@ -9,7 +9,7 @@ def get_config():
     parser = argparse.ArgumentParser(description="WebPush tool")
     parser.add_argument("--data", '-d', help="Data file")
     parser.add_argument("--info", "-i", help="Subscription Info JSON file")
-    parser.add_argument("--head",  help="Header Info JSON file")
+    parser.add_argument("--head", help="Header Info JSON file")
     parser.add_argument("--claims", help="Vapid claim file")
     parser.add_argument("--key", help="Vapid private key file path")
     parser.add_argument("--curl", help="Don't send, display as curl command",
@@ -56,7 +56,8 @@ def main():
             vapid_claims=args.claims,
             curl=args.curl,
             content_encoding=args.encoding,
-            verbose=args.verbose)
+            verbose=args.verbose,
+            headers=args.head)
         print(result)
     except Exception as ex:
         print("ERROR: {}".format(ex))

pywebpush/tests/test_webpush.py

@@ -154,7 +154,8 @@ def test_send_vapid(self, mock_post):
             data=data,
             vapid_private_key=self.vapid_key,
             vapid_claims={"sub": "mailto:ops@example.com"},
-            content_encoding="aesgcm"
+            content_encoding="aesgcm",
+            headers={"Test-Header": "test-value"}
         )
         eq_(subscription_info.get('endpoint'), mock_post.call_args[0][0])
         pheaders = mock_post.call_args[1].get('headers')
@@ -173,6 +174,7 @@ def repad(str):
         ckey = pheaders.get('crypto-key')
         ok_('dh=' in ckey)
         eq_(pheaders.get('content-encoding'), 'aesgcm')
+        eq_(pheaders.get('test-header'), 'test-value')
 
     @patch.object(WebPusher, "send")
     @patch.object(py_vapid.Vapid, "sign")