diff --git a/PLUGINS/MCP/llmfunc.py b/PLUGINS/MCP/llmfunc.py index e3e0812..83aea67 100644 --- a/PLUGINS/MCP/llmfunc.py +++ b/PLUGINS/MCP/llmfunc.py @@ -6,8 +6,8 @@ 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, QueryOutput, IndexInfo, SchemaIndexSummary, \ - DiscoverIndexFieldsInput, DiscoverIndexFieldsOutput +from PLUGINS.SIEM.models import AdaptiveQueryInput, ESQLQueryInput, KeywordSearchInput, SchemaExplorerInput, QueryOutput, IndexInfo, SchemaIndexSummary, \ + DiscoverIndexFieldsInput, DiscoverIndexFieldsOutput, SPLQueryInput 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 @@ -370,6 +370,42 @@ def siem_discover_index_fields( return SIEMToolKit.discover_index_fields(input_data) +def siem_execute_spl( + query: Annotated[str, Field(description="Raw Splunk SPL query string (原始 Splunk SPL 查询语句)")], + limit: Annotated[int, Field(description="Maximum number of records to return, default 100 (最大返回记录数,默认 100)")] = 100, + time_range_start: Annotated[Optional[str], Field(description="Optional UTC start time in ISO8601 (可选的 UTC 开始时间,ISO8601 格式)")] = None, + time_range_end: Annotated[Optional[str], Field(description="Optional UTC end time in ISO8601 (可选的 UTC 结束时间,ISO8601 格式)")] = None, + time_field: Annotated[str, Field(description="Time field for range filtering (时间范围过滤字段名)")] = "@timestamp", + index_name: Annotated[Optional[str], Field(description="Index name for output labeling; omit if specified in SPL (用于输出标记的索引名,SPL 中已指定时可不填)")] = None, +) -> Annotated[str, Field(description="Raw SPL query result as JSON string (原始 SPL 查询结果 JSON 字符串)")]: + """Execute a raw Splunk SPL query and return results. (执行原始 Splunk SPL 查询并返回结果)""" + input_data = SPLQueryInput( + query=query, limit=limit, + time_range_start=time_range_start, time_range_end=time_range_end, + time_field=time_field, index_name=index_name, + ) + result = SIEMToolKit.execute_spl(input_data) + return result.model_dump_json() + + +def siem_execute_esql( + query: Annotated[str, Field(description="Raw ELK ES|QL query string (原始 ELK ES|QL 查询语句)")], + limit: Annotated[int, Field(description="Maximum number of records to return, default 100 (最大返回记录数,默认 100)")] = 100, + time_range_start: Annotated[Optional[str], Field(description="Optional UTC start time in ISO8601 (可选的 UTC 开始时间,ISO8601 格式)")] = None, + time_range_end: Annotated[Optional[str], Field(description="Optional UTC end time in ISO8601 (可选的 UTC 结束时间,ISO8601 格式)")] = None, + time_field: Annotated[str, Field(description="Time field for range filtering (时间范围过滤字段名)")] = "@timestamp", + index_name: Annotated[Optional[str], Field(description="Index name for output labeling; omit if specified in ES|QL (用于输出标记的索引名,ES|QL 中已指定时可不填)")] = None, +) -> Annotated[str, Field(description="Raw ES|QL query result as JSON string (原始 ES|QL 查询结果 JSON 字符串)")]: + """Execute a raw ELK ES|QL query and return results. (执行原始 ELK ES|QL 查询并返回结果)""" + input_data = ESQLQueryInput( + query=query, limit=limit, + time_range_start=time_range_start, time_range_end=time_range_end, + time_field=time_field, index_name=index_name, + ) + result = SIEMToolKit.execute_esql(input_data) + return result.model_dump_json() + + def get_current_time() -> Annotated[str, Field(description="Current local time string with UTC (当前本地时间UTC字符串)")]: """Get current system UTC time. (获取当前系统 UTC 时间)""" return datetime.now(timezone.utc).isoformat(timespec='seconds').replace('+00:00', 'Z') @@ -418,6 +454,8 @@ REGISTERED_MCP_TOOLS = [ siem_adaptive_query, siem_keyword_search, siem_discover_index_fields, + siem_execute_spl, + siem_execute_esql, # ThreatIntelligence ti_query, diff --git a/PLUGINS/SIEM/backends.py b/PLUGINS/SIEM/backends.py index 53da8b5..48fb6b9 100644 --- a/PLUGINS/SIEM/backends.py +++ b/PLUGINS/SIEM/backends.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import re from dataclasses import dataclass, field from typing import Any, Literal @@ -19,9 +20,11 @@ from PLUGINS.SIEM.models import ( AdaptiveQueryInput, DiscoveredFieldInfo, DiscoverIndexFieldsOutput, + ESQLQueryInput, FieldStat, KeywordSearchInput, SAMPLE_THRESHOLD, + SPLQueryInput, ) from PLUGINS.SIEM.query_builders import ( build_elk_keyword_clauses, @@ -156,6 +159,40 @@ class ELKQueryBackend: params["aggs"] = aggs return client.search(**params) + @classmethod + def execute_esql_query(cls, input_data: ESQLQueryInput) -> BackendQueryResult: + client = ELKClient.get_client() + query = input_data.query + + if input_data.time_range_start and input_data.time_range_end: + time_clause = ( + f'| WHERE {input_data.time_field} >= "{input_data.time_range_start}"' + f' AND {input_data.time_field} < "{input_data.time_range_end}"' + ) + limit_match = re.search(r'\|\s*LIMIT\s+\d+', query, re.IGNORECASE) + if limit_match: + query = query[:limit_match.start()] + time_clause + " " + query[limit_match.start():] + else: + query = query + " " + time_clause + + if not re.search(r'\|\s*LIMIT\s+\d+', query, re.IGNORECASE): + query = f"{query} | LIMIT {input_data.limit}" + + response = client.esql.query(query=query) + body = response.body + columns = [col["name"] for col in body.get("columns", [])] + rows = body.get("values", []) + raw_records = [dict(zip(columns, row)) for row in rows] + + return BackendQueryResult( + backend="ELK", + index_name=input_data.index_name or "unknown", + total_hits=len(raw_records), + aggregation_fields=[], + statistics=[], + raw_records=raw_records, + ) + class SplunkQueryBackend: backend_name: Literal["ELK", "Splunk"] = "Splunk" @@ -281,3 +318,23 @@ class SplunkQueryBackend: backend=cls.backend_name, index_name=index_name, total_fields=len(discovered), fields=discovered, ) + + @classmethod + def execute_spl_query(cls, input_data: SPLQueryInput) -> BackendQueryResult: + service = SplunkClient.get_service() + if input_data.time_range_start and input_data.time_range_end: + start_time, end_time = parse_time_range(input_data.time_range_start, input_data.time_range_end) + else: + start_time, end_time = None, None + + job = create_and_wait_splunk_job(service, input_data.query, start_time, end_time) + total_hits = int(job["eventCount"]) + + return BackendQueryResult( + backend="Splunk", + index_name=input_data.index_name or "unknown", + total_hits=total_hits, + aggregation_fields=[], + statistics=[], + raw_records=fetch_splunk_records(job, input_data.limit) if total_hits > 0 else [], + ) diff --git a/PLUGINS/SIEM/data_extractors.py b/PLUGINS/SIEM/data_extractors.py index 17bf893..d23931d 100644 --- a/PLUGINS/SIEM/data_extractors.py +++ b/PLUGINS/SIEM/data_extractors.py @@ -78,8 +78,13 @@ def fetch_splunk_top_stats(service, search_query: str, start_time: float, end_ti 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") +def create_and_wait_splunk_job(service, search_query: str, start_time: float | None = None, end_time: float | None = None): + time_kwargs = {} + if start_time is not None: + time_kwargs["earliest_time"] = start_time + if end_time is not None: + time_kwargs["latest_time"] = end_time + job = service.jobs.create(search_query, exec_mode="normal", **time_kwargs) while not job.is_done(): time.sleep(0.2) return job diff --git a/PLUGINS/SIEM/models.py b/PLUGINS/SIEM/models.py index 385f548..16a92f2 100644 --- a/PLUGINS/SIEM/models.py +++ b/PLUGINS/SIEM/models.py @@ -165,6 +165,59 @@ class KeywordSearchInput(BaseModel): raise ValueError("keyword must be a string or a list of strings") +class _RawQueryInput(BaseModel): + query: str = Field(..., description="Raw query string to execute") + index_name: Optional[str] = Field( + default=None, + description="Index/source name for output labeling. If omitted, defaults to 'unknown' in the response.", + ) + limit: int = Field( + default=100, + ge=1, + le=10000, + description="Maximum number of records to return.", + ) + time_range_start: Optional[str] = Field( + default=None, + description="Optional start time. Accepts common datetime strings, normalized to UTC ISO8601.", + ) + time_range_end: Optional[str] = Field( + default=None, + description="Optional end time. Accepts common datetime strings, normalized to UTC ISO8601.", + ) + time_field: str = Field( + default="@timestamp", + description="Field used for time range filtering when time_range_start/end are provided.", + ) + + @model_validator(mode="before") + @classmethod + def normalize_time_range_inputs(cls, data: Any) -> Any: + return normalize_time_range_inputs(data) + + @model_validator(mode="after") + def validate_time_range_order(self): + if self.time_range_start is not None and self.time_range_end is not None: + validate_time_range_order(self.time_range_start, self.time_range_end) + return self + + @field_validator("query") + @classmethod + def validate_query(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("query must not be empty") + return stripped + + +class SPLQueryInput(_RawQueryInput): + pass + + +class ESQLQueryInput(_RawQueryInput): + pass + + class FieldStat(BaseModel): field_name: str = Field(..., description="Name of the field for which statistics are computed") top_values: Dict[Union[str, int], int] = Field( diff --git a/PLUGINS/SIEM/response.py b/PLUGINS/SIEM/response.py index efd50ba..01de38e 100644 --- a/PLUGINS/SIEM/response.py +++ b/PLUGINS/SIEM/response.py @@ -5,10 +5,12 @@ from typing import Any, Optional from PLUGINS.SIEM.backends import BackendQueryResult from PLUGINS.SIEM.models import ( AdaptiveQueryInput, + ESQLQueryInput, KeywordSearchInput, QueryOutput, SAMPLE_COUNT, SAMPLE_THRESHOLD, + SPLQueryInput, ) from PLUGINS.SIEM.registry import get_default_agg_fields @@ -74,3 +76,23 @@ def _get_nested(record: dict[str, Any], field_path: str) -> Any: return _MISSING current = current[segment] return current + + +def build_raw_query_output( + input_data: SPLQueryInput | ESQLQueryInput, + result: BackendQueryResult, + *, + limit: int = 100, +) -> QueryOutput: + records = result.raw_records + return QueryOutput( + backend=result.backend, + index_name=result.index_name, + status="records", + total_hits=result.total_hits, + returned_records=len(records), + truncated=len(records) >= limit, + message=f"Executed raw {result.backend} query against {result.index_name}. Returned {len(records)} records.", + statistics=[], + records=records, + ) diff --git a/PLUGINS/SIEM/tools.py b/PLUGINS/SIEM/tools.py index 61195b5..cd82b69 100644 --- a/PLUGINS/SIEM/tools.py +++ b/PLUGINS/SIEM/tools.py @@ -5,13 +5,15 @@ from PLUGINS.SIEM.models import ( AdaptiveQueryInput, DiscoverIndexFieldsInput, DiscoverIndexFieldsOutput, + ESQLQueryInput, KeywordSearchInput, QueryOutput, SchemaExplorerInput, SchemaIndexSummary, IndexInfo, + SPLQueryInput, ) from PLUGINS.SIEM.registry import get_backend_type, get_default_agg_fields, get_index_info, list_indices -from PLUGINS.SIEM.response import build_query_output +from PLUGINS.SIEM.response import build_query_output, build_raw_query_output class SIEMToolKit: @@ -88,6 +90,31 @@ class SIEMToolKit: return results + @classmethod + def execute_spl(cls, input_data: SPLQueryInput) -> QueryOutput: + """ + Execute a raw Splunk SPL query and return results. + + The `limit` parameter controls the maximum number of records returned (default 100). + Time range is optional; if omitted, Splunk defaults to all time. + """ + query_backend = cls._get_query_backend("Splunk") + backend_result = query_backend.execute_spl_query(input_data) + return build_raw_query_output(input_data, backend_result, limit=input_data.limit) + + @classmethod + def execute_esql(cls, input_data: ESQLQueryInput) -> QueryOutput: + """ + Execute a raw ELK ES|QL query and return results. + + The `limit` parameter controls the maximum number of records returned (default 100). + If the query has no LIMIT clause, one is appended automatically. + Time range is optional; if provided, a WHERE clause is injected into the query. + """ + query_backend = cls._get_query_backend("ELK") + backend_result = query_backend.execute_esql_query(input_data) + return build_raw_query_output(input_data, backend_result, limit=input_data.limit) + @classmethod def discover_index_fields(cls, input_data: DiscoverIndexFieldsInput) -> DiscoverIndexFieldsOutput: """