added CSP middleware

piccolo_api/middleware/csp.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from functools import wraps
+import typing as t
+
+from starlette.types import ASGIApp, Send, Receive, Scope, Message
+
+
+@dataclass
+class CSPConfig():
+    report_uri: t.Optional[bytes] = None
+
+
+class CSPMiddleware():
+    """
+    Adds Content Security Policy headers to the response.
+    """
+
+    def __init__(self, app: ASGIApp, config: CSPConfig = CSPConfig()):
+        self.app = app
+        self.config = config
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        @wraps(send)
+        async def wrapped_send(message: Message):
+            if message["type"] == "http.response.start":
+                headers = message.get('headers', [])
+                header_value = b"default-src 'self'"
+                if self.config.report_uri:
+                    header_value = (
+                        header_value +
+                        b'; report-uri ' +
+                        self.config.report_uri
+                    )
+                headers.append(
+                    [
+                        b'Content-Security-Policy',
+                        header_value
+                    ]
+                )
+                message['headers'] = headers
+
+            await send(message)
+
+        await self.app(scope, receive, wrapped_send)

test-requirements.txt

@@ -0,0 +1,2 @@
+pytest==5.0.1
+

tests/run_tests.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+# To run all in a folder tests/
+# To run all in a file tests/test_foo.py
+# To run all in a class tests/test_foo.py::TestFoo
+# To run a single test tests/test_foo.py::TestFoo::test_foo
+cd .. && python -m pytest -s $@

tests/test_csp.py

@@ -0,0 +1,53 @@
+from unittest import TestCase
+
+from piccolo_api.middleware.csp import CSPMiddleware, CSPConfig
+from starlette.testclient import TestClient
+
+
+async def app(scope, receive, send):
+    await send({
+        'type': 'http.response.start',
+        'status': 200,
+        'headers': [
+            [b'content-type', b'text/plain'],
+        ],
+    })
+    await send({
+        'type': 'http.response.body',
+        'body': b'Hello, world!',
+    })
+
+
+class TestCSPMiddleware(TestCase):
+
+    def test_headers(self):
+        wrapped_app = CSPMiddleware(app)
+
+        client = TestClient(wrapped_app)
+        response = client.request('GET', '/')
+
+        # Make sure the headers got added:
+        self.assertTrue(
+            'Content-Security-Policy' in response.headers.keys()
+        )
+
+        # Make sure the original headers are still intact:
+        self.assertTrue(
+            'content-type' in response.headers.keys()
+        )
+
+    def test_report_uri(self):
+        wrapped_app = CSPMiddleware(
+            app,
+            config=CSPConfig(
+                report_uri=b'foo.com'
+            )
+        )
+
+        client = TestClient(wrapped_app)
+        response = client.request('GET', '/')
+
+        header = response.headers['Content-Security-Policy']
+        self.assertTrue(
+            'report-uri' in header
+        )