feat: add AuthInterceptor for gRPC authentication handling

Author: nicosammito

main.py

@@ -5,6 +5,8 @@
 
 import grpc
 
+from src.interceptor.auth_interceptor import AuthInterceptor
+
 current_dir = os.path.dirname(os.path.abspath(__file__))
 venv_site_packages = os.path.join(current_dir, ".venv", "lib", "python3.12", "site-packages")
 
@@ -25,7 +27,10 @@
 
 if __name__ == '__main__':
 
-    server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=10))
+    server = grpc.server(
+        concurrent.futures.ThreadPoolExecutor(max_workers=10),
+        interceptors=[AuthInterceptor()]
+    )
     generate_pb2_grpc.add_GenerateServiceServicer_to_server(GenerateService(), server)
     info_pb2_grpc.add_InfoServiceServicer_to_server(InfoService(), server)
 

src/interceptor/auth_interceptor.py

@@ -0,0 +1,36 @@
+import os
+
+import grpc
+import jwt
+
+VALID_TOKEN = os.getenv('SECURITY_TOKEN', 'velorum-internal-communication-secret')
+
+class AuthInterceptor(grpc.ServerInterceptor):
+
+    def intercept_service(self, continuation, handler_call_details):
+        metadata = dict(handler_call_details.invocation_metadata or [])
+        auth_token = metadata.get('authorization', '')
+
+        if not auth_token:
+            return self._abort_handler()
+
+        try:
+            jwt.decode(
+                auth_token,
+                VALID_TOKEN,
+                algorithms=['HS256']
+            )
+        except Exception as e:
+            print(e)
+            return self._abort_handler()
+
+        return continuation(handler_call_details)
+
+    def _abort_handler(self, details="Invalid or missing authentication token"):
+        def abort_handler(request, context):
+            context.abort(
+                code=grpc.StatusCode.UNAUTHENTICATED,
+                details=details
+            )
+
+        return grpc.unary_unary_rpc_method_handler(abort_handler)