fix API endpoint, fix global vars bug, use API model

libs/common/common/models2/api.py

@@ -4,7 +4,7 @@ from datetime import UTC, datetime
 from typing import Annotated, Literal
 from uuid import UUID
 
-from pydantic import BaseModel, BeforeValidator, Field, field_validator, field_serializer
+from pydantic import BaseModel, BeforeValidator, Field, field_serializer, field_validator
 
 logger = logging.getLogger(__name__)
 
@@ -130,7 +130,7 @@ class FileMetadata(BaseModel):
         }
     }
 
-    @field_serializer('timestamp', 'expiration', when_used='unless-none')
+    @field_serializer("timestamp", "expiration", when_used="unless-none")
     def serialize_datetime(self, dt: datetime, _info):
         return dt.isoformat()
 

libs/common/common/models2/enrichments.py

@@ -0,0 +1,18 @@
+"""API models for the /enrichments route."""
+
+from pydantic import BaseModel
+
+
+class EnrichmentRequest(BaseModel):
+    object_id: str
+
+
+class EnrichmentResponse(BaseModel):
+    status: str
+    message: str
+    object_id: str
+    instance_id: str
+
+
+class ModulesListResponse(BaseModel):
+    modules: list[str]

projects/file_enrichment/file_enrichment/activities/enrichment_modules.py

@@ -4,6 +4,7 @@ import json
 import os
 
 import common.helpers as helpers
+import file_enrichment.global_vars as global_vars
 from common.logger import get_logger
 from common.models import EnrichmentResult
 from common.workflows.setup import workflow_activity
@@ -15,7 +16,6 @@ from ..tracing import get_tracer
 logger = get_logger(__name__)
 
 # Global module map - will be set during initialization
-global_module_map = {}
 
 
 @workflow_activity
@@ -97,11 +97,11 @@ def determine_modules_to_process(object_id: str, temp_file_path: str, execution_
     modules_to_process = []
 
     for module_name in execution_order:
-        if module_name not in global_module_map:
+        if module_name not in global_vars.global_module_map:
             logger.warning("Module not found", module_name=module_name)
             continue
 
-        module = global_module_map[module_name]
+        module = global_vars.global_module_map[module_name]
         try:
             should_process = module.should_process(object_id, temp_file_path)
 
@@ -116,7 +116,7 @@ def determine_modules_to_process(object_id: str, temp_file_path: str, execution_
 
 async def execute_enrichment_module(object_id: str, temp_file_path: str, module_name: str) -> EnrichmentResult | None:
     """Second pass: process a single module and return its result."""
-    module = global_module_map[module_name]
+    module = global_vars.global_module_map[module_name]
     logger.debug("Starting module processing", module_name=module_name)
 
     # Check if the module's process method returns a coroutine (async)

projects/file_enrichment/file_enrichment/controller.py

@@ -191,6 +191,6 @@ async def process_nosey_parker_results(event: CloudEvent[NoseyParkerOutput]):
 async def process_bulk_enrichment_task(event: CloudEvent[BulkEnrichmentTask]):
     """Handler for individual bulk enrichment tasks"""
     global workflow_manager
-    from .workflow import global_module_map
+    import file_enrichment.global_vars as global_vars
 
-    await process_bulk_enrichment_event(event.data, workflow_manager, global_module_map)
+    await process_bulk_enrichment_event(event.data, workflow_manager, global_vars.global_module_map)

projects/file_enrichment/file_enrichment/routes/enrichments.py

@@ -9,9 +9,8 @@ from typing import TYPE_CHECKING
 import common.helpers as helpers
 import file_enrichment.global_vars as global_vars
 from common.logger import get_logger
-from fastapi import APIRouter, Body, HTTPException, Path, Request
-from file_enrichment.global_vars import global_module_map
-from pydantic import BaseModel
+from common.models2.enrichments import EnrichmentRequest, EnrichmentResponse, ModulesListResponse
+from fastapi import APIRouter, Body, HTTPException, Path
 
 if TYPE_CHECKING:
     pass
@@ -21,15 +20,11 @@ logger = get_logger(__name__)
 router = APIRouter(tags=["enrichments"])
 
 
-class EnrichmentRequest(BaseModel):
-    object_id: str
-
-
-@router.get("/llm_enrichments")
-async def list_enabled_llm_enrichments():
+@router.get("/llm_enrichments", response_model=ModulesListResponse)
+async def list_enabled_llm_enrichments() -> ModulesListResponse:
     """List the enabled LLM enrichments based on environment variables."""
     try:
-        if not global_module_map:
+        if not global_vars.global_module_map:
             raise HTTPException(status_code=503, detail="Modules not initialized")
 
         llm_enrichments = []
@@ -40,51 +35,50 @@ async def list_enabled_llm_enrichments():
         if os.getenv("RIGGING_GENERATOR_TRIAGE"):
             llm_enrichments.append("finding_triage")
 
-        return {"modules": llm_enrichments}
+        return ModulesListResponse(modules=llm_enrichments)
 
     except Exception as e:
         logger.exception(e, message="Error listing enabled LLM enrichment modules", pid=os.getpid())
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") from e
 
 
-@router.get("/enrichments")
-async def list_enrichments():
+@router.get("/enrichments", response_model=ModulesListResponse)
+async def list_enrichments() -> ModulesListResponse:
     """List all available enrichment modules."""
     try:
-        if not global_module_map:
+        if not global_vars.global_module_map:
             raise HTTPException(status_code=503, detail="Modules not initialized")
 
-        module_names = list(global_module_map.keys())
-        return {"modules": module_names}
+        module_names = list(global_vars.global_module_map.keys())
+        return ModulesListResponse(modules=module_names)
 
     except Exception as e:
         logger.exception(e, message="Error listing enrichment modules", pid=os.getpid())
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") from e
 
 
-@router.post("/enrichments/{enrichment_name}")
+@router.post("/enrichments/{enrichment_name}", response_model=EnrichmentResponse)
 async def run_enrichment(
     enrichment_name: str = Path(..., description="Name of the enrichment module to run"),
     enrichment_request: EnrichmentRequest = Body(..., description="The enrichment request containing the object ID"),
-    request: Request = None,
-):
+) -> EnrichmentResponse:
     """Run a specific enrichment module directly."""
     try:
-        if enrichment_name not in global_module_map:
+        if enrichment_name not in global_vars.global_module_map:
             raise HTTPException(status_code=404, detail=f"Enrichment module '{enrichment_name}' not found")
 
         # Get the module
-        module = global_module_map[enrichment_name]
+        module = global_vars.global_module_map[enrichment_name]
 
         # Check if we should process this file - run in thread since it might use sync operations
         should_process = await asyncio.to_thread(module.should_process, enrichment_request.object_id)
         if not should_process:
-            return {
-                "status": "skipped",
-                "message": f"Module {enrichment_name} decided to skip processing",
-                "object_id": enrichment_request.object_id,
-                "instance_id": "",
-            }
+            return EnrichmentResponse(
+                status="skipped",
+                message=f"Module {enrichment_name} decided to skip processing",
+                object_id=enrichment_request.object_id,
+                instance_id="",
+            )
 
         # Process the file in a separate thread to avoid event loop conflicts
         result = await asyncio.to_thread(module.process, enrichment_request.object_id)
@@ -135,15 +129,15 @@ async def run_enrichment(
                             finding.origin_type,
                             finding.origin_name,
                             json.dumps(finding.raw_data),
-                            json.dumps([obj.model_dump_json() for obj in finding.data]),
+                            json.dumps([obj.model_dump() for obj in finding.data]),
                         )
 
-        return {
-            "status": "success",
-            "message": f"Completed enrichment with module '{enrichment_name}'",
-            "object_id": enrichment_request.object_id,
-            "instance_id": str(uuid.uuid4()),  # Generate a unique instance ID
-        }
+        return EnrichmentResponse(
+            status="success",
+            message=f"Completed enrichment with module '{enrichment_name}'",
+            object_id=enrichment_request.object_id,
+            instance_id=str(uuid.uuid4()),  # Generate a unique instance ID
+        )
 
     except HTTPException:
         raise

projects/file_enrichment/file_enrichment/subscriptions/bulk_enrichment.py

@@ -12,9 +12,7 @@ async def process_bulk_enrichment_event(task: BulkEnrichmentTask, workflow_manag
         enrichment_name = task.enrichment_name
         object_id = task.object_id
 
-        logger.debug(
-            "Received bulk enrichment task", enrichment_name=enrichment_name, object_id=object_id
-        )
+        logger.debug("Received bulk enrichment task", enrichment_name=enrichment_name, object_id=object_id)
 
         # Check if module exists
         if not global_module_map:

projects/file_enrichment/file_enrichment/workflow.py

@@ -16,7 +16,6 @@ from .activities import (
     publish_findings_alerts,
     run_enrichment_modules,
 )
-from .activities.enrichment_modules import global_module_map
 
 logger = get_logger(__name__)
 
@@ -240,9 +239,8 @@ async def initialize_workflow_runtime(dpapi_manager: DpapiManager):
     module_loader = ModuleLoader()
     await module_loader.load_modules()
     # Update the global_module_map in the enrichment_modules activity
-    from .activities import enrichment_modules
 
-    enrichment_modules.global_module_map = module_loader.modules
+    global_vars.global_module_map = module_loader.modules
 
     asyncio_loop = asyncio.get_running_loop()
 
@@ -288,7 +286,7 @@ def reload_yara_rules():
     """Reloads all disk/state yara rules."""
 
     logger.debug("workflow/workflow.py reloading Yara rules")
-    global_module_map["yara"].rule_manager.load_rules()
+    global_vars.global_module_map["yara"].rule_manager.load_rules()
 
 
 # endregion