Allow to send headers to proxied documents

src/h_vialib/client.py

@@ -5,7 +5,7 @@
 
 from webob.multidict import MultiDict
 
-from h_vialib.secure import ViaSecureURL
+from h_vialib.secure import Encryption, ViaSecureURL
 
 
 class ViaDoc:
@@ -45,6 +45,7 @@ def __init__(self, secret, service_url=None, html_service_url=None):
         :param service_url: Location of the via server
         :param html_service_url: Location of the Via HTML presenter
         """
+        self._secure_secrets = Encryption(secret.encode("utf-8"))
         self._secure_url = ViaSecureURL(secret)
         self._service_url = urlparse(service_url) if service_url else None
         self._html_service_url = html_service_url
@@ -56,7 +57,10 @@ def __init__(self, secret, service_url=None, html_service_url=None):
             "via.external_link_mode": "new-tab",
         }
 
-    def url_for(self, url, content_type=None, options=None, blocked_for=None):
+    # pylint:disable=too-many-arguments
+    def url_for(
+        self, url, content_type=None, options=None, blocked_for=None, headers=None
+    ):
         """Generate a Via URL to display a given URL.
 
         If provided, the options will be merged with default Via options.
@@ -66,6 +70,10 @@ def url_for(self, url, content_type=None, options=None, blocked_for=None):
             or "html")
         :param options: Any additional params to add to the URL
         :param blocked_for: context for the blocked pages
+        :param headers: Any headers needed to make the request to `url`.
+            The headers are sent encrypted to Via.
+            Note that the headers are only used when the content is proxied with
+            python, ie only PDFs from certain sources.
         :return: Full Via URL suitable for redirecting a user to
         """
         doc = ViaDoc(url, content_type)
@@ -74,6 +82,9 @@ def url_for(self, url, content_type=None, options=None, blocked_for=None):
         if options:
             query.update(options)
 
+        if headers:
+            query["via.secret.headers"] = self._secure_secrets.encrypt_dict(headers)
+
         if blocked_for:
             query["via.blocked_for"] = blocked_for
 

src/h_vialib/secure/__init__.py

@@ -1,5 +1,6 @@
 """Security helpers."""
 
+from h_vialib.secure.encryption import Encryption
 from h_vialib.secure.expiry import quantized_expiry
 from h_vialib.secure.token import SecureToken
 from h_vialib.secure.url import SecureURL, ViaSecureURL

src/h_vialib/secure/encryption.py

@@ -0,0 +1,26 @@
+import json
+
+from jose import constants, jwe
+
+
+class Encryption:
+    JWE_ALGORITHM = constants.ALGORITHMS.DIR
+    JWE_ENCRYPTION = constants.ALGORITHMS.A128CBC_HS256
+
+    def __init__(self, secret: bytes):
+        self._secret = secret.ljust(32)[:32]
+
+    def encrypt_dict(self, payload: dict) -> str:
+        """Encrypt a dictionary as a JWE."""
+
+        return jwe.encrypt(
+            json.dumps(payload),
+            self._secret,
+            algorithm=self.JWE_ALGORITHM,
+            encryption=self.JWE_ENCRYPTION,
+        ).decode("utf-8")
+
+    def decrypt_dict(self, payload: str) -> dict:
+        """Decrypts payloads created by `encrypt_dict`."""
+
+        return json.loads(jwe.decrypt(payload, self._secret))

tests/unit/h_vialib/client_test.py

@@ -120,6 +120,17 @@ def test_url_for_allows_you_to_override_options(self, client, content_type):
 
         assert final_url == Any.url.containing_query(override)
 
+    def test_url_for_with_headers(self, client, Encryption):
+        headers = {"some": "header"}
+        Encryption.return_value.encrypt_dict.return_value = "secure headers"
+
+        final_url = client.url_for("http://example.com", headers=headers)
+
+        Encryption.return_value.encrypt_dict.assert_called_once_with(headers)
+        assert final_url == Any.url.containing_query(
+            {"via.secret.headers": "secure headers"}
+        )
+
     @pytest.mark.parametrize("content_type", (None, "pdf", "html"))
     def test_url_for_raises_without_a_service_url(self, content_type):
         client = ViaClient(
@@ -147,7 +158,11 @@ def test_it_fixes_html_urls_with_bare_hostnames(self, client, url, path):
         assert signed_url == Any.url.with_path(path)
 
     @pytest.fixture
-    def client(self):
+    def Encryption(self, patch):
+        return patch("h_vialib.client.Encryption")
+
+    @pytest.fixture
+    def client(self, Encryption):  # pylint:disable=unused-argument
         return ViaClient(
             service_url=self.VIA_URL,
             html_service_url=self.VIAHTML_URL,

tests/unit/h_vialib/secure/encryption_test.py

@@ -0,0 +1,50 @@
+import pytest
+from jose import constants
+
+from h_vialib.secure import Encryption
+
+
+class TestEncryption:
+    def test_encrypt_dict_round_trip(self, encryption):
+        payload_dict = {"some": "data"}
+
+        encrypted = encryption.encrypt_dict(payload_dict)
+
+        assert encryption.decrypt_dict(encrypted) == payload_dict
+
+    def test_encrypt_dict(self, encryption, secret, json, jwe):
+        payload_dict = {"some": "data"}
+
+        encrypted = encryption.encrypt_dict(payload_dict)
+
+        json.dumps.assert_called_with(payload_dict)
+        jwe.encrypt.assert_called_once_with(
+            json.dumps.return_value,
+            secret.ljust(32),
+            algorithm=constants.ALGORITHMS.DIR,
+            encryption=constants.ALGORITHMS.A128CBC_HS256,
+        )
+        assert encrypted == jwe.encrypt.return_value.decode.return_value
+
+    def test_decrypt_dict(self, encryption, secret, jwe, json):
+        plain_text_dict = encryption.decrypt_dict("payload")
+
+        jwe.decrypt.assert_called_once_with("payload", secret.ljust(32))
+        json.loads.assert_called_once_with(jwe.decrypt.return_value)
+        assert plain_text_dict == json.loads.return_value
+
+    @pytest.fixture
+    def secret(self):
+        return b"VERY SECRET"
+
+    @pytest.fixture
+    def encryption(self, secret):
+        return Encryption(secret)
+
+    @pytest.fixture
+    def json(self, patch):
+        return patch("h_vialib.secure.encryption.json")
+
+    @pytest.fixture
+    def jwe(self, patch):
+        return patch("h_vialib.secure.encryption.jwe")