Added healthcheck and signature

main.py

@@ -9,7 +9,8 @@ from scripts.db.psql.create_default_tables import create_default_psql_dependenci
 from scripts.services import route
 from scripts.services.tag_list_services import tag_service_router
 from scripts.utils.security_utils.decorators import CookieAuthentication
-
+from scripts.utils.security_utils.signature import SignatureVerificationMiddleware
+from scripts.config import Service
 auth = CookieAuthentication()
 
 
@@ -26,6 +27,9 @@ class FastAPIConfig:
 app = FastAPI(**FastAPIConfig().__dict__)
 
 
+if Service.verify_signature in [True, 'True', 'true']:
+    app.add_middleware(SignatureVerificationMiddleware)
+
 @app.get("/api/oee_services/health_check")
 async def ping():
     return {"status": 200}

scripts/config/__init__.py

@@ -41,6 +41,9 @@ class Service:
     PORT = config.getint("SERVICE", "port")
     WORKERS = config.getint("SERVICE", "workers")
     secure_cookie = config.get("SERVICE", "secure_cookie")
+    verify_signature = os.environ.get("VERIFY_SIGNATURE", False)
+    protected_hosts = os.environ.get("PROTECTED_HOSTS", "").split(",")
+
 
 
 class DBConf:

scripts/constants/__init__.py

@@ -57,6 +57,7 @@ class Secrets:
     token = '8674cd1d-2578-4a62-8ab7-d3ee5f9a'
     issuer = "ilens"
     alg = "RS256"
+    signature_key = 'kliLensKLiLensKL'
 
 
 class CommonKeys:

scripts/utils/security_utils/signature.py

+
+from scripts.config import Service
+from scripts.constants import Secrets
+import json
+import logging
+
+import jwt
+from fastapi import Request
+from jwt.exceptions import (
+    InvalidSignatureError,
+    ExpiredSignatureError,
+    MissingRequiredClaimError,
+    DecodeError
+)
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+ENFORCE_DOMAIN_WILDCARD = "Domain wildcard patterns must be like '*.example.com'."
+
+protect_hosts = Service.protected_hosts
+if not protect_hosts:
+    protect_hosts = ["*"]
+
+for _pattern in protect_hosts:
+    assert "*" not in _pattern[1:], ENFORCE_DOMAIN_WILDCARD
+    if _pattern.startswith("*") and _pattern != "*":
+        assert _pattern.startswith("*."), ENFORCE_DOMAIN_WILDCARD
+
+
+class SignatureVerificationMiddleware(BaseHTTPMiddleware):
+
+    async def set_body(self, request: Request):
+
+        async def verify_signature():
+            receive_ = await request.receive()
+            signature = bytearray()
+            signature.extend(receive_.get('body'))
+            while receive_['more_body']:
+                receive_ = await request.receive()
+                signature.extend(receive_['body'])
+
+            signature = bytes(signature)
+            try:
+                signature = jwt.decode(signature.decode(), Secrets.signature_key, algorithms=['HS256'])
+            except (
+                    InvalidSignatureError,
+                    ExpiredSignatureError,
+                    MissingRequiredClaimError,
+                    DecodeError
+            ) as inv_exp:
+                logging.error(inv_exp)
+                signature = {}
+            signature = json.dumps(signature).encode()
+            return {"type": receive_["type"], "body": signature, "more_body": False}
+
+        if request.headers.get("Content-Type") == "application/json":
+            host = request.headers.get("host", "").split(":")[0]
+            is_protected_host = False
+            for pattern in protect_hosts:
+                if host == pattern or (
+                        pattern.startswith("*") and host.endswith(pattern[1:])
+                ):
+                    is_protected_host = True
+                    break
+            if is_protected_host:
+                return Request(request.scope, verify_signature, request._send)
+            else:
+                return request
+
+    async def dispatch(self, request, call_next):
+        request = await self.set_body(request)
+        response = await call_next(request)
+        return response