From 8ecf41b82cdc531e82b5d006143995cee008dfda Mon Sep 17 00:00:00 2001 From: Lee Chagolla-Christensen Date: Mon, 27 Oct 2025 11:23:27 -0700 Subject: [PATCH] fix API endpoint, fix global vars bug, use API model --- libs/common/common/models2/api.py | 4 +- libs/common/common/models2/enrichments.py | 18 ++++++ .../activities/enrichment_modules.py | 8 +-- .../file_enrichment/controller.py | 4 +- .../file_enrichment/routes/enrichments.py | 62 +++++++++---------- .../subscriptions/bulk_enrichment.py | 4 +- .../file_enrichment/workflow.py | 6 +- 7 files changed, 57 insertions(+), 49 deletions(-) create mode 100644 libs/common/common/models2/enrichments.py diff --git a/libs/common/common/models2/api.py b/libs/common/common/models2/api.py index 5520735..fcd57bd 100644 --- a/libs/common/common/models2/api.py +++ b/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() diff --git a/libs/common/common/models2/enrichments.py b/libs/common/common/models2/enrichments.py new file mode 100644 index 0000000..03c2c04 --- /dev/null +++ b/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] diff --git a/projects/file_enrichment/file_enrichment/activities/enrichment_modules.py b/projects/file_enrichment/file_enrichment/activities/enrichment_modules.py index 92f70c0..7e45cfa 100644 --- a/projects/file_enrichment/file_enrichment/activities/enrichment_modules.py +++ b/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) diff --git a/projects/file_enrichment/file_enrichment/controller.py b/projects/file_enrichment/file_enrichment/controller.py index d5b0771..122e3a1 100644 --- a/projects/file_enrichment/file_enrichment/controller.py +++ b/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) diff --git a/projects/file_enrichment/file_enrichment/routes/enrichments.py b/projects/file_enrichment/file_enrichment/routes/enrichments.py index d46ed9d..fba55d2 100644 --- a/projects/file_enrichment/file_enrichment/routes/enrichments.py +++ b/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 diff --git a/projects/file_enrichment/file_enrichment/subscriptions/bulk_enrichment.py b/projects/file_enrichment/file_enrichment/subscriptions/bulk_enrichment.py index d8f3e0f..d07d650 100644 --- a/projects/file_enrichment/file_enrichment/subscriptions/bulk_enrichment.py +++ b/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: diff --git a/projects/file_enrichment/file_enrichment/workflow.py b/projects/file_enrichment/file_enrichment/workflow.py index b714805..32c840e 100644 --- a/projects/file_enrichment/file_enrichment/workflow.py +++ b/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