diff --git a/DATA/Agent_SIEM/system_prompt.md b/DATA/Agent_SIEM/system_prompt.md index 8a33299..72b56d8 100644 --- a/DATA/Agent_SIEM/system_prompt.md +++ b/DATA/Agent_SIEM/system_prompt.md @@ -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 diff --git a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py index 2132137..b84365b 100644 --- a/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py +++ b/MODULES/Cloud-01-AWS-IAM-Privilege-Escalation-via-AttachUserPolicy.py @@ -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() diff --git a/PLUGINS/MCP/llmfunc.py b/PLUGINS/MCP/llmfunc.py index a46729a..c426245 100644 --- a/PLUGINS/MCP/llmfunc.py +++ b/PLUGINS/MCP/llmfunc.py @@ -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( diff --git a/PLUGINS/SIEM/models.py b/PLUGINS/SIEM/models.py index b13c48e..fa48f49 100644 --- a/PLUGINS/SIEM/models.py +++ b/PLUGINS/SIEM/models.py @@ -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] diff --git a/PLUGINS/SIEM/query_backends.py b/PLUGINS/SIEM/query_backends.py new file mode 100644 index 0000000..2cbcec0 --- /dev/null +++ b/PLUGINS/SIEM/query_backends.py @@ -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 diff --git a/PLUGINS/SIEM/registry.py b/PLUGINS/SIEM/registry.py index 2bbf054..edfaac9 100644 --- a/PLUGINS/SIEM/registry.py +++ b/PLUGINS/SIEM/registry.py @@ -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 diff --git a/PLUGINS/SIEM/response_builder.py b/PLUGINS/SIEM/response_builder.py new file mode 100644 index 0000000..acbc07b --- /dev/null +++ b/PLUGINS/SIEM/response_builder.py @@ -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 diff --git a/PLUGINS/SIEM/tools.py b/PLUGINS/SIEM/tools.py index e4f58ea..571fa15 100644 --- a/PLUGINS/SIEM/tools.py +++ b/PLUGINS/SIEM/tools.py @@ -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}") diff --git a/pyproject.toml b/pyproject.toml index 14dc6e5..c8797a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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/" diff --git a/uv.lock b/uv.lock index 6610c5f..2c491fa 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },