External/agentic-soc-platform
1
0
Fork 0
mirror of https://github.com/FunnyWolf/agentic-soc-platform.git synced 2026-08-22 13:12:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

tmp

Browse Source
This commit is contained in:
rookit
2026-04-23 10:10:08 +08:00
parent c6fd22413b
commit e73fdf5a17
10 changed files with 762 additions and 758 deletions
Download Patch File Download Diff File Expand all files Collapse all files
DATA/Agent_SIEM/system_prompt.md
+22 -19
View File
@@ -44,12 +44,12 @@ This tool supports a step-by-step refinement approach:
- Request statistics on additional fields to drill deeper
4. **Final retrieval**: Once you've identified the specific logs you need
- The tool automatically returns full records when result volume is small
- Or use the statistics from "sample" and "summary" responses to guide your analysis
- The tool returns projected records when result volume is small
- Or use the statistics from `sample` and `summary` responses to guide your analysis
**Key benefit:** The tool automatically adjusts its response format:
- Returns all records when there are few results (easy analysis)
- Returns projected records when there are few results (easy analysis)
- Returns statistics + sample records for medium volumes (pattern identification)
- Returns statistics only for large volumes (efficient insights)
@@ -65,13 +65,13 @@ Execute keyword-based full-text search across SIEM backends with intelligent res
- Supports searching by IP addresses, hostnames, usernames, or any arbitrary string
- Automatically applies the same adaptive response strategy as execute_adaptive_query
- Control result volume by expanding or shrinking the time range and by adding or removing keywords
- If results are too large, narrow the time range or add more precise keywords until the response reaches `full`
- If results are too large, narrow the time range or add more precise keywords until the response reaches `records`
- If results are too small or empty, expand the time range or remove restrictive keywords to recover relevant logs
- The preferred end state is `full`, because it returns the complete original raw logs for the matched events
- The preferred end state is `records`, because it returns the projected event records for the matched events
**Key benefit:** The tool automatically adjusts its response format:
- Returns all records when there are few results (easy analysis)
- Returns projected records when there are few results (easy analysis)
- Returns statistics + sample records for medium volumes (pattern identification)
- Returns statistics only for large volumes (efficient insights)
- When searching across all indices, provides distribution metrics showing hit count per index
@@ -82,7 +82,7 @@ Execute keyword-based full-text search across SIEM backends with intelligent res
2. **Assess the returned status**:
- If status is `summary`, there are too many hits to inspect directly
- If status is `sample`, use the samples and statistics to refine further
- If status is `full`, you already have the complete raw logs and can stop refining
- If status is `records`, you already have the projected event records and can stop refining
3. **Reduce hit volume when needed**:
- Shrink the time range around the suspected activity window
- Add another keyword to the list so all keywords must match
@@ -91,26 +91,29 @@ Execute keyword-based full-text search across SIEM backends with intelligent res
- Expand the time range if the event may have happened earlier or later
- Remove one restrictive keyword from the AND list
- Fall back from a specific index to cross-index search if the source is uncertain
5. **Aim for `full`**: Keep refining until the response reaches `full`, then use the returned records as the complete original log set for that query
5. **Aim for `records`**: Keep refining until the response reaches `records`, then use the returned records as the working event set for that query
**Aggregation fields:** When an index_name is provided, the tool automatically returns statistics for default
aggregation fields defined for that index. This helps identify patterns and distributions.
**Response includes:**
- `status`: Response type ("full", "sample", or "summary")
- `status`: Response type (`records`, `sample`, or `summary`)
- `index_name`: The index/source represented by the result
- `total_hits`: Total number of matching events
- `returned_records`: Number of records actually included in the response
- `truncated`: Whether the tool omitted events or fields to keep the payload LLM-safe
- `index_distribution`: Shows count of results per index (when searching across indices)
- `statistics`: Top values for aggregation fields
- `records`: Sample or full records depending on volume
- `records`: Projected records depending on volume
- `backend`: Which backend returned the results (ELK or Splunk)
## Query Execution Strategy
1. **Receive request**: Get query parameters (index, filters, time range, fields)
2. **Execute query**: Run the appropriate tool with specified parameters
3. **Return results**: Return raw query results in structured format
4. **If data exceeds practical limits**: Suggest refinements such as narrower time ranges or more precise keywords until `full` mode is reached
3. **Return results**: Return structured query results in a machine-readable format
4. **If data exceeds practical limits**: Suggest refinements such as narrower time ranges or more precise keywords until `records` mode is reached
5. **Never perform independent analysis**: Only return the data requested
## Handling Large Log Volumes
@@ -122,7 +125,7 @@ When querying returns excessive data:
- Narrow the time range
- Add more specific keywords or convert a single keyword into an AND keyword list
- Focus on specific event outcomes or behaviors
- Continue refining until `keyword_search` returns `full`, which contains the complete raw logs
- Continue refining until `keyword_search` returns `records`, which contains the projected event records
3. **Apply intelligent compression** (only when needed to maintain efficiency):
- Remove non-essential fields while preserving investigative value
- Group highly repetitive events with occurrence counts and time ranges
@@ -193,7 +196,7 @@ keyword_search(
)
```
→ If status is `summary` or `sample`, there are still too many logs to retrieve as a complete raw set
→ If status is `summary` or `sample`, there are still too many logs to retrieve as a projected record set
**Step 2: Narrow by adding another keyword and shrinking the time window**
@@ -207,7 +210,7 @@ keyword_search(
→ Use AND semantics plus a smaller window to reduce the hit count
**Step 3: Keep refining until `full`**
**Step 3: Keep refining until `records`**
```
keyword_search(
@@ -218,7 +221,7 @@ keyword_search(
)
```
→ When status becomes `full`, use the returned records as the complete original raw logs for final analysis
→ When status becomes `records`, use the returned records as the working event set for final analysis
### Example 5: Investigating Security Events (Progressive Approach)
@@ -262,7 +265,7 @@ execute_adaptive_query(
)
```
→ Get full records for final analysis
→ Get projected records for final analysis
## Important Notes
@@ -270,12 +273,12 @@ execute_adaptive_query(
- If no time range is given, default to a recent window such as 5, 15, or 60 minutes based on query urgency
- The progressive query approach helps you narrow down large datasets efficiently
- For `keyword_search`, prefer iterative refinement by shrinking or expanding the time range and adding or removing keywords
- In keyword-based investigations, the practical target is `full` mode so you can retrieve the complete original raw logs
- In keyword-based investigations, the practical target is `records` mode so you can retrieve the projected event records
- Use explore_schema(target_index="index_name") to discover specific field names when needed
## Output Guidance
- When data volume is manageable: Return complete records in structured format for parent agent analysis
- When data volume is manageable: Return projected records in structured format for parent agent analysis
- When data volume is excessive: Suggest query refinement rather than aggressive compression
- Always provide query context and statistics alongside results
- Use clear structured formats (JSON, tables) to facilitate parent agent processing
MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py
+3 -3
View File
@@ -220,9 +220,9 @@ if __name__ == "__main__":
django.setup()
# 单独测试某条告警
module = Module()
module.debug_message_id = "1776753170359-0"
module.run()
# module = Module()
# module.debug_message_id = "1776753170359-0"
# module.run()
# 批量测试最早的100条告警
module = Module()
PLUGINS/MCP/llmfunc.py
+8 -7
View File
@@ -1,18 +1,18 @@
import json
from datetime import datetime
from typing import Annotated, Optional
from typing import Annotated, Optional, Union, List
from pydantic import Field
from Lib.playbookloader import PlaybookLoader
from PLUGINS.Redis.redis_stream_api import RedisStreamAPI
from PLUGINS.SIEM.models import AdaptiveQueryInput, KeywordSearchInput, SchemaExplorerInput
from PLUGINS.SIEM.models import AdaptiveQueryInput, KeywordSearchInput, SchemaExplorerInput, KeywordSearchOutput, IndexInfo, SchemaIndexSummary
from PLUGINS.SIEM.tools import SIEMToolKit
from PLUGINS.SIRP.nocolymodel import Group, Condition, Operator
from PLUGINS.SIRP.sirpapi import Alert, Artifact, Case, Enrichment, Knowledge, Playbook, Ticket
from PLUGINS.SIRP.sirpextramodel import PlaybookType, KnowledgeSource, PlaybookJobStatus, KnowledgeAction
from PLUGINS.SIRP.sirpcoremodel import TicketStatus, TicketType, ArtifactType, ArtifactRole, ArtifactReputationScore, Severity, AttackStage, Confidence, \
AlertStatus, CaseStatus, CaseVerdict, EnrichmentModel, TicketModel, ArtifactModel
from PLUGINS.SIRP.sirpextramodel import PlaybookType, KnowledgeSource, PlaybookJobStatus, KnowledgeAction
def _dump_models_for_ai(models, limit: int) -> list[dict]:
@@ -485,11 +485,11 @@ def search_knowledge(
def siem_explore_schema(
target_index: Annotated[Optional[str], Field(
description="Target SIEM index to inspect; omit to list all available indices (目标 SIEM 索引名称,不填则列出所有可用索引)")] = None
) -> Annotated[str, Field(description="Schema exploration result as JSON string (索引 Schema 探查结果 JSON 字符串)")]:
) -> Annotated[Union[IndexInfo, List[SchemaIndexSummary]], Field(description="Schema exploration result(索引 Schema 探查结果)")]:
"""Explore available SIEM indices or inspect one index schema. (探查可用的 SIEM 索引列表或指定索引的 Schema)"""
input_data = SchemaExplorerInput(target_index=target_index)
result = SIEMToolKit.explore_schema(input_data)
return json.dumps(result, ensure_ascii=False)
return result
def siem_keyword_search(
@@ -499,7 +499,7 @@ def siem_keyword_search(
time_field: Annotated[str, Field(description="Time field used for range filtering (用于时间范围过滤的字段名)")] = "@timestamp",
index_name: Annotated[
Optional[str], Field(description="Target SIEM index or source; None means all indices (目标 SIEM 索引或数据源,None 表示全部索引)")] = None
) -> Annotated[list[str], Field(description="Search hits as JSON strings (命中的事件列表,每条为 JSON 字符串)")]:
) -> Annotated[list[KeywordSearchOutput], Field(description="Search hits as JSON strings (命中的事件列表,每条为 JSON 字符串)")]:
"""Search SIEM events by keyword and time range. (按关键词和时间范围搜索 SIEM 事件)"""
input_data = KeywordSearchInput(
keyword=keyword,
@@ -509,7 +509,8 @@ def siem_keyword_search(
index_name=index_name
)
results = SIEMToolKit.keyword_search(input_data)
return [item.model_dump_json() for item in results]
# return [item.model_dump_json() for item in results]
return results
def siem_adaptive_query(
PLUGINS/SIEM/models.py
+115 -127
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
from datetime import datetime
from typing import List, Dict, Any, Optional, Union
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field, field_validator
@@ -8,138 +10,140 @@ SAMPLE_THRESHOLD = 100
SAMPLE_COUNT = 5
# --- Input Models ---
class SchemaIndexSummary(BaseModel):
name: str = Field(..., description="Registered SIEM index/source name")
backend: Literal["ELK", "Splunk"] = Field(..., description="Backend that owns this index")
description: str = Field(..., description="Human-readable description of the index")
default_aggregation_fields: List[str] = Field(
default_factory=list,
description="Registry key fields used as default aggregation fields for this index",
)
class SchemaFieldInfo(BaseModel):
name: str = Field(..., description="Field name")
type: str = Field(..., description="Field type declared in the SIEM registry")
description: str = Field(..., description="Human-readable field description")
is_key_field: bool = Field(
default=False,
description="Whether the field is marked as a key field in the registry",
)
class SchemaExplorerInput(BaseModel):
target_index: Optional[str] = Field(
default=None,
description=(
"Target index to explore. "
"If None: returns a list of all available indices with descriptions (list of dicts with 'name' and 'description'). "
"If provided: returns detailed field metadata for that specific index (list of field schemas with 'name', 'type', 'description', etc.)"
)
"If None: returns summaries for all registered indices. "
"If provided: returns field metadata for that specific index."
),
)
class AdaptiveQueryInput(BaseModel):
index_name: str = Field(
...,
description="Target SIEM index/source name. Examples: 'logs-security', 'main', 'logs-endpoint'"
description="Target SIEM index/source name. Examples: 'logs-security', 'main', 'logs-endpoint'",
)
time_field: str = Field(
default="@timestamp",
description=(
"The field to apply time range filter on. "
"Commonly used fields: '@timestamp', 'event.created', '_time'. "
"Must be a Date/DateTime type in your SIEM."
)
"Field used for time range filtering. "
"The field must exist in the target index and be queryable as a timestamp field."
),
)
time_range_start: str = Field(
...,
description="Start time in UTC ISO8601 format. Format: 'YYYY-MM-DDTHH:MM:SSZ'. Example: '2026-02-04T06:00:00Z'"
description="Start time in UTC ISO8601 format: YYYY-MM-DDTHH:MM:SSZ",
)
time_range_end: str = Field(
...,
description="End time in UTC ISO8601 format. Format: 'YYYY-MM-DDTHH:MM:SSZ'. Example: '2026-02-04T07:00:00Z'"
description="End time in UTC ISO8601 format: YYYY-MM-DDTHH:MM:SSZ",
)
filters: Dict[str, Union[str, List[str]]] = Field(
default_factory=dict,
description=(
"Key-value pairs for exact matching filters (term/exact match, not full-text search). "
"Supports both single values (exact match) and lists (OR logic within list). "
"Examples: "
"{'event.outcome': 'success', 'source.ip': '45.33.22.11'} OR "
"{'event.outcome': ['success', 'failed'], 'source.ip': '45.33.22.11'}"
)
"Exact-match filters. "
"String values mean single exact match; list values mean OR semantics within that field."
),
)
aggregation_fields: List[str] = Field(
default_factory=list,
description=(
"Fields to get top-N statistics for. "
"If empty, uses backend-specific default key fields. "
"Example: ['event.outcome', 'source.ip', 'process.name']"
)
"Fields used for top-N aggregation statistics. "
"If empty, the tool uses the registry key fields for the target index."
),
)
@field_validator('time_range_start', 'time_range_end')
@field_validator("time_range_start", "time_range_end")
@classmethod
def validate_utc_format(cls, v):
def validate_utc_format(cls, value: str) -> str:
try:
if not v.endswith("Z"):
if not value.endswith("Z"):
raise ValueError("Time must end with 'Z' to indicate UTC.")
datetime.strptime(v, "%Y-%m-%dT%H:%M:%SZ")
except ValueError:
raise ValueError("Invalid format. Must be UTC ISO8601: YYYY-MM-DDTHH:MM:SSZ")
return v
datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ")
except ValueError as exc:
raise ValueError("Invalid format. Must be UTC ISO8601: YYYY-MM-DDTHH:MM:SSZ") from exc
return value
class KeywordSearchInput(BaseModel):
keyword: Union[str, List[str]] = Field(
...,
description=(
"Search keyword or a list of keywords. "
"A single string performs a standard full-text search. "
"A list performs an AND search, meaning every keyword in the list must match. "
"Keywords can be IP addresses, hostnames, usernames, or arbitrary strings."
)
"Search keyword or keyword list. "
"A list uses AND semantics, so every keyword in the list must match."
),
)
time_range_start: str = Field(
...,
description="Start time in UTC ISO8601 format. Format: 'YYYY-MM-DDTHH:MM:SSZ'. Example: '2026-02-04T06:00:00Z'"
description="Start time in UTC ISO8601 format: YYYY-MM-DDTHH:MM:SSZ",
)
time_range_end: str = Field(
...,
description="End time in UTC ISO8601 format. Format: 'YYYY-MM-DDTHH:MM:SSZ'. Example: '2026-02-04T07:00:00Z'"
description="End time in UTC ISO8601 format: YYYY-MM-DDTHH:MM:SSZ",
)
time_field: str = Field(
default="@timestamp",
description=(
"The field to apply time range filter on. "
"Commonly used fields: '@timestamp', 'event.created', '_time'. "
"Must be a Date/DateTime type in your SIEM."
)
"Field used for time range filtering. "
"The field must exist in the target index and be queryable as a timestamp field."
),
)
index_name: Optional[str] = Field(
default=None,
description=(
"Target SIEM index/source name. "
"If None or empty: searches across all indices. "
"If provided: searches only in specified index. "
"Examples: 'logs-security', 'main', 'logs-endpoint'"
)
"If omitted, the tool first discovers hit indices across the registered backends."
),
)
@field_validator('time_range_start', 'time_range_end')
@field_validator("time_range_start", "time_range_end")
@classmethod
def validate_utc_format(cls, v):
def validate_utc_format(cls, value: str) -> str:
try:
if not v.endswith("Z"):
if not value.endswith("Z"):
raise ValueError("Time must end with 'Z' to indicate UTC.")
datetime.strptime(v, "%Y-%m-%dT%H:%M:%SZ")
except ValueError:
raise ValueError("Invalid format. Must be UTC ISO8601: YYYY-MM-DDTHH:MM:SSZ")
return v
datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ")
except ValueError as exc:
raise ValueError("Invalid format. Must be UTC ISO8601: YYYY-MM-DDTHH:MM:SSZ") from exc
return value
@field_validator('keyword')
@field_validator("keyword")
@classmethod
def validate_keyword(cls, v):
if isinstance(v, str):
keyword = v.strip()
def validate_keyword(cls, value: Union[str, List[str]]) -> Union[str, List[str]]:
if isinstance(value, str):
keyword = value.strip()
if not keyword:
raise ValueError("keyword must not be empty")
return keyword
if isinstance(v, list):
if not v:
if isinstance(value, list):
if not value:
raise ValueError("keyword list must not be empty")
normalized_keywords = []
for item in v:
for item in value:
if not isinstance(item, str):
raise ValueError("keyword list must contain only strings")
keyword = item.strip()
@@ -151,100 +155,84 @@ class KeywordSearchInput(BaseModel):
raise ValueError("keyword must be a string or a list of strings")
# --- Output Models ---
class FieldStat(BaseModel):
field_name: str = Field(
...,
description="Name of the field for which statistics are computed"
)
field_name: str = Field(..., description="Name of the field for which statistics are computed")
top_values: Dict[Union[str, int], int] = Field(
...,
description="Top-N value distribution for the field (key: value, int: count)"
description="Top-N value distribution for the field (value -> count)",
)
class AdaptiveQueryOutput(BaseModel):
status: str = Field(
backend: Literal["ELK", "Splunk"] = Field(..., description="Backend that executed the query")
index_name: str = Field(..., description="Index/source queried by the tool")
status: Literal["records", "sample", "summary"] = Field(
...,
description=(
"Response type indicator based on result volume. "
f"Possible values: 'full' (complete logs, < {SAMPLE_THRESHOLD} results), "
f"'sample' (statistics + sample records, {SAMPLE_THRESHOLD}-{SUMMARY_THRESHOLD} results), "
f"'summary' (statistics only, > {SUMMARY_THRESHOLD} results)"
)
f"'records' returns up to {SAMPLE_THRESHOLD} projected records, "
f"'sample' returns statistics plus up to {SAMPLE_COUNT} projected records, "
f"'summary' returns statistics only."
),
)
total_hits: int = Field(
total_hits: int = Field(..., description="Total number of matching records in the SIEM backend")
returned_records: int = Field(..., description="Number of records included in the response payload")
truncated: bool = Field(
...,
description="Total number of matching records in the SIEM backend"
)
message: str = Field(
...,
description="Human-readable status message describing the response"
description="Whether the tool omitted matching events or record fields to keep the payload LLM-safe",
)
message: str = Field(..., description="Human-readable status message describing the response")
statistics: List[FieldStat] = Field(
...,
description=(
"Top-N value distribution for each aggregation field. "
"Each FieldStat contains field_name and top_values (dict mapping values to their counts)"
)
default_factory=list,
description="Top-N value distribution for each aggregation field",
)
records: List[Dict[str, Any]] = Field(
...,
default_factory=list,
description=(
"Actual log records returned based on status: "
"'full' status returns all records up to SAMPLE_THRESHOLD; "
"'sample' status returns first 3 representative records; "
"'summary' status returns empty list"
)
"Projected log records. "
"These records may omit non-essential fields to control response size."
),
)
class KeywordSearchOutput(BaseModel):
status: str = Field(
backend: Literal["ELK", "Splunk"] = Field(..., description="Backend that executed the search")
index_name: str = Field(..., description="Index/source represented by this result")
status: Literal["records", "sample", "summary"] = Field(
...,
description=(
"Response type indicator based on result volume. "
f"Possible values: 'full' (complete logs, < {SAMPLE_THRESHOLD} results), "
f"'sample' (statistics + sample records, {SAMPLE_THRESHOLD}-{SUMMARY_THRESHOLD} results), "
f"'summary' (statistics only, > {SUMMARY_THRESHOLD} results)"
)
f"'records' returns up to {SAMPLE_THRESHOLD} projected records, "
f"'sample' returns statistics plus up to {SAMPLE_COUNT} projected records, "
f"'summary' returns statistics only."
),
)
total_hits: int = Field(
total_hits: int = Field(..., description="Total number of matching records for this result set")
returned_records: int = Field(..., description="Number of records included in the response payload")
truncated: bool = Field(
...,
description="Total number of matching records across all indices"
)
message: str = Field(
...,
description="Human-readable status message describing the response"
description="Whether the tool omitted matching events or record fields to keep the payload LLM-safe",
)
message: str = Field(..., description="Human-readable status message describing the response")
index_distribution: Dict[str, int] = Field(
...,
description=(
"Distribution of hits across indices. "
"Key: index name, Value: number of hits in that index. "
"When searching a specific index, this will contain only one entry. "
"When searching all indices (*), this shows which indices contain matching data"
)
default_factory=dict,
description="Distribution of hits across indices seen by this search result",
)
statistics: List[FieldStat] = Field(
default_factory=list,
description=(
"Top-N value distribution for each aggregation field. "
"When searching across all indices without specifying aggregation_fields, this may be empty or contain only common fields. "
"When searching a specific index, this contains statistics for the specified or default aggregation fields"
)
description="Top-N value distribution for each aggregation field",
)
records: List[Dict[str, Any]] = Field(
...,
default_factory=list,
description=(
"Actual log records returned based on status. "
"Each record includes '_index' field to indicate its source index. "
"'full' status returns all records up to SAMPLE_THRESHOLD; "
"'sample' status returns first 3 representative records; "
"'summary' status returns empty list"
)
)
backend: str = Field(
default="",
description="Backend type: 'ELK' or 'Splunk'. Empty when searching a specific index."
"Projected log records. "
"These records may omit non-essential fields to control response size."
),
)
class IndexInfo(BaseModel):
name: str
backend: Literal["ELK", "Splunk"]
description: str
fields: List[SchemaFieldInfo]
PLUGINS/SIEM/query_backends.py
+386
View File
@@ -0,0 +1,386 @@
from __future__ import annotations
import json
import re
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from functools import lru_cache
from typing import Any, Literal
from splunklib.results import JSONResultsReader
from PLUGINS.ELK.client import ELKClient
from PLUGINS.SIEM.models import (
AdaptiveQueryInput,
FieldStat,
KeywordSearchInput,
SAMPLE_COUNT,
SAMPLE_THRESHOLD,
)
from PLUGINS.SIEM.registry import get_default_agg_fields
from PLUGINS.Splunk.client import SplunkClient
@dataclass(slots=True)
class BackendQueryResult:
backend: Literal["ELK", "Splunk"]
index_name: str
total_hits: int
aggregation_fields: list[str]
statistics: list[FieldStat]
raw_records: list[dict[str, Any]]
index_distribution: dict[str, int] = field(default_factory=dict)
def normalize_keywords(keyword_input: str | list[str]) -> list[str]:
if isinstance(keyword_input, str):
return [keyword_input]
return keyword_input
def parse_time_range(time_range_start: str, time_range_end: str) -> tuple[float, float]:
utc_format = "%Y-%m-%dT%H:%M:%SZ"
try:
start = datetime.strptime(time_range_start, utc_format).replace(tzinfo=timezone.utc)
end = datetime.strptime(time_range_end, utc_format).replace(tzinfo=timezone.utc)
except ValueError as exc:
raise ValueError("Invalid UTC format.") from exc
return start.timestamp(), end.timestamp()
def _extract_elk_records(hits: list[dict[str, Any]], include_index: bool = False) -> list[dict[str, Any]]:
records: list[dict[str, Any]] = []
for hit in hits:
record = hit["_source"].copy() if include_index else hit["_source"]
if include_index:
record["_index"] = hit["_index"]
records.append(record)
return records
def _extract_elk_stats(response: dict[str, Any], agg_fields: list[str]) -> list[FieldStat]:
stats_output: list[FieldStat] = []
aggregations = response.get("aggregations", {})
for field in agg_fields:
agg_key = f"{field}.keyword" if f"{field}.keyword" in aggregations else field
if agg_key not in aggregations:
continue
buckets = aggregations[agg_key].get("buckets", [])
if buckets:
stats_output.append(
FieldStat(field_name=field, top_values={bucket["key"]: bucket["doc_count"] for bucket in buckets})
)
return stats_output
def _build_time_range_clause(time_field: str, time_range_start: str, time_range_end: str) -> dict[str, Any]:
return {
"range": {
time_field: {
"gte": time_range_start,
"lt": time_range_end,
}
}
}
def _build_elk_keyword_clauses(keyword_input: str | list[str]) -> list[dict[str, Any]]:
return [
{"multi_match": {"query": keyword, "type": "best_fields", "fuzziness": "AUTO"}}
for keyword in normalize_keywords(keyword_input)
]
def _format_splunk_keyword(keyword: str) -> str:
if re.fullmatch(r"[A-Za-z0-9._:@/\\-]+", keyword):
return keyword
escaped_keyword = keyword.replace("\\", "\\\\").replace('"', '\\"')
return f'"{escaped_keyword}"'
def _build_splunk_keyword_clause(keyword_input: str | list[str]) -> str:
return " AND ".join(_format_splunk_keyword(keyword) for keyword in normalize_keywords(keyword_input))
def _clean_splunk_record(log: dict[str, Any]) -> dict[str, Any]:
clean_record: dict[str, Any] = {}
for key, value in log.items():
if not key.startswith("_") and key not in ["splunk_server", "host", "source", "sourcetype"]:
clean_record[key] = value
if "_time" in log:
clean_record["@timestamp"] = log["_time"]
if "_raw" in log:
try:
raw_parsed = json.loads(log["_raw"])
except (json.JSONDecodeError, TypeError):
raw_parsed = None
if isinstance(raw_parsed, dict):
for key, value in raw_parsed.items():
if key not in clean_record:
clean_record[key] = value
return clean_record
def _fetch_splunk_records(job, count: int) -> list[dict[str, Any]]:
records: list[dict[str, Any]] = []
results = job.results(count=count, output_mode="json")
for result in results:
payload = json.loads(result)
for log in payload.get("results", []):
records.append(_clean_splunk_record(log))
return records
def _fetch_splunk_top_stats(service, search_query: str, start_time: float, end_time: float, agg_fields: list[str]) -> list[FieldStat]:
stats_output: list[FieldStat] = []
for field in agg_fields:
stats_query = f"{search_query} | top limit={SAMPLE_COUNT} {field}"
oneshot = service.jobs.oneshot(stats_query, earliest_time=start_time, latest_time=end_time, output_mode="json")
reader = JSONResultsReader(oneshot)
top_values: dict[str | int, int] = {}
for item in reader:
if isinstance(item, dict) and field in item:
top_values[item[field]] = int(item["count"])
if top_values:
stats_output.append(FieldStat(field_name=field, top_values=top_values))
return stats_output
def _create_and_wait_splunk_job(service, search_query: str, start_time: float, end_time: float):
job = service.jobs.create(search_query, earliest_time=start_time, latest_time=end_time, exec_mode="normal")
while not job.is_done():
time.sleep(0.2)
return job
def _extract_field_types(properties: dict[str, Any], prefix: str, result: dict[str, str]) -> None:
for field_name, field_info in properties.items():
full_name = f"{prefix}{field_name}" if prefix else field_name
if "type" in field_info:
result[full_name] = field_info["type"]
if "properties" in field_info:
_extract_field_types(field_info["properties"], f"{full_name}.", result)
@lru_cache(maxsize=64)
def _get_elk_field_types(index_name: str) -> dict[str, str]:
client = ELKClient.get_client()
field_types: dict[str, str] = {}
try:
mapping_resp = client.indices.get_mapping(index=index_name)
except Exception:
return field_types
for _, index_mapping in mapping_resp.items():
properties = index_mapping.get("mappings", {}).get("properties", {})
_extract_field_types(properties, "", field_types)
return field_types
def _build_safe_aggs(agg_fields: list[str], index_name: str) -> dict[str, Any]:
field_types = _get_elk_field_types(index_name)
safe_aggs: dict[str, Any] = {}
for field in agg_fields:
field_type = field_types.get(field)
if field_type in (None, "text"):
agg_field = f"{field}.keyword"
agg_key = agg_field
else:
agg_field = field
agg_key = field
safe_aggs[agg_key] = {"terms": {"field": agg_field, "size": SAMPLE_COUNT}}
return safe_aggs
class ELKQueryBackend:
backend_name = "ELK"
@classmethod
def execute_structured_query(cls, input_data: AdaptiveQueryInput) -> BackendQueryResult:
client = ELKClient.get_client()
aggregation_fields = input_data.aggregation_fields or get_default_agg_fields(input_data.index_name)
must_clauses: list[dict[str, Any]] = [
_build_time_range_clause(input_data.time_field, input_data.time_range_start, input_data.time_range_end)
]
for field, value in input_data.filters.items():
if isinstance(value, list):
must_clauses.append({"terms": {field: value}})
else:
must_clauses.append({"term": {field: value}})
query_body = {"bool": {"must": must_clauses}}
response = client.search(
index=input_data.index_name,
query=query_body,
aggs=_build_safe_aggs(aggregation_fields, input_data.index_name),
size=SAMPLE_THRESHOLD,
track_total_hits=True,
)
return BackendQueryResult(
backend=cls.backend_name,
index_name=input_data.index_name,
total_hits=response["hits"]["total"]["value"],
aggregation_fields=aggregation_fields,
statistics=_extract_elk_stats(response, aggregation_fields),
raw_records=_extract_elk_records(response["hits"]["hits"]),
)
@classmethod
def execute_keyword_query(cls, input_data: KeywordSearchInput) -> BackendQueryResult:
client = ELKClient.get_client()
effective_index = input_data.index_name or "*"
aggregation_fields = get_default_agg_fields(input_data.index_name) if input_data.index_name else []
query_body = {
"bool": {
"must": [
_build_time_range_clause(input_data.time_field, input_data.time_range_start, input_data.time_range_end),
*_build_elk_keyword_clauses(input_data.keyword),
]
}
}
aggs = {"_index": {"terms": {"field": "_index", "size": 50}}}
if aggregation_fields:
aggs.update(_build_safe_aggs(aggregation_fields, effective_index))
response = client.search(
index=effective_index,
query=query_body,
aggs=aggs,
size=SAMPLE_THRESHOLD,
track_total_hits=True,
)
buckets = response.get("aggregations", {}).get("_index", {}).get("buckets", [])
index_distribution = {bucket["key"]: bucket["doc_count"] for bucket in buckets}
if input_data.index_name and input_data.index_name not in index_distribution:
index_distribution[input_data.index_name] = response["hits"]["total"]["value"]
return BackendQueryResult(
backend=cls.backend_name,
index_name=input_data.index_name or effective_index,
total_hits=response["hits"]["total"]["value"],
aggregation_fields=aggregation_fields,
statistics=_extract_elk_stats(response, aggregation_fields),
raw_records=_extract_elk_records(response["hits"]["hits"], include_index=True),
index_distribution=index_distribution,
)
@classmethod
def discover_keyword_hit_indices(cls, input_data: KeywordSearchInput, indices: list[str]) -> list[str]:
if not indices:
return []
client = ELKClient.get_client()
response = client.search(
index=",".join(indices),
query={
"bool": {
"must": [
_build_time_range_clause(
input_data.time_field,
input_data.time_range_start,
input_data.time_range_end,
),
*_build_elk_keyword_clauses(input_data.keyword),
]
}
},
aggs={"_index": {"terms": {"field": "_index", "size": 50}}},
size=0,
track_total_hits=True,
)
buckets = response.get("aggregations", {}).get("_index", {}).get("buckets", [])
return [bucket["key"] for bucket in buckets if bucket["doc_count"] > 0]
class SplunkQueryBackend:
backend_name = "Splunk"
@classmethod
def execute_structured_query(cls, input_data: AdaptiveQueryInput) -> BackendQueryResult:
service = SplunkClient.get_service()
start_time, end_time = parse_time_range(input_data.time_range_start, input_data.time_range_end)
search_query = f"search index=\"{input_data.index_name}\""
for field, value in input_data.filters.items():
if isinstance(value, list):
or_clause = " OR ".join(f'{field}=\"{item}\"' for item in value)
search_query += f" ({or_clause})"
else:
search_query += f" {field}=\"{value}\""
job = _create_and_wait_splunk_job(service, search_query, start_time, end_time)
total_hits = int(job["eventCount"])
aggregation_fields = input_data.aggregation_fields or get_default_agg_fields(input_data.index_name)
return BackendQueryResult(
backend=cls.backend_name,
index_name=input_data.index_name,
total_hits=total_hits,
aggregation_fields=aggregation_fields,
statistics=_fetch_splunk_top_stats(service, search_query, start_time, end_time, aggregation_fields)
if total_hits > 0
else [],
raw_records=_fetch_splunk_records(job, SAMPLE_THRESHOLD) if total_hits > 0 else [],
)
@classmethod
def execute_keyword_query(cls, input_data: KeywordSearchInput) -> BackendQueryResult:
service = SplunkClient.get_service()
start_time, end_time = parse_time_range(input_data.time_range_start, input_data.time_range_end)
effective_index = input_data.index_name or "*"
search_query = f"search index=\"{effective_index}\" ({_build_splunk_keyword_clause(input_data.keyword)})"
job = _create_and_wait_splunk_job(service, search_query, start_time, end_time)
total_hits = int(job["eventCount"])
index_distribution: dict[str, int] = {}
if total_hits > 0:
stats_query = f"{search_query} | stats count by index"
oneshot = service.jobs.oneshot(stats_query, earliest_time=start_time, latest_time=end_time, output_mode="json")
reader = JSONResultsReader(oneshot)
for item in reader:
if isinstance(item, dict) and "index" in item and "count" in item:
index_distribution[item["index"]] = int(item["count"])
if input_data.index_name and input_data.index_name not in index_distribution:
index_distribution[input_data.index_name] = total_hits
aggregation_fields = get_default_agg_fields(input_data.index_name) if input_data.index_name else []
statistics = (
_fetch_splunk_top_stats(service, search_query, start_time, end_time, aggregation_fields)
if total_hits > 0 and aggregation_fields
else []
)
return BackendQueryResult(
backend=cls.backend_name,
index_name=input_data.index_name or effective_index,
total_hits=total_hits,
aggregation_fields=aggregation_fields,
statistics=statistics,
raw_records=_fetch_splunk_records(job, SAMPLE_THRESHOLD) if total_hits > 0 else [],
index_distribution=index_distribution,
)
@classmethod
def discover_keyword_hit_indices(cls, input_data: KeywordSearchInput, indices: list[str]) -> list[str]:
if not indices:
return []
service = SplunkClient.get_service()
start_time, end_time = parse_time_range(input_data.time_range_start, input_data.time_range_end)
index_clause = " OR ".join(f'index=\"{index}\"' for index in indices)
search_query = f"search ({index_clause}) ({_build_splunk_keyword_clause(input_data.keyword)}) | stats count by index"
oneshot = service.jobs.oneshot(search_query, earliest_time=start_time, latest_time=end_time, output_mode="json")
reader = JSONResultsReader(oneshot)
hit_indices: list[str] = []
for item in reader:
if isinstance(item, dict) and "index" in item and "count" in item and int(item["count"]) > 0:
hit_indices.append(item["index"])
return hit_indices
PLUGINS/SIEM/registry.py
+27 -34
View File
@@ -1,30 +1,18 @@
from __future__ import annotations
from functools import lru_cache
from pathlib import Path
from typing import List, Dict, Literal
from typing import Dict, List
import yaml
from pydantic import BaseModel
from Lib.log import logger
from PLUGINS.SIEM.models import SchemaFieldInfo, IndexInfo
class FieldInfo(BaseModel):
name: str
type: str
description: str
is_key_field: bool = False
class IndexInfo(BaseModel):
name: str
backend: Literal["ELK", "Splunk"]
description: str
fields: List[FieldInfo]
@lru_cache(maxsize=1)
def _load_yaml_configs() -> Dict[str, IndexInfo]:
registry = {}
registry: Dict[str, IndexInfo] = {}
script_path = Path(__file__).resolve()
project_root = script_path.parents[2]
@@ -35,33 +23,38 @@ def _load_yaml_configs() -> Dict[str, IndexInfo]:
for yaml_file in indexs_dir.glob("*.yaml"):
try:
with open(yaml_file, 'r', encoding='utf-8') as f:
data = yaml.safe_load(f)
with open(yaml_file, "r", encoding="utf-8") as file:
data = yaml.safe_load(file) or {}
fields = [FieldInfo(**field) for field in data.get('fields', [])]
fields = [SchemaFieldInfo(**field) for field in data.get("fields", [])]
index_info = IndexInfo(
name=data['name'],
backend=data['backend'],
description=data['description'],
fields=fields
name=data["name"],
backend=data["backend"],
description=data["description"],
fields=fields,
)
registry[index_info.name] = index_info
except Exception as e:
logger.exception(f"Error loading YAML file {yaml_file}: {e}")
except Exception as exc:
logger.exception(f"Error loading YAML file {yaml_file}: {exc}")
return registry
def get_default_agg_fields(index_name: str) -> List[str]:
def list_indices() -> List[IndexInfo]:
return list(_load_yaml_configs().values())
def get_index_info(index_name: str) -> IndexInfo:
registry = _load_yaml_configs()
if index_name not in registry:
return []
result = [f.name for f in registry[index_name].fields if f.is_key_field]
return result
raise ValueError(f"Index {index_name} not found.")
return registry[index_name]
def get_default_agg_fields(index_name: str) -> List[str]:
index_info = get_index_info(index_name)
return [field.name for field in index_info.fields if field.is_key_field]
def get_backend_type(index_name: str) -> str:
registry = _load_yaml_configs()
if index_name in registry:
return registry[index_name].backend
return "ELK"
return get_index_info(index_name).backend
PLUGINS/SIEM/response_builder.py
+144
View File
@@ -0,0 +1,144 @@
from __future__ import annotations
from typing import Any, Literal
from PLUGINS.SIEM.models import (
AdaptiveQueryInput,
AdaptiveQueryOutput,
KeywordSearchInput,
KeywordSearchOutput,
SAMPLE_COUNT,
SAMPLE_THRESHOLD,
SUMMARY_THRESHOLD,
)
from PLUGINS.SIEM.query_backends import BackendQueryResult
from PLUGINS.SIEM.registry import get_default_agg_fields
def resolve_status(total_hits: int) -> Literal["records", "sample", "summary"]:
if total_hits > SUMMARY_THRESHOLD:
return "summary"
if total_hits > SAMPLE_THRESHOLD:
return "sample"
return "records"
def build_adaptive_output(input_data: AdaptiveQueryInput, result: BackendQueryResult) -> AdaptiveQueryOutput:
status = resolve_status(result.total_hits)
records = _project_records(
result.raw_records[: _record_limit_for_status(status)],
index_name=result.index_name,
time_field=input_data.time_field,
explicit_fields=list(input_data.filters.keys()) + result.aggregation_fields,
)
return AdaptiveQueryOutput(
backend=result.backend,
index_name=result.index_name,
status=status,
total_hits=result.total_hits,
returned_records=len(records),
truncated=_is_truncated(result.raw_records, records, result.total_hits, status),
message=_build_message(result.backend, result.index_name, result.total_hits, status),
statistics=result.statistics,
records=records,
)
def build_keyword_output(input_data: KeywordSearchInput, result: BackendQueryResult) -> KeywordSearchOutput:
status = resolve_status(result.total_hits)
records = _project_records(
result.raw_records[: _record_limit_for_status(status)],
index_name=result.index_name,
time_field=input_data.time_field,
explicit_fields=result.aggregation_fields,
)
index_distribution = result.index_distribution or {result.index_name: result.total_hits}
return KeywordSearchOutput(
backend=result.backend,
index_name=result.index_name,
status=status,
total_hits=result.total_hits,
returned_records=len(records),
truncated=_is_truncated(result.raw_records, records, result.total_hits, status),
message=_build_message(result.backend, result.index_name, result.total_hits, status),
index_distribution=index_distribution,
statistics=result.statistics,
records=records,
)
def _record_limit_for_status(status: Literal["records", "sample", "summary"]) -> int:
if status == "records":
return SAMPLE_THRESHOLD
if status == "sample":
return SAMPLE_COUNT
return SAMPLE_COUNT
def _build_message(backend: str, index_name: str, total_hits: int, status: str) -> str:
if status == "summary":
return f"Matched {total_hits} events in {index_name} ({backend}). Returning statistics only."
if status == "sample":
return f"Matched {total_hits} events in {index_name} ({backend}). Returning statistics and projected samples."
return f"Matched {total_hits} events in {index_name} ({backend}). Returning projected records."
def _is_truncated(
raw_records: list[dict[str, Any]],
projected_records: list[dict[str, Any]],
total_hits: int,
status: str,
) -> bool:
if status != "records":
return total_hits > len(projected_records)
if total_hits > len(projected_records):
return True
return any(len(projected) < len(raw) for raw, projected in zip(raw_records[: len(projected_records)], projected_records))
def _project_records(
records: list[dict[str, Any]],
*,
index_name: str,
time_field: str,
explicit_fields: list[str],
) -> list[dict[str, Any]]:
projection_fields = _build_projection_fields(index_name=index_name, time_field=time_field, explicit_fields=explicit_fields)
return [_project_record(record, projection_fields) for record in records]
def _build_projection_fields(*, index_name: str, time_field: str, explicit_fields: list[str]) -> list[str]:
ordered_fields: list[str] = []
for field in [time_field, *explicit_fields, *get_default_agg_fields(index_name)]:
if field and field not in ordered_fields:
ordered_fields.append(field)
return ordered_fields
def _project_record(record: dict[str, Any], projection_fields: list[str]) -> dict[str, Any]:
projected: dict[str, Any] = {}
for field in projection_fields:
found, value = _extract_field_value(record, field)
if found:
projected[field] = value
if "_index" in record:
projected["_index"] = record["_index"]
if projected:
return projected
fallback_fields = list(record.keys())[: min(len(record), SAMPLE_COUNT)]
return {field: record[field] for field in fallback_fields}
def _extract_field_value(record: dict[str, Any], field_path: str) -> tuple[bool, Any]:
if field_path in record:
return True, record[field_path]
current: Any = record
for segment in field_path.split("."):
if not isinstance(current, dict) or segment not in current:
return False, None
current = current[segment]
return True, current
PLUGINS/SIEM/tools.py
+54 -568
View File
@@ -1,616 +1,102 @@
import json
import re
import time
from datetime import datetime, timezone
from typing import List
from typing import List, Union
from splunklib.results import JSONResultsReader
from PLUGINS.ELK.client import ELKClient
from PLUGINS.SIEM.models import (
SchemaExplorerInput,
AdaptiveQueryInput,
KeywordSearchInput,
AdaptiveQueryOutput,
KeywordSearchInput,
KeywordSearchOutput,
FieldStat,
SUMMARY_THRESHOLD,
SAMPLE_THRESHOLD, SAMPLE_COUNT
SchemaExplorerInput,
SchemaIndexSummary, IndexInfo,
)
from PLUGINS.SIEM.registry import _load_yaml_configs, get_default_agg_fields, get_backend_type
from PLUGINS.Splunk.client import SplunkClient
from PLUGINS.SIEM.query_backends import ELKQueryBackend, SplunkQueryBackend
from PLUGINS.SIEM.registry import get_backend_type, get_default_agg_fields, get_index_info, list_indices
from PLUGINS.SIEM.response_builder import build_adaptive_output, build_keyword_output
def get_indices_by_backend() -> dict:
registry = _load_yaml_configs()
result = {"ELK": [], "Splunk": []}
for idx_name, idx_info in registry.items():
if idx_info.backend in result:
result[idx_info.backend].append(idx_name)
for index_info in list_indices():
result.setdefault(index_info.backend, []).append(index_info.name)
return result
class SIEMToolKit(object):
class SIEMToolKit:
@classmethod
def explore_schema(cls, input_data: SchemaExplorerInput = SchemaExplorerInput(target_index=None)):
def explore_schema(cls, input_data: SchemaExplorerInput) -> Union[IndexInfo, List[SchemaIndexSummary]]:
"""
Explore available SIEM indices and their field schemas.
Explore registered SIEM indices and their declared schemas.
This tool helps agents discover what data sources are available and what fields they contain.
It supports two modes based on the target_index parameter in SchemaExplorerInput:
1. List all indices (when target_index is None)
2. Get detailed field information for a specific index
See SchemaExplorerInput for detailed parameter documentation.
When `target_index` is omitted, the tool returns summaries for all registered indices.
When `target_index` is provided, the tool returns the field definitions for that index.
Raises:
ValueError: If the specified target_index is not found in the registry.
Example Usage by Agent:
# List all indices
explore_schema()
# Get details on "logs-security" index
explore_schema(SchemaExplorerInput(target_index="logs-security"))
ValueError: If `target_index` is not present in the SIEM registry.
"""
if not input_data.target_index:
registry = _load_yaml_configs()
result = [
{"name": k, "description": v.description}
for k, v in registry.items()
return [
SchemaIndexSummary(
name=index_info.name,
backend=index_info.backend,
description=index_info.description,
default_aggregation_fields=get_default_agg_fields(index_info.name),
)
for index_info in list_indices()
]
return result
registry = _load_yaml_configs()
if input_data.target_index not in registry:
raise ValueError(f"Index {input_data.target_index} not found.")
idx_info = registry[input_data.target_index]
result = [f.model_dump() for f in idx_info.fields]
return result
index_info = get_index_info(input_data.target_index)
return index_info
@classmethod
def execute_adaptive_query(cls, input_data: AdaptiveQueryInput) -> AdaptiveQueryOutput:
"""
Execute adaptive queries against SIEM backends (ELK or Splunk) with intelligent response formatting.
Execute an exact-match SIEM query and return an LLM-safe response.
This tool executes queries with automatic backend detection and response optimization:
- Automatically adjusts response format based on result volume:
* Full logs: Complete log records (for < 20 results)
* Sample: Statistics + sample records (for 20-1000 results)
* Summary: Statistics only (for > 1000 results)
- Provides top-N statistics for specified aggregation fields
- Handles time range filtering with UTC ISO8601 timestamps
Raises:
ValueError: If time format is invalid or backend is unsupported
ConnectionError: If SIEM backend is unreachable
Example Usage by Agent:
# Query security logs from last hour
input_data = AdaptiveQueryInput(
index_name="logs-security",
time_range_start="2026-02-04T06:00:00Z",
time_range_end="2026-02-04T07:00:00Z",
filters={"event.outcome": "failure"},
aggregation_fields=["event.action", "user.name"]
)
result = execute_adaptive_query(input_data)
# Agent can then analyze result.statistics for patterns
# and if needed, drill down with result.records
The response uses three status levels:
- `records`: returns projected records when result volume is small
- `sample`: returns statistics and projected sample records
- `summary`: returns statistics only
"""
backend = get_backend_type(input_data.index_name)
if backend == "ELK":
result = cls._execute_elk(input_data)
return result
elif backend == "Splunk":
result = cls._execute_splunk(input_data)
return result
else:
raise ValueError(f"Unsupported backend: {backend}")
query_backend = cls._get_query_backend(backend)
backend_result = query_backend.execute_structured_query(input_data)
return build_adaptive_output(input_data, backend_result)
@classmethod
def keyword_search(cls, input_data: KeywordSearchInput) -> List[KeywordSearchOutput]:
"""
Execute keyword-based search across SIEM backends with intelligent response formatting.
Execute keyword search against SIEM data and return one result per matched index.
This tool performs full-text search using one keyword or a list of keywords across all fields (or specified index):
- Supports searching by IP, hostname, username, or any arbitrary string
- When a keyword list is provided, all keywords must match in the same search
- When index_name is not specified, searches BOTH ELK and Splunk backends and returns results from each
- Applies the same adaptive response strategy as execute_adaptive_query:
* Full logs: < 100 results
* Sample: 100-1000 results (statistics + samples)
* Summary: > 1000 results (statistics only)
- Provides top-N statistics for specified aggregation fields
- Handles time range filtering with UTC ISO8601 timestamps
Raises:
ValueError: If time format is invalid or backend is unsupported
ConnectionError: If SIEM backend is unreachable
Example Usage by Agent:
# Search for an IP across all indices (returns results from both ELK and Splunk)
input_data = KeywordSearchInput(
keyword="192.168.1.100",
time_range_start="2026-02-04T06:00:00Z",
time_range_end="2026-02-04T07:00:00Z"
)
result = keyword_search(input_data)
# Search for multiple terms with AND semantics
input_data = KeywordSearchInput(
keyword=["alice", "10.10.10.15"],
time_range_start="2026-02-04T06:00:00Z",
time_range_end="2026-02-04T07:00:00Z"
)
result = keyword_search(input_data)
# Search for hostname in specific index
input_data = KeywordSearchInput(
keyword="DESKTOP-ABC123",
time_range_start="2026-02-04T06:00:00Z",
time_range_end="2026-02-04T07:00:00Z",
index_name="logs-endpoint"
)
result = keyword_search(input_data)
If `index_name` is provided, the tool queries only that index and returns a single-item list.
If `index_name` is omitted, the tool first discovers hit indices across the registered backends and then
runs per-index searches so each response stays small and attributable to a single source.
"""
if input_data.index_name:
backend = get_backend_type(input_data.index_name)
if backend == "ELK":
return [cls._keyword_search_elk(input_data)]
elif backend == "Splunk":
return [cls._keyword_search_splunk(input_data)]
else:
raise ValueError(f"Unsupported backend: {backend}")
backend_result = cls._get_query_backend(backend).execute_keyword_query(input_data)
return [build_keyword_output(input_data, backend_result)]
results: list[KeywordSearchOutput] = []
indices_by_backend = get_indices_by_backend()
results = []
elk_indices = indices_by_backend.get("ELK", [])
if elk_indices:
hit_indices = cls._discover_elk_hit_indices(input_data, elk_indices)
for idx_name in hit_indices:
modified_input = KeywordSearchInput(
for backend_name, indices in indices_by_backend.items():
query_backend = cls._get_query_backend(backend_name)
for index_name in query_backend.discover_keyword_hit_indices(input_data, indices):
per_index_input = KeywordSearchInput(
keyword=input_data.keyword,
time_range_start=input_data.time_range_start,
time_range_end=input_data.time_range_end,
time_field=input_data.time_field,
index_name=idx_name
index_name=index_name,
)
result = cls._keyword_search_elk(modified_input)
result.backend = "ELK"
results.append(result)
splunk_indices = indices_by_backend.get("Splunk", [])
if splunk_indices:
hit_indices = cls._discover_splunk_hit_indices(input_data, splunk_indices)
for idx_name in hit_indices:
modified_input = KeywordSearchInput(
keyword=input_data.keyword,
time_range_start=input_data.time_range_start,
time_range_end=input_data.time_range_end,
time_field=input_data.time_field,
index_name=idx_name
)
result = cls._keyword_search_splunk(modified_input)
result.backend = "Splunk"
results.append(result)
backend_result = query_backend.execute_keyword_query(per_index_input)
results.append(build_keyword_output(per_index_input, backend_result))
return results
@classmethod
def _build_time_range_clause(cls, time_field: str, time_range_start: str, time_range_end: str) -> dict:
return {
"range": {
time_field: {
"gte": time_range_start,
"lt": time_range_end
}
}
}
@classmethod
def _normalize_keywords(cls, keyword_input: str | list[str]) -> list[str]:
if isinstance(keyword_input, str):
return [keyword_input]
return keyword_input
@classmethod
def _build_elk_keyword_clauses(cls, keyword_input: str | list[str]) -> list[dict]:
return [
{"multi_match": {"query": keyword, "type": "best_fields", "fuzziness": "AUTO"}}
for keyword in cls._normalize_keywords(keyword_input)
]
@classmethod
def _format_splunk_keyword(cls, keyword: str) -> str:
if re.fullmatch(r"[A-Za-z0-9._:@/\\-]+", keyword):
return keyword
escaped_keyword = keyword.replace("\\", "\\\\").replace('"', '\\"')
return f'"{escaped_keyword}"'
@classmethod
def _build_splunk_keyword_clause(cls, keyword_input: str | list[str]) -> str:
keywords = cls._normalize_keywords(keyword_input)
return " AND ".join(cls._format_splunk_keyword(keyword) for keyword in keywords)
@classmethod
def _extract_elk_records(cls, hits: list, include_index: bool = False) -> list[dict]:
records = []
for hit in hits:
record = hit["_source"].copy() if include_index else hit["_source"]
if include_index:
record["_index"] = hit["_index"]
records.append(record)
return records
@classmethod
def _extract_elk_stats(cls, response: dict, agg_fields: list) -> list[FieldStat]:
stats_output = []
if "aggregations" not in response:
return stats_output
for field in agg_fields:
agg_key = f"{field}.keyword" if f"{field}.keyword" in response["aggregations"] else field
if agg_key in response["aggregations"]:
buckets = response["aggregations"][agg_key]["buckets"]
if buckets:
stats_output.append(FieldStat(
field_name=field,
top_values={b["key"]: b["doc_count"] for b in buckets}
))
return stats_output
@classmethod
def _parse_time_range(cls, time_range_start: str, time_range_end: str) -> tuple[float, float]:
utc_format = "%Y-%m-%dT%H:%M:%SZ"
try:
dt_start = datetime.strptime(time_range_start, utc_format).replace(tzinfo=timezone.utc)
dt_end = datetime.strptime(time_range_end, utc_format).replace(tzinfo=timezone.utc)
return dt_start.timestamp(), dt_end.timestamp()
except ValueError:
raise ValueError("Invalid UTC format.")
@classmethod
def _clean_splunk_record(cls, log: dict) -> dict:
clean_record = {}
for k, v in log.items():
if not k.startswith("_") and k not in ["_raw", "splunk_server", "host", "source", "sourcetype"]:
clean_record[k] = v
if "_time" in log:
clean_record["@timestamp"] = log["_time"]
if "_raw" in log:
try:
raw_parsed = json.loads(log["_raw"])
if isinstance(raw_parsed, dict):
for rk, rv in raw_parsed.items():
if rk not in clean_record:
clean_record[rk] = rv
except (json.JSONDecodeError, TypeError):
pass
return clean_record
@classmethod
def _fetch_splunk_records(cls, job, count: int) -> list[dict]:
records = []
results = job.results(count=count, output_mode="json")
for result in results:
result = json.loads(result)
for log in result.get("results", []):
records.append(cls._clean_splunk_record(log))
return records
@classmethod
def _fetch_splunk_top_stats(cls, service, search_query: str, t_start: float, t_end: float, agg_fields: list) -> list[FieldStat]:
stats_output = []
for field in agg_fields:
stats_spl = f"{search_query} | top limit={SAMPLE_COUNT} {field}"
rr = service.jobs.oneshot(stats_spl, earliest_time=t_start, latest_time=t_end, output_mode="json")
reader = JSONResultsReader(rr)
top_vals = {}
for item in reader:
if isinstance(item, dict) and field in item:
top_vals[item[field]] = int(item['count'])
if top_vals:
stats_output.append(FieldStat(field_name=field, top_values=top_vals))
return stats_output
@classmethod
def _create_and_wait_splunk_job(cls, service, search_query: str, t_start: float, t_end: float):
job = service.jobs.create(search_query, earliest_time=t_start, latest_time=t_end, exec_mode="normal")
while not job.is_done():
time.sleep(0.2)
return job
@classmethod
def _execute_elk(cls, input_data: AdaptiveQueryInput) -> AdaptiveQueryOutput:
client = ELKClient.get_client()
must_clauses = [cls._build_time_range_clause(input_data.time_field, input_data.time_range_start, input_data.time_range_end)]
for k, v in input_data.filters.items():
if isinstance(v, list):
must_clauses.append({"terms": {k: v}})
else:
must_clauses.append({"term": {k: v}})
query_body = {"bool": {"must": must_clauses}}
agg_fields = input_data.aggregation_fields or get_default_agg_fields(input_data.index_name)
aggs_dsl = cls._build_safe_aggs(agg_fields, input_data.index_name)
response = client.search(
index=input_data.index_name, query=query_body, aggs=aggs_dsl, size=SAMPLE_COUNT, track_total_hits=True
)
total_hits = response["hits"]["total"]["value"]
hits_data = cls._extract_elk_records(response["hits"]["hits"])
stats_output = cls._extract_elk_stats(response, agg_fields)
return cls._apply_funnel_strategy(total_hits, stats_output, hits_data, input_data, client, query_body)
@classmethod
def _keyword_search_elk(cls, input_data: KeywordSearchInput) -> KeywordSearchOutput:
client = ELKClient.get_client()
effective_index = input_data.index_name or "*"
must_clauses = [
cls._build_time_range_clause(input_data.time_field, input_data.time_range_start, input_data.time_range_end),
*cls._build_elk_keyword_clauses(input_data.keyword)
]
query_body = {"bool": {"must": must_clauses}}
aggs_dsl = {"_index": {"terms": {"field": "_index", "size": 50}}}
agg_fields = []
if input_data.index_name:
agg_fields = get_default_agg_fields(input_data.index_name)
field_aggs = cls._build_safe_aggs(agg_fields, input_data.index_name)
aggs_dsl.update(field_aggs)
response = client.search(
index=effective_index, query=query_body, aggs=aggs_dsl, size=SAMPLE_COUNT, track_total_hits=True
)
total_hits = response["hits"]["total"]["value"]
hits_data = cls._extract_elk_records(response["hits"]["hits"], include_index=True)
index_distribution = {}
if "aggregations" in response and "_index" in response["aggregations"]:
buckets = response["aggregations"]["_index"]["buckets"]
index_distribution = {b["key"]: b["doc_count"] for b in buckets}
stats_output = cls._extract_elk_stats(response, agg_fields)
status = cls._resolve_funnel_status(total_hits)
idx_count = len(index_distribution)
if status == "summary":
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=[],
message=f"Found {total_hits} events across {idx_count} index(es). Showing statistics only."
)
elif status == "sample":
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=hits_data,
message=f"Found {total_hits} events across {idx_count} index(es). Showing statistics + samples."
)
else:
final_records = hits_data
if total_hits > SAMPLE_COUNT:
resp = client.search(index=effective_index, query=query_body, size=SAMPLE_THRESHOLD)
final_records = cls._extract_elk_records(resp["hits"]["hits"], include_index=True)
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=final_records,
message=f"Found {total_hits} events. Returning full logs."
)
@classmethod
def _execute_splunk(cls, input_data: AdaptiveQueryInput) -> AdaptiveQueryOutput:
service = SplunkClient.get_service()
t_start, t_end = cls._parse_time_range(input_data.time_range_start, input_data.time_range_end)
search_query = f"search index=\"{input_data.index_name}\""
for k, v in input_data.filters.items():
if isinstance(v, list):
or_clause = " OR ".join([f'{k}="{val}"' for val in v])
search_query += f" ({or_clause})"
else:
search_query += f" {k}=\"{v}\""
job = cls._create_and_wait_splunk_job(service, search_query, t_start, t_end)
total_hits = int(job["eventCount"])
agg_fields = input_data.aggregation_fields or get_default_agg_fields(input_data.index_name)
stats_output = cls._fetch_splunk_top_stats(service, search_query, t_start, t_end, agg_fields) if total_hits > 0 else []
hits_data = cls._fetch_splunk_records(job, SAMPLE_COUNT) if total_hits > 0 else []
status = cls._resolve_funnel_status(total_hits)
if status == "summary":
return AdaptiveQueryOutput(
status=status, total_hits=total_hits, statistics=stats_output, records=[],
message=f"Found {total_hits} events in Splunk. Showing statistics only."
)
elif status == "sample":
return AdaptiveQueryOutput(
status=status, total_hits=total_hits, statistics=stats_output, records=hits_data,
message=f"Found {total_hits} events in Splunk. Showing statistics + samples."
)
else:
final_records = cls._fetch_splunk_records(job, SAMPLE_THRESHOLD)
return AdaptiveQueryOutput(
status=status, total_hits=total_hits, statistics=stats_output, records=final_records,
message="Low volume. Returning full logs."
)
@classmethod
def _keyword_search_splunk(cls, input_data: KeywordSearchInput) -> KeywordSearchOutput:
service = SplunkClient.get_service()
t_start, t_end = cls._parse_time_range(input_data.time_range_start, input_data.time_range_end)
effective_index = input_data.index_name or "*"
keyword_clause = cls._build_splunk_keyword_clause(input_data.keyword)
search_query = f"search index=\"{effective_index}\" ({keyword_clause})"
job = cls._create_and_wait_splunk_job(service, search_query, t_start, t_end)
total_hits = int(job["eventCount"])
index_distribution = {}
if total_hits > 0:
index_stats_query = f"{search_query} | stats count by index"
rr = service.jobs.oneshot(index_stats_query, earliest_time=t_start, latest_time=t_end, output_mode="json")
reader = JSONResultsReader(rr)
for item in reader:
if isinstance(item, dict) and "index" in item and "count" in item:
index_distribution[item["index"]] = int(item["count"])
agg_fields = []
stats_output = []
if input_data.index_name:
agg_fields = get_default_agg_fields(input_data.index_name)
if total_hits > 0:
stats_output = cls._fetch_splunk_top_stats(service, search_query, t_start, t_end, agg_fields)
hits_data = cls._fetch_splunk_records(job, SAMPLE_COUNT) if total_hits > 0 else []
status = cls._resolve_funnel_status(total_hits)
idx_count = len(index_distribution)
if status == "summary":
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=[],
message=f"Found {total_hits} events across {idx_count} index(es) in Splunk. Showing statistics only."
)
elif status == "sample":
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=hits_data,
message=f"Found {total_hits} events across {idx_count} index(es) in Splunk. Showing statistics + samples."
)
else:
final_records = cls._fetch_splunk_records(job, SAMPLE_THRESHOLD)
return KeywordSearchOutput(
status=status, total_hits=total_hits, index_distribution=index_distribution,
statistics=stats_output, records=final_records,
message=f"Found {total_hits} events in Splunk. Returning full logs."
)
@classmethod
def _discover_elk_hit_indices(cls, input_data: KeywordSearchInput, elk_indices: list) -> list:
client = ELKClient.get_client()
index_pattern = ",".join(elk_indices)
must_clauses = [
cls._build_time_range_clause(input_data.time_field, input_data.time_range_start, input_data.time_range_end),
*cls._build_elk_keyword_clauses(input_data.keyword)
]
query_body = {"bool": {"must": must_clauses}}
aggs_dsl = {"_index": {"terms": {"field": "_index", "size": 50}}}
response = client.search(
index=index_pattern, query=query_body, aggs=aggs_dsl, size=0, track_total_hits=True
)
hit_indices = []
if "aggregations" in response and "_index" in response["aggregations"]:
buckets = response["aggregations"]["_index"]["buckets"]
hit_indices = [b["key"] for b in buckets if b["doc_count"] > 0]
return hit_indices
@classmethod
def _discover_splunk_hit_indices(cls, input_data: KeywordSearchInput, splunk_indices: list) -> list:
service = SplunkClient.get_service()
t_start, t_end = cls._parse_time_range(input_data.time_range_start, input_data.time_range_end)
index_clause = " OR ".join([f'index="{idx}"' for idx in splunk_indices])
keyword_clause = cls._build_splunk_keyword_clause(input_data.keyword)
search_query = f"search ({index_clause}) ({keyword_clause}) | stats count by index"
rr = service.jobs.oneshot(search_query, earliest_time=t_start, latest_time=t_end, output_mode="json")
reader = JSONResultsReader(rr)
hit_indices = []
for item in reader:
if isinstance(item, dict) and "index" in item and "count" in item:
if int(item["count"]) > 0:
hit_indices.append(item["index"])
return hit_indices
@classmethod
def _build_safe_aggs(cls, agg_fields, index_name="*"):
client = ELKClient.get_client()
field_types = {}
try:
mapping_resp = client.indices.get_mapping(index=index_name)
for idx_name, idx_mapping in mapping_resp.items():
properties = idx_mapping.get("mappings", {}).get("properties", {})
cls._extract_field_types(properties, "", field_types)
except Exception:
pass
safe_aggs = {}
for f in agg_fields:
field_type = field_types.get(f)
if field_type == "text":
agg_field = f"{f}.keyword"
agg_key = f"{f}.keyword"
elif field_type in (None,):
agg_field = f"{f}.keyword"
agg_key = f"{f}.keyword"
else:
agg_field = f
agg_key = f
safe_aggs[agg_key] = {"terms": {"field": agg_field, "size": 5}}
return safe_aggs
@classmethod
def _extract_field_types(cls, properties: dict, prefix: str, result: dict):
for field_name, field_info in properties.items():
full_name = f"{prefix}{field_name}" if prefix else field_name
if "type" in field_info:
result[full_name] = field_info["type"]
if "properties" in field_info:
cls._extract_field_types(field_info["properties"], f"{full_name}.", result)
@classmethod
def _apply_funnel_strategy(cls, total, stats, initial_hits, input_data, client, query_body, index_name=None):
effective_index = index_name if index_name is not None else input_data.index_name
status = cls._resolve_funnel_status(total)
if status == "summary":
return AdaptiveQueryOutput(
status="summary", total_hits=total, statistics=stats, records=[],
message=f"Matches {total} records (ELK). High volume."
)
if status == "sample":
return AdaptiveQueryOutput(
status="sample", total_hits=total, statistics=stats, records=initial_hits,
message=f"Matches {total} records (ELK). Showing samples."
)
final_recs = initial_hits
if total > SAMPLE_COUNT:
resp = client.search(index=effective_index, query=query_body, size=SAMPLE_THRESHOLD)
final_recs = [h["_source"] for h in resp["hits"]["hits"]]
return AdaptiveQueryOutput(
status="full", total_hits=total, statistics=stats, records=final_recs,
message="Low volume. Returning full logs."
)
@classmethod
def _resolve_funnel_status(cls, total_hits: int) -> str:
if total_hits > SUMMARY_THRESHOLD:
return "summary"
if SAMPLE_THRESHOLD < total_hits <= SUMMARY_THRESHOLD:
return "sample"
return "full"
@staticmethod
def _get_query_backend(backend: str):
if backend == "ELK":
return ELKQueryBackend
if backend == "Splunk":
return SplunkQueryBackend
raise ValueError(f"Unsupported backend: {backend}")
pyproject.toml
+1
View File
@@ -44,6 +44,7 @@ dependencies = [
"uvicorn>=0.42.0",
"fastmcp>=3.2.0",
"python-dateutil>=2.9.0.post0",
"pyyaml>=6.0.3",
]
[[tool.uv.index]]
url = "https://mirrors.aliyun.com/pypi/simple/"
uv.lock
Generated
+2
View File
@@ -178,6 +178,7 @@ dependencies = [
{ name = "pycryptodome" },
{ name = "pydantic" },
{ name = "python-dateutil" },
{ name = "pyyaml" },
{ name = "redis" },
{ name = "requests" },
{ name = "rsa" },
@@ -221,6 +222,7 @@ requires-dist = [
{ name = "pycryptodome", specifier = ">=3.23.0" },
{ name = "pydantic", specifier = ">=2.12.5" },
{ name = "python-dateutil", specifier = ">=2.9.0.post0" },
{ name = "pyyaml", specifier = ">=6.0.3" },
{ name = "redis", specifier = ">=7.4.0" },
{ name = "requests", specifier = ">=2.32.5" },
{ name = "rsa", specifier = ">=4.9.1" },