Merge pull request #52 from RobGeada/PassthroughHeaders

Feat: Add optional passthrough headers to custom functions

Author: RobGeada

detectors/built_in/app.py

@@ -1,6 +1,6 @@
 import logging
 
-from fastapi import HTTPException
+from fastapi import HTTPException, Request
 from contextlib import asynccontextmanager
 from base_detector_registry import BaseDetectorRegistry
 from regex_detectors import RegexDetectorRegistry
@@ -29,9 +29,11 @@ async def lifespan(app: FastAPI):
 
 
 @app.post("/api/v1/text/contents", response_model=ContentsAnalysisResponse)
-def detect_content(request: ContentAnalysisHttpRequest):
+def detect_content(request: ContentAnalysisHttpRequest, raw_request: Request):
     logger.info(f"Request for {request.detector_params}")
 
+    headers = dict(raw_request.headers)
+
     detections = []
     for content in request.contents:
         message_detections = []
@@ -43,7 +45,7 @@ def detect_content(request: ContentAnalysisHttpRequest):
                 raise TypeError(f"Detector {detector_kind} is not a valid BaseDetectorRegistry")
             else:
                 try:
-                    message_detections += detector_registry.handle_request(content, request.detector_params)
+                    message_detections += detector_registry.handle_request(content, request.detector_params, headers)
                 except HTTPException as e:
                     raise e
                 except Exception as e:

detectors/built_in/base_detector_registry.py

@@ -8,7 +8,7 @@ def __init__(self):
         self.registry = None
 
     @abstractmethod
-    def handle_request(self, content: str, detector_params: dict) -> List[ContentAnalysisResponse]:
+    def handle_request(self, content: str, detector_params: dict, headers: dict) -> List[ContentAnalysisResponse]:
         pass
 
     def get_registry(self):

detectors/built_in/custom_detectors/custom_detectors.py

@@ -4,3 +4,6 @@ def over_100_characters(text: str) -> bool:
 
 def contains_word(text: str) -> bool:
     return "apple" in text.lower()
+
+def function_that_needs_headers(text: str, headers: dict) -> bool:
+    return headers['magic-key'] != "123"

detectors/built_in/custom_detectors_wrapper.py

@@ -12,10 +12,15 @@
 
 logger = logging.getLogger(__name__)
 
-def custom_func_wrapper(func: Callable, func_name: str, s: str) -> Optional[ContentAnalysisResponse]:
+def custom_func_wrapper(func: Callable, func_name: str, s: str, headers: dict) -> Optional[ContentAnalysisResponse]:
     """Convert a some f(text)->bool into a Detector response"""
+    sig = inspect.signature(func)
     try:
-        result = func(s)
+        if headers is not None:
+            result = func(s, headers)
+        else:
+            result = func(s)
+
     except Exception as e:
         logging.error(f"Error when computing custom detector function {func_name}: {e}")
         raise e
@@ -96,17 +101,19 @@ def __init__(self):
         self.registry = {name: obj for name, obj
                          in inspect.getmembers(custom_detectors, inspect.isfunction)
                          if not name.startswith("_")}
+        self.function_needs_headers = {name: "headers" in inspect.signature(obj).parameters for name, obj in self.registry.items() }
         logger.info(f"Registered the following custom detectors: {self.registry.keys()}")
 
-    def handle_request(self, content: str, detector_params: dict) -> List[ContentAnalysisResponse]:
+    def handle_request(self, content: str, detector_params: dict, headers: dict) -> List[ContentAnalysisResponse]:
         detections = []
         if "custom" in detector_params and isinstance(detector_params["custom"], (list, str)):
             custom_functions = detector_params["custom"]
             custom_functions = [custom_functions] if isinstance(custom_functions, str) else custom_functions
             for custom_function in custom_functions:
                 if self.registry.get(custom_function):
                     try:
-                        result = custom_func_wrapper(self.registry[custom_function], custom_function, content)
+                        func_headers = headers if self.function_needs_headers.get(custom_function) else None
+                        result = custom_func_wrapper(self.registry[custom_function], custom_function, content, func_headers)
                         if result is not None:
                             detections.append(result)
                     except Exception as e:

detectors/built_in/file_type_detectors.py

@@ -176,7 +176,7 @@ def __init__(self):
             "yaml-with-schema:$SCHEMA": is_valid_yaml_schema,
         }
 
-    def handle_request(self, content: str, detector_params: dict) -> List[ContentAnalysisResponse]:
+    def handle_request(self, content: str, detector_params: dict, headers: dict) -> List[ContentAnalysisResponse]:
         detections = []
         if "file_type" in detector_params and isinstance(detector_params["file_type"], (list, str)):
             file_types = detector_params["file_type"]

detectors/built_in/regex_detectors.py

@@ -131,7 +131,7 @@ def __init__(self):
             "$CUSTOM_REGEX": custom_regex_documenter,
         }
 
-    def handle_request(self, content: str, detector_params: dict) -> List[ContentAnalysisResponse]:
+    def handle_request(self, content: str, detector_params: dict, headers: dict) -> List[ContentAnalysisResponse]:
         detections = []
         if "regex" in detector_params and isinstance(detector_params["regex"], (list, str)):
             regexes = detector_params["regex"]

tests/detectors/builtIn/test_custom.py

@@ -12,13 +12,8 @@
     "../../../detectors/built_in/custom_detectors/custom_detectors.py"
 )
 
-SAFE_CODE = """
-def over_100_characters(text: str) -> bool:
-    return len(text)>100
-
-def contains_word(text: str) -> bool:
-    return "apple" in text.lower()
-"""
+with open(CUSTOM_DETECTORS_PATH) as f:
+    SAFE_CODE = f.read()
 
 UNSAFE_CODE = '''
 import os
@@ -80,6 +75,27 @@ def test_custom_detectors_not_match(self, client):
         texts = [d["text"] for d in resp.json()[0]]
         assert msg not in texts
 
+    def test_custom_detectors_need_header(self, client):
+        msg = "What is an banana?"
+        payload = {
+            "contents": [msg],
+            "detector_params": {"custom": ["function_that_needs_headers"]}
+        }
+
+        # shouldn't flag
+        headers = {"magic-key": "123"}
+        resp = client.post("/api/v1/text/contents", json=payload, headers=headers)
+        assert resp.status_code == 200
+        texts = [d["text"] for d in resp.json()[0]]
+        assert msg not in texts
+
+        # should flag
+        headers = {"magic-key": "wrong"}
+        resp = client.post("/api/v1/text/contents", json=payload, headers=headers)
+        assert resp.status_code == 200
+        texts = [d["text"] for d in resp.json()[0]]
+        assert msg in texts
+
     def test_unsafe_code(self, client):
         write_code_to_custom_detectors(UNSAFE_CODE)
         from detectors.built_in.custom_detectors_wrapper import CustomDetectorRegistry