From 93a8be3dcaf389acda6108746a9138cbb117cd32 Mon Sep 17 00:00:00 2001
From: Grigory Frantsuzov <wooferclaw@users.noreply.github.com>
Date: Wed, 13 Sep 2023 01:42:05 +0400
Subject: [PATCH] Added security headers library for fastApi (#198)

* Added security headers library for fastApi

* Fix TOML syntax and add to "all" list.

* Add secure library installation to apiserver image.

* Add test of secure middleware that checks the inserted headers.

* Update test to capture more output.

---------

Co-authored-by: Grigory Frantsuzov <grigoriy.frantsuzov@quantori.com>
Co-authored-by: James Mathews <mail@jmath.read-books.org>
---
 build/apiserver/Dockerfile                    |  1 +
 pyproject.toml.unversioned                    |  6 ++--
 spatialprofilingtoolbox/apiserver/app/main.py |  9 ++++++
 .../unit_tests/expected_headers_example.txt   | 14 +++++++++
 .../unit_tests/test_secure_headers.sh         | 29 +++++++++++++++++++
 5 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 test/apiserver/unit_tests/expected_headers_example.txt
 create mode 100644 test/apiserver/unit_tests/test_secure_headers.sh

diff --git a/build/apiserver/Dockerfile b/build/apiserver/Dockerfile
index f65f03361..6765aa992 100644
--- a/build/apiserver/Dockerfile
+++ b/build/apiserver/Dockerfile
@@ -14,6 +14,7 @@ RUN python -m pip install pyshp==2.2.0
 RUN python -m pip install scikit-learn==1.2.2
 RUN python -m pip install Pillow==9.5.0
 RUN python -m pip install pydantic==2.0.2
+RUN python -m pip install secure==0.3.0
 ARG version
 ARG service_name
 ARG WHEEL_FILENAME
diff --git a/pyproject.toml.unversioned b/pyproject.toml.unversioned
index 9dc413808..ef04864db 100644
--- a/pyproject.toml.unversioned
+++ b/pyproject.toml.unversioned
@@ -39,7 +39,8 @@ apiserver = [
     "pyshp==2.2.0",
     "scikit-learn==1.2.2",
     "Pillow==9.5.0",
-    "pydantic==2.0.2"
+    "pydantic==2.0.2",
+    "secure==0.3.0"
 ]
 cggnn = [
     "cg-gnn"
@@ -85,7 +86,8 @@ all = [
     "pydantic==2.0.2",
     "fastapi==0.100.0",
     "Pillow==9.5.0",
-    "squidpy==1.3.0"
+    "squidpy==1.3.0",
+    "secure==0.3.0"
 ]
 dev = [
     "autopep8",
diff --git a/spatialprofilingtoolbox/apiserver/app/main.py b/spatialprofilingtoolbox/apiserver/app/main.py
index cc3eb8eec..b065955ca 100644
--- a/spatialprofilingtoolbox/apiserver/app/main.py
+++ b/spatialprofilingtoolbox/apiserver/app/main.py
@@ -9,6 +9,8 @@
 from fastapi import Response
 from fastapi.responses import StreamingResponse
 
+import secure
+
 from spatialprofilingtoolbox.ondemand.service_client import OnDemandRequester
 from spatialprofilingtoolbox.db.exchange_data_formats.study import StudyHandle
 from spatialprofilingtoolbox.db.exchange_data_formats.study import StudySummary
@@ -76,6 +78,13 @@ def custom_openapi():
 
 setattr(app, 'openapi', custom_openapi)
 
+secure_headers = secure.Secure()
+
+@app.middleware("http")
+async def set_secure_headers(request, call_next):
+    response = await call_next(request)
+    secure_headers.framework.fastapi(response)
+    return response
 
 @app.get("/")
 async def get_root():
diff --git a/test/apiserver/unit_tests/expected_headers_example.txt b/test/apiserver/unit_tests/expected_headers_example.txt
new file mode 100644
index 000000000..05e2643b1
--- /dev/null
+++ b/test/apiserver/unit_tests/expected_headers_example.txt
@@ -0,0 +1,14 @@
+< HTTP/1.1 200 OK
+< server: uvicorn
+< content-length: 97
+< content-type: application/json
+< strict-transport-security: max-age=63072000; includeSubdomains
+< x-frame-options: SAMEORIGIN
+< x-xss-protection: 0
+< x-content-type-options: nosniff
+< referrer-policy: no-referrer, strict-origin-when-cross-origin
+< cache-control: no-store
+< 
+{ [97 bytes data]
+* Connection #0 to host spt-apiserver-testing left intact
+[{"handle":"Melanoma intralesional IL2","display_name_detail":"Cancer Immunology Research 2022"}]
diff --git a/test/apiserver/unit_tests/test_secure_headers.sh b/test/apiserver/unit_tests/test_secure_headers.sh
new file mode 100644
index 000000000..eba523d8d
--- /dev/null
+++ b/test/apiserver/unit_tests/test_secure_headers.sh
@@ -0,0 +1,29 @@
+
+query=http://spt-apiserver-testing:8080/study-names/
+
+curl -s $query 1>/dev/null 2>/dev/null;
+if [ "$?" -gt 0 ];
+then
+    echo "Error with apiserver query."
+    curl $query
+    exit 1
+fi
+
+curl -s --verbose $query 2>&1 | tail -n +9 | grep -v 'date: ' | tr -d '\b\r' > response.txt
+
+diff unit_tests/expected_headers_example.txt response.txt
+status=$?
+[ $status -eq 0 ] || (echo "API query for headers inspection failed."; )
+
+if [ $status -eq 0 ];
+then
+    echo "Response headers were as expected:"
+    echo
+    cat response.txt
+	rm response.txt
+    exit 0
+else
+	cat response.txt
+	rm response.txt
+    exit 1
+fi