mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
from integrations.siem.models import (
|
|
AdaptiveQueryInput,
|
|
ESQLQueryInput,
|
|
KeywordSearchInput,
|
|
QueryOutput,
|
|
SAMPLE_COUNT,
|
|
SAMPLE_THRESHOLD,
|
|
SPLQueryInput,
|
|
)
|
|
from integrations.siem.registry import get_default_agg_fields
|
|
|
|
|
|
def build_query_output(
|
|
input_data: AdaptiveQueryInput | KeywordSearchInput,
|
|
result: BackendQueryResult,
|
|
*,
|
|
index_distribution: Optional[dict[str, int]] = None,
|
|
) -> QueryOutput:
|
|
status = "records" if result.total_hits <= SAMPLE_THRESHOLD else "summary"
|
|
record_limit = SAMPLE_THRESHOLD if status == "records" else SAMPLE_COUNT
|
|
|
|
if isinstance(input_data, AdaptiveQueryInput):
|
|
explicit_fields = list(input_data.filters.keys()) + result.aggregation_fields
|
|
else:
|
|
explicit_fields = result.aggregation_fields
|
|
|
|
fields_to_project = list(dict.fromkeys(
|
|
[input_data.time_field, *explicit_fields, *get_default_agg_fields(result.index_name)]
|
|
))
|
|
|
|
records = [
|
|
project_record(r, fields_to_project)
|
|
for r in result.raw_records[:record_limit]
|
|
]
|
|
|
|
return QueryOutput(
|
|
backend=result.backend,
|
|
index_name=result.index_name,
|
|
status=status,
|
|
total_hits=result.total_hits,
|
|
returned_records=len(records),
|
|
truncated=result.total_hits > len(records),
|
|
message=f"Matched {result.total_hits} events in {result.index_name} ({result.backend}). "
|
|
+ ("Returning projected records." if status == "records" else "Returning statistics and samples."),
|
|
index_distribution=index_distribution,
|
|
statistics=result.statistics,
|
|
records=records,
|
|
)
|
|
|
|
|
|
def project_record(record: dict[str, Any], fields: list[str]) -> dict[str, Any]:
|
|
projected: dict[str, Any] = {}
|
|
for field_path in fields:
|
|
value = _get_nested(record, field_path)
|
|
if value is not _MISSING:
|
|
projected[field_path] = value
|
|
if "_index" in record:
|
|
projected["_index"] = record["_index"]
|
|
return projected
|
|
|
|
|
|
_MISSING = object()
|
|
|
|
|
|
def _get_nested(record: dict[str, Any], field_path: str) -> Any:
|
|
if field_path in record:
|
|
return record[field_path]
|
|
current: Any = record
|
|
for segment in field_path.split("."):
|
|
if not isinstance(current, dict) or segment not in current:
|
|
return _MISSING
|
|
current = current[segment]
|
|
return current
|
|
|
|
|
|
def build_raw_query_output(
|
|
input_data: SPLQueryInput | ESQLQueryInput,
|
|
result: BackendQueryResult,
|
|
*,
|
|
limit: int = 100,
|
|
) -> QueryOutput:
|
|
records = result.raw_records
|
|
return QueryOutput(
|
|
backend=result.backend,
|
|
index_name=result.index_name,
|
|
status="records",
|
|
total_hits=result.total_hits,
|
|
returned_records=len(records),
|
|
truncated=len(records) >= limit,
|
|
message=f"Executed raw {result.backend} query against {result.index_name}. Returned {len(records)} records.",
|
|
statistics=[],
|
|
records=records,
|
|
)
|