mirror of
https://github.com/SpecterOps/Nemesis
synced 2026-06-08 12:36:42 +00:00
more psycopg to asyncpg pools in file enrichment
- Replace psycopg_pool.ConnectionPool with asyncpg.Pool throughout file_enrichment service - Store asyncpg pool in FastAPI app.state instead of global variable - Update type hints - Fix enrichment success tracking query typo in noseyparker - Improve formatting in dotnet finding summaries (use code formatting for method names)
This commit is contained in:
Lee Chagolla-Christensen
@@ -14,7 +14,6 @@ from file_enrichment.postgres_notifications import postgres_notify_listener
|
||||
from file_enrichment.workflow_recovery import recover_interrupted_workflows
|
||||
from nemesis_dpapi import DpapiManager as NemesisDpapiManager
|
||||
from nemesis_dpapi.eventing import DaprDpapiEventPublisher
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from .debug_utils import setup_debug_signals
|
||||
from .routes.dpapi import dpapi_background_monitor, dpapi_router
|
||||
@@ -34,8 +33,6 @@ max_workflow_execution_time = int(
|
||||
|
||||
logger.info(f"max_workflow_execution_time: {max_workflow_execution_time}", pid=os.getpid())
|
||||
|
||||
pool = ConnectionPool(get_postgres_connection_str(), open=True)
|
||||
|
||||
module_execution_order = []
|
||||
workflow_manager: WorkflowManager = None
|
||||
|
||||
@@ -60,7 +57,6 @@ async def lifespan(app: FastAPI):
|
||||
app.state.event_loop = asyncio.get_running_loop()
|
||||
set_fastapi_loop(asyncio.get_event_loop())
|
||||
|
||||
# Create asyncpg connection pool for WorkflowManager and workflow activities
|
||||
dapr_client = DaprClient()
|
||||
postgres_connection_string = get_postgres_connection_str(dapr_client)
|
||||
|
||||
@@ -69,7 +65,7 @@ async def lifespan(app: FastAPI):
|
||||
min_size=5,
|
||||
max_size=15,
|
||||
)
|
||||
logger.info("AsyncPG pool created", pid=os.getpid())
|
||||
app.state.asyncpg_pool = asyncpg_pool
|
||||
|
||||
# Initialize global DpapiManager for the application lifetime
|
||||
dpapi_manager = NemesisDpapiManager(
|
||||
@@ -107,7 +103,7 @@ async def lifespan(app: FastAPI):
|
||||
logger.info("Started masterkey watcher task", pid=os.getpid())
|
||||
|
||||
# Recover any interrupted workflows before starting normal processing
|
||||
await recover_interrupted_workflows(pool)
|
||||
await recover_interrupted_workflows(asyncpg_pool)
|
||||
|
||||
logger.info(
|
||||
"Workflow runtime initialized successfully",
|
||||
@@ -249,19 +245,19 @@ async def debug_tasks():
|
||||
async def process_file(event: CloudEvent[File]):
|
||||
"""Handler for incoming file events"""
|
||||
global workflow_manager
|
||||
await process_file_event(event.data, workflow_manager, module_execution_order, pool)
|
||||
await process_file_event(event.data, workflow_manager, module_execution_order, app.state.asyncpg_pool)
|
||||
|
||||
|
||||
@dapr_app.subscribe(pubsub="pubsub", topic="dotnet-output")
|
||||
async def process_dotnet_results(event: CloudEvent):
|
||||
"""Handler for incoming .NET processing results from the dotnet_service."""
|
||||
await process_dotnet_event(event.data, pool)
|
||||
await process_dotnet_event(event.data, app.state.asyncpg_pool)
|
||||
|
||||
|
||||
@dapr_app.subscribe(pubsub="pubsub", topic="noseyparker-output")
|
||||
async def process_nosey_parker_results(event: CloudEvent):
|
||||
"""Handler for incoming Nosey Parker scan results"""
|
||||
await process_noseyparker_event(event.data, pool)
|
||||
await process_noseyparker_event(event.data, app.state.asyncpg_pool)
|
||||
|
||||
|
||||
@dapr_app.subscribe(pubsub="pubsub", topic="bulk-enrichment-task")
|
||||
|
||||
@@ -1,16 +1,16 @@
|
||||
# src/workflow/dotnet.py
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import psycopg
|
||||
import asyncpg
|
||||
from common.helpers import sanitize_for_jsonb
|
||||
from common.logger import get_logger
|
||||
from common.models import (
|
||||
DotNetAssemblyAnalysis,
|
||||
EnrichmentResult,
|
||||
File,
|
||||
FileEnriched,
|
||||
FileObject,
|
||||
Finding,
|
||||
FindingCategory,
|
||||
@@ -18,7 +18,6 @@ from common.models import (
|
||||
Transform,
|
||||
)
|
||||
from dapr.clients import DaprClient
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -59,9 +58,9 @@ def create_dotnet_finding_summary(analysis: DotNetAssemblyAnalysis) -> str:
|
||||
summary_content += f"\n#### {category}\n"
|
||||
for method in methods:
|
||||
if hasattr(method, "MethodName"):
|
||||
summary_content += f"* **{method.MethodName}** (Level: {method.FilterLevel})\n"
|
||||
summary_content += f"* `{method.MethodName}` (Level: {method.FilterLevel})\n"
|
||||
else:
|
||||
summary_content += f"* {method}\n"
|
||||
summary_content += f"* `{method}`\n"
|
||||
summary_content += "\n"
|
||||
return summary_content
|
||||
return ""
|
||||
@@ -78,10 +77,10 @@ def create_dotnet_finding_summary(analysis: DotNetAssemblyAnalysis) -> str:
|
||||
|
||||
async def store_dotnet_results(
|
||||
object_id: str,
|
||||
decompilation_object_id: str = None,
|
||||
analysis: DotNetAssemblyAnalysis = None,
|
||||
pool: ConnectionPool = None,
|
||||
file_enriched=None,
|
||||
decompilation_object_id: str | None = None,
|
||||
analysis: DotNetAssemblyAnalysis | None = None,
|
||||
pool: asyncpg.Pool | None = None,
|
||||
file_enriched: FileEnriched | None = None,
|
||||
):
|
||||
"""
|
||||
Store DotNet analysis results in the database, including creating findings and transforms.
|
||||
@@ -90,23 +89,22 @@ async def store_dotnet_results(
|
||||
object_id (str): The object ID of the file that was analyzed
|
||||
decompilation_object_id (str, optional): Object ID of the decompiled source ZIP
|
||||
analysis (DotNetAssemblyAnalysis, optional): Assembly analysis results
|
||||
pool (ConnectionPool, optional): Database connection pool
|
||||
pool (asyncpg.Pool, optional): Database connection pool
|
||||
file_enriched: The FileEnriched object for the original file
|
||||
"""
|
||||
try:
|
||||
# Update workflow success status
|
||||
try:
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE workflows
|
||||
SET enrichments_success = array_append(enrichments_success, %s)
|
||||
WHERE object_id = %s
|
||||
""",
|
||||
("dotnet_service", object_id),
|
||||
)
|
||||
conn.commit()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE workflows
|
||||
SET enrichments_success = array_append(enrichments_success, $1)
|
||||
WHERE object_id = $2
|
||||
""",
|
||||
"dotnet_service",
|
||||
object_id,
|
||||
)
|
||||
except Exception as db_error:
|
||||
logger.error(f"Failed to update dotnet_service enrichment success in database: {str(db_error)}")
|
||||
|
||||
@@ -203,71 +201,63 @@ async def store_dotnet_results(
|
||||
# Add findings to enrichment result
|
||||
enrichment_result.findings = findings_list
|
||||
|
||||
def store_in_db():
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json")))
|
||||
cur.execute(
|
||||
# Store in database
|
||||
async with pool.acquire() as conn:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json")))
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES ($1, $2, $3)
|
||||
""",
|
||||
object_id,
|
||||
"dotnet_service",
|
||||
results_escaped,
|
||||
)
|
||||
|
||||
# Store any transforms
|
||||
if enrichment_result.transforms:
|
||||
for transform in enrichment_result.transforms:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES (%s, %s, %s)
|
||||
INSERT INTO transforms (object_id, type, transform_object_id, metadata)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
(object_id, "dotnet_service", results_escaped),
|
||||
object_id,
|
||||
transform.type,
|
||||
transform.object_id,
|
||||
json.dumps(transform.metadata) if transform.metadata else None,
|
||||
)
|
||||
|
||||
# Store any transforms
|
||||
if enrichment_result.transforms:
|
||||
for transform in enrichment_result.transforms:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO transforms (object_id, type, transform_object_id, metadata)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
object_id,
|
||||
transform.type,
|
||||
transform.object_id,
|
||||
json.dumps(transform.metadata) if transform.metadata else None,
|
||||
),
|
||||
)
|
||||
# Store any findings
|
||||
for finding in findings_list:
|
||||
# Convert each FileObject to a JSON string
|
||||
data_as_strings = []
|
||||
for obj in finding.data:
|
||||
# Convert the model to a dict first
|
||||
if hasattr(obj, "model_dump"):
|
||||
obj_dict = obj.model_dump()
|
||||
else:
|
||||
obj_dict = obj
|
||||
sanitized_obj = sanitize_for_jsonb(obj_dict)
|
||||
data_as_strings.append(json.dumps(sanitized_obj))
|
||||
|
||||
# Store any findings
|
||||
for finding in findings_list:
|
||||
# Convert each FileObject to a JSON string
|
||||
data_as_strings = []
|
||||
for obj in finding.data:
|
||||
# Convert the model to a dict first
|
||||
if hasattr(obj, "model_dump"):
|
||||
obj_dict = obj.model_dump()
|
||||
else:
|
||||
obj_dict = obj
|
||||
sanitized_obj = sanitize_for_jsonb(obj_dict)
|
||||
data_as_strings.append(json.dumps(sanitized_obj))
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(sanitize_for_jsonb(finding.raw_data)),
|
||||
json.dumps(data_as_strings), # Store as array of JSON strings
|
||||
),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# Run database operations in thread
|
||||
await asyncio.to_thread(store_in_db)
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
""",
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(sanitize_for_jsonb(finding.raw_data)),
|
||||
json.dumps(data_as_strings), # Store as array of JSON strings
|
||||
)
|
||||
|
||||
logger.info("Successfully stored DotNet results", object_id=object_id, has_findings=len(findings_list) > 0)
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
@@ -7,11 +6,10 @@ import uuid
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import psycopg
|
||||
import asyncpg
|
||||
from common.helpers import sanitize_for_jsonb
|
||||
from common.logger import get_logger
|
||||
from common.models import EnrichmentResult, FileObject, Finding, FindingCategory, FindingOrigin, MatchInfo, ScanStats
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -162,7 +160,7 @@ async def store_noseyparker_results(
|
||||
object_id: str,
|
||||
matches: list[MatchInfo],
|
||||
scan_stats: ScanStats,
|
||||
pool: ConnectionPool,
|
||||
pool: asyncpg.Pool,
|
||||
):
|
||||
"""
|
||||
Store Nosey Parker results in the database, including creating findings.
|
||||
@@ -171,21 +169,20 @@ async def store_noseyparker_results(
|
||||
object_id (str): The object ID of the file that was scanned
|
||||
matches (List[MatchInfo]): List of match information from Nosey Parker
|
||||
scan_stats (dict, optional): Statistics about the scan
|
||||
pool (ConnectionPool): Database connection pool
|
||||
pool (asyncpg.Pool): Database connection pool
|
||||
"""
|
||||
try:
|
||||
try:
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE workflows
|
||||
SET enrichments_success = array_append(enrichments_failure, %s)
|
||||
WHERE object_id = %s
|
||||
""",
|
||||
("noseyparker", object_id),
|
||||
)
|
||||
conn.commit()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE workflows
|
||||
SET enrichments_success = array_append(enrichments_failure, $1)
|
||||
WHERE object_id = $2
|
||||
""",
|
||||
"noseyparker",
|
||||
object_id,
|
||||
)
|
||||
except Exception as db_error:
|
||||
logger.error(f"Failed to update noseyparker enrichment success in database: {str(db_error)}")
|
||||
|
||||
@@ -233,55 +230,49 @@ async def store_noseyparker_results(
|
||||
# Add findings to enrichment result
|
||||
enrichment_result.findings = findings_list
|
||||
|
||||
def store_in_db():
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json")))
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES (%s, %s, %s)
|
||||
""",
|
||||
(object_id, "noseyparker", results_escaped),
|
||||
)
|
||||
# Store in database
|
||||
async with pool.acquire() as conn:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json")))
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES ($1, $2, $3)
|
||||
""",
|
||||
object_id,
|
||||
"noseyparker",
|
||||
results_escaped,
|
||||
)
|
||||
|
||||
# Store any findings
|
||||
for finding in findings_list:
|
||||
# Convert each FileObject to a JSON string
|
||||
data_as_strings = []
|
||||
for obj in finding.data:
|
||||
# Convert the model to a dict first
|
||||
if hasattr(obj, "model_dump"):
|
||||
obj_dict = obj.model_dump()
|
||||
else:
|
||||
obj_dict = obj
|
||||
sanitized_obj = sanitize_for_jsonb(obj_dict)
|
||||
data_as_strings.append(json.dumps(sanitized_obj))
|
||||
# Store any findings
|
||||
for finding in findings_list:
|
||||
# Convert each FileObject to a JSON string
|
||||
data_as_strings = []
|
||||
for obj in finding.data:
|
||||
# Convert the model to a dict first
|
||||
if hasattr(obj, "model_dump"):
|
||||
obj_dict = obj.model_dump()
|
||||
else:
|
||||
obj_dict = obj
|
||||
sanitized_obj = sanitize_for_jsonb(obj_dict)
|
||||
data_as_strings.append(json.dumps(sanitized_obj))
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(sanitize_for_jsonb(finding.raw_data)),
|
||||
json.dumps(data_as_strings), # Store as array of JSON strings
|
||||
),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# Run database operations in thread
|
||||
await asyncio.to_thread(store_in_db)
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
""",
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(sanitize_for_jsonb(finding.raw_data)),
|
||||
json.dumps(data_as_strings), # Store as array of JSON strings
|
||||
)
|
||||
|
||||
logger.info("Successfully stored NoseyParker results", object_id=object_id, match_count=len(matches))
|
||||
|
||||
|
||||
@@ -5,9 +5,10 @@ import json
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import asyncpg
|
||||
import common.helpers as helpers
|
||||
from common.logger import get_logger
|
||||
from fastapi import APIRouter, Body, HTTPException, Path
|
||||
from fastapi import APIRouter, Body, HTTPException, Path, Request
|
||||
from file_enrichment.workflow import wf_runtime
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -60,13 +61,13 @@ async def list_enrichments():
|
||||
@router.post("/enrichments/{enrichment_name}")
|
||||
async def run_enrichment(
|
||||
enrichment_name: str = Path(..., description="Name of the enrichment module to run"),
|
||||
request: EnrichmentRequest = Body(..., description="The enrichment request containing the object ID"),
|
||||
enrichment_request: EnrichmentRequest = Body(..., description="The enrichment request containing the object ID"),
|
||||
request: Request = None,
|
||||
):
|
||||
"""Run a specific enrichment module directly."""
|
||||
# Import pool from controller module to avoid circular imports
|
||||
from file_enrichment import controller
|
||||
|
||||
try:
|
||||
# Get asyncpg pool from app state
|
||||
asyncpg_pool: asyncpg.Pool = request.app.state.asyncpg_pool
|
||||
# Check if module
|
||||
if not wf_runtime or not wf_runtime.modules:
|
||||
raise HTTPException(status_code=503, detail="Workflow runtime or modules not initialized")
|
||||
@@ -78,80 +79,71 @@ async def run_enrichment(
|
||||
module = wf_runtime.modules[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, request.object_id)
|
||||
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": request.object_id,
|
||||
"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, request.object_id)
|
||||
result = await asyncio.to_thread(module.process, enrichment_request.object_id)
|
||||
|
||||
if result:
|
||||
# Store enrichment result in database
|
||||
def store_results():
|
||||
with controller.pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(helpers.sanitize_for_jsonb(result.model_dump(mode="json")))
|
||||
cur.execute(
|
||||
async with asyncpg_pool.acquire() as conn:
|
||||
# Store main enrichment result
|
||||
results_escaped = json.dumps(helpers.sanitize_for_jsonb(result.model_dump(mode="json")))
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES ($1, $2, $3)
|
||||
""",
|
||||
enrichment_request.object_id,
|
||||
enrichment_name,
|
||||
results_escaped,
|
||||
)
|
||||
|
||||
# Store any transforms
|
||||
if result.transforms:
|
||||
for transform in result.transforms:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO enrichments (object_id, module_name, result_data)
|
||||
VALUES (%s, %s, %s)
|
||||
INSERT INTO transforms (object_id, type, transform_object_id, metadata)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
(request.object_id, enrichment_name, results_escaped),
|
||||
enrichment_request.object_id,
|
||||
transform.type,
|
||||
transform.object_id,
|
||||
json.dumps(transform.metadata) if transform.metadata else None,
|
||||
)
|
||||
|
||||
# Store any transforms
|
||||
if result.transforms:
|
||||
for transform in result.transforms:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO transforms (object_id, type, transform_object_id, metadata)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
request.object_id,
|
||||
transform.type,
|
||||
transform.object_id,
|
||||
json.dumps(transform.metadata) if transform.metadata else None,
|
||||
),
|
||||
)
|
||||
|
||||
# Store any findings
|
||||
if result.findings:
|
||||
for finding in result.findings:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
request.object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(finding.raw_data),
|
||||
json.dumps([obj.model_dump_json() for obj in finding.data]),
|
||||
),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# Run database operations in thread
|
||||
await asyncio.to_thread(store_results)
|
||||
# Store any findings
|
||||
if result.findings:
|
||||
for finding in result.findings:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO findings (
|
||||
finding_name, category, severity, object_id,
|
||||
origin_type, origin_name, raw_data, data
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
""",
|
||||
finding.finding_name,
|
||||
finding.category,
|
||||
finding.severity,
|
||||
enrichment_request.object_id,
|
||||
finding.origin_type,
|
||||
finding.origin_name,
|
||||
json.dumps(finding.raw_data),
|
||||
json.dumps([obj.model_dump_json() for obj in finding.data]),
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Completed enrichment with module '{enrichment_name}'",
|
||||
"object_id": request.object_id,
|
||||
"object_id": enrichment_request.object_id,
|
||||
"instance_id": str(uuid.uuid4()), # Generate a unique instance ID
|
||||
}
|
||||
|
||||
@@ -162,7 +154,7 @@ async def run_enrichment(
|
||||
e,
|
||||
message="Error running enrichment module",
|
||||
enrichment_name=enrichment_name,
|
||||
object_id=request.object_id,
|
||||
object_id=enrichment_request.object_id,
|
||||
pid=os.getpid(),
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") from e
|
||||
|
||||
@@ -3,16 +3,16 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import asyncpg
|
||||
from common.logger import get_logger
|
||||
from common.models import DotNetOutput
|
||||
from common.state_helpers import get_file_enriched_async
|
||||
from file_enrichment.dotnet import store_dotnet_results
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def process_dotnet_event(raw_data, pool: ConnectionPool):
|
||||
async def process_dotnet_event(raw_data, pool: asyncpg.Pool):
|
||||
"""Process incoming .NET processing results from the dotnet_service"""
|
||||
try:
|
||||
logger.debug(f"Received DotNet output event: {raw_data}", pid=os.getpid())
|
||||
@@ -36,10 +36,10 @@ async def process_dotnet_event(raw_data, pool: ConnectionPool):
|
||||
logger.debug(f"Processing dotnet results for object {object_id}", pid=os.getpid())
|
||||
|
||||
# Get the file enriched data for creating transforms
|
||||
file_enriched = None
|
||||
try:
|
||||
file_enriched = await get_file_enriched_async(object_id)
|
||||
except Exception as e:
|
||||
file_enriched = None
|
||||
logger.warning(f"Could not get file_enriched for {object_id}: {e}", pid=os.getpid())
|
||||
|
||||
# Store the results in the database using our helper function
|
||||
|
||||
@@ -1,17 +1,16 @@
|
||||
"""Handler for file subscription events."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
import asyncpg
|
||||
from common.logger import get_logger
|
||||
from common.models import File
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def save_file_message(file: File, pool: ConnectionPool):
|
||||
async def save_file_message(file: File, pool: asyncpg.Pool):
|
||||
"""Save the file message to the database for recovery purposes"""
|
||||
try:
|
||||
# Only save files that are not nested (originating files)
|
||||
@@ -24,53 +23,47 @@ async def save_file_message(file: File, pool: ConnectionPool):
|
||||
)
|
||||
return
|
||||
|
||||
def save_to_db():
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
query = """
|
||||
INSERT INTO files (
|
||||
object_id, agent_id, source, project, timestamp, expiration,
|
||||
path, originating_object_id, originating_container_id, nesting_level,
|
||||
file_creation_time, file_access_time, file_modification_time
|
||||
) VALUES (
|
||||
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
||||
) ON CONFLICT (object_id) DO UPDATE SET
|
||||
agent_id = EXCLUDED.agent_id,
|
||||
source = EXCLUDED.source,
|
||||
project = EXCLUDED.project,
|
||||
timestamp = EXCLUDED.timestamp,
|
||||
expiration = EXCLUDED.expiration,
|
||||
path = EXCLUDED.path,
|
||||
originating_object_id = EXCLUDED.originating_object_id,
|
||||
originating_container_id = EXCLUDED.originating_container_id,
|
||||
nesting_level = EXCLUDED.nesting_level,
|
||||
file_creation_time = EXCLUDED.file_creation_time,
|
||||
file_access_time = EXCLUDED.file_access_time,
|
||||
file_modification_time = EXCLUDED.file_modification_time,
|
||||
updated_at = CURRENT_TIMESTAMP;
|
||||
"""
|
||||
query = """
|
||||
INSERT INTO files (
|
||||
object_id, agent_id, source, project, timestamp, expiration,
|
||||
path, originating_object_id, originating_container_id, nesting_level,
|
||||
file_creation_time, file_access_time, file_modification_time
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13
|
||||
) ON CONFLICT (object_id) DO UPDATE SET
|
||||
agent_id = EXCLUDED.agent_id,
|
||||
source = EXCLUDED.source,
|
||||
project = EXCLUDED.project,
|
||||
timestamp = EXCLUDED.timestamp,
|
||||
expiration = EXCLUDED.expiration,
|
||||
path = EXCLUDED.path,
|
||||
originating_object_id = EXCLUDED.originating_object_id,
|
||||
originating_container_id = EXCLUDED.originating_container_id,
|
||||
nesting_level = EXCLUDED.nesting_level,
|
||||
file_creation_time = EXCLUDED.file_creation_time,
|
||||
file_access_time = EXCLUDED.file_access_time,
|
||||
file_modification_time = EXCLUDED.file_modification_time,
|
||||
updated_at = CURRENT_TIMESTAMP;
|
||||
"""
|
||||
|
||||
cur.execute(
|
||||
query,
|
||||
(
|
||||
file.object_id,
|
||||
file.agent_id,
|
||||
file.source,
|
||||
file.project,
|
||||
file.timestamp,
|
||||
file.expiration,
|
||||
file.path,
|
||||
file.originating_object_id,
|
||||
getattr(file, "originating_container_id", None),
|
||||
file.nesting_level,
|
||||
datetime.fromisoformat(file.creation_time) if file.creation_time else None,
|
||||
datetime.fromisoformat(file.access_time) if file.access_time else None,
|
||||
datetime.fromisoformat(file.modification_time) if file.modification_time else None,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
query,
|
||||
file.object_id,
|
||||
file.agent_id,
|
||||
file.source,
|
||||
file.project,
|
||||
file.timestamp,
|
||||
file.expiration,
|
||||
file.path,
|
||||
file.originating_object_id,
|
||||
getattr(file, "originating_container_id", None),
|
||||
file.nesting_level,
|
||||
datetime.fromisoformat(file.creation_time) if file.creation_time else None,
|
||||
datetime.fromisoformat(file.access_time) if file.access_time else None,
|
||||
datetime.fromisoformat(file.modification_time) if file.modification_time else None,
|
||||
)
|
||||
|
||||
await asyncio.to_thread(save_to_db)
|
||||
logger.debug("Successfully saved file message to database", object_id=file.object_id, pid=os.getpid())
|
||||
|
||||
except Exception as e:
|
||||
@@ -78,7 +71,7 @@ async def save_file_message(file: File, pool: ConnectionPool):
|
||||
raise
|
||||
|
||||
|
||||
async def process_file_event(file: File, workflow_manager, module_execution_order: list, pool: ConnectionPool):
|
||||
async def process_file_event(file: File, workflow_manager, module_execution_order: list, pool: asyncpg.Pool):
|
||||
"""Process incoming file events"""
|
||||
try:
|
||||
# Save the file message to database first for recovery purposes
|
||||
|
||||
@@ -3,15 +3,15 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import asyncpg
|
||||
from common.logger import get_logger
|
||||
from common.models import NoseyParkerOutput
|
||||
from file_enrichment.noseyparker import store_noseyparker_results
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def process_noseyparker_event(raw_data, pool: ConnectionPool):
|
||||
async def process_noseyparker_event(raw_data, pool: asyncpg.Pool):
|
||||
"""Process incoming Nosey Parker scan results"""
|
||||
try:
|
||||
# Extract the raw data
|
||||
|
||||
@@ -27,68 +27,18 @@ async def recover_interrupted_workflows(pool) -> None:
|
||||
|
||||
logger.info("Starting workflow recovery process...", pid=os.getpid())
|
||||
|
||||
def get_and_delete_running_workflows():
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Atomic DELETE with RETURNING - only one worker will get the interrupted workflows
|
||||
cur.execute("""
|
||||
DELETE FROM workflows
|
||||
WHERE status = 'RUNNING'
|
||||
RETURNING object_id
|
||||
""")
|
||||
running_ids = [row[0] for row in cur.fetchall()]
|
||||
conn.commit()
|
||||
# Get interrupted workflows atomically using asyncpg
|
||||
async with pool.acquire() as conn:
|
||||
# Atomic DELETE with RETURNING - only one worker will get the interrupted workflows
|
||||
running_ids = await conn.fetch("""
|
||||
DELETE FROM workflows
|
||||
WHERE status = 'RUNNING'
|
||||
RETURNING object_id
|
||||
""")
|
||||
running_object_ids = [row['object_id'] for row in running_ids]
|
||||
|
||||
if running_ids:
|
||||
logger.info(f"Atomically claimed {len(running_ids)} interrupted workflows", pid=os.getpid())
|
||||
|
||||
return running_ids
|
||||
|
||||
def get_file_data_and_cleanup(object_ids):
|
||||
recovered_files = []
|
||||
with pool.connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
for object_id in object_ids:
|
||||
# Get file data for reconstruction
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT object_id, agent_id, source, project, timestamp, expiration,
|
||||
path, originating_object_id, originating_container_id, nesting_level,
|
||||
file_creation_time, file_access_time, file_modification_time
|
||||
FROM files WHERE object_id = %s
|
||||
""",
|
||||
(object_id,),
|
||||
)
|
||||
|
||||
row = cur.fetchone()
|
||||
if row:
|
||||
# Convert database row to File-compatible dict
|
||||
file_data = {
|
||||
"object_id": str(row[0]),
|
||||
"agent_id": row[1],
|
||||
"source": row[2],
|
||||
"project": row[3],
|
||||
"timestamp": row[4],
|
||||
"expiration": row[5],
|
||||
"path": row[6],
|
||||
"originating_object_id": str(row[7]) if row[7] else None,
|
||||
"originating_container_id": str(row[8]) if row[8] else None,
|
||||
"nesting_level": row[9],
|
||||
"creation_time": row[10].isoformat() if row[10] else None,
|
||||
"access_time": row[11].isoformat() if row[11] else None,
|
||||
"modification_time": row[12].isoformat() if row[12] else None,
|
||||
}
|
||||
recovered_files.append(file_data)
|
||||
logger.debug("Recovered file data for workflow", object_id=object_id, pid=os.getpid())
|
||||
else:
|
||||
logger.warning("No file data found for workflow", object_id=object_id, pid=os.getpid())
|
||||
|
||||
conn.commit()
|
||||
|
||||
return recovered_files
|
||||
|
||||
# Get interrupted workflows
|
||||
running_object_ids = await asyncio.to_thread(get_and_delete_running_workflows)
|
||||
if running_object_ids:
|
||||
logger.info(f"Atomically claimed {len(running_object_ids)} interrupted workflows", pid=os.getpid())
|
||||
|
||||
if not running_object_ids:
|
||||
logger.info("No interrupted workflows found", pid=os.getpid())
|
||||
@@ -97,7 +47,41 @@ async def recover_interrupted_workflows(pool) -> None:
|
||||
logger.info(f"Found {len(running_object_ids)} interrupted workflows to recover", pid=os.getpid())
|
||||
|
||||
# Get file data and clean up partial results
|
||||
recovered_files = await asyncio.to_thread(get_file_data_and_cleanup, running_object_ids)
|
||||
recovered_files = []
|
||||
async with pool.acquire() as conn:
|
||||
for object_id in running_object_ids:
|
||||
# Get file data for reconstruction
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT object_id, agent_id, source, project, timestamp, expiration,
|
||||
path, originating_object_id, originating_container_id, nesting_level,
|
||||
file_creation_time, file_access_time, file_modification_time
|
||||
FROM files WHERE object_id = $1
|
||||
""",
|
||||
object_id,
|
||||
)
|
||||
|
||||
if row:
|
||||
# Convert database row to File-compatible dict
|
||||
file_data = {
|
||||
"object_id": str(row['object_id']),
|
||||
"agent_id": row['agent_id'],
|
||||
"source": row['source'],
|
||||
"project": row['project'],
|
||||
"timestamp": row['timestamp'],
|
||||
"expiration": row['expiration'],
|
||||
"path": row['path'],
|
||||
"originating_object_id": str(row['originating_object_id']) if row['originating_object_id'] else None,
|
||||
"originating_container_id": str(row['originating_container_id']) if row['originating_container_id'] else None,
|
||||
"nesting_level": row['nesting_level'],
|
||||
"creation_time": row['file_creation_time'].isoformat() if row['file_creation_time'] else None,
|
||||
"access_time": row['file_access_time'].isoformat() if row['file_access_time'] else None,
|
||||
"modification_time": row['file_modification_time'].isoformat() if row['file_modification_time'] else None,
|
||||
}
|
||||
recovered_files.append(file_data)
|
||||
logger.debug("Recovered file data for workflow", object_id=object_id, pid=os.getpid())
|
||||
else:
|
||||
logger.warning("No file data found for workflow", object_id=object_id, pid=os.getpid())
|
||||
|
||||
if not recovered_files:
|
||||
logger.warning("No file data found for interrupted workflows", pid=os.getpid())
|
||||
|
||||
Reference in New Issue
Block a user