From ebd0e49478d6d1afc1d462438c07eb8fd14a7bdf Mon Sep 17 00:00:00 2001 From: Dan Dye Date: Tue, 2 Jun 2026 16:05:32 -0400 Subject: [PATCH] Add GEPA optimization for MCP Tools Integrate Pareto-optimized tool definitions and relationship constraints under genetic-pareto framework Update copyright headers to 2026 for touched files optimize(soar-mcp): integrate gepa library optimization suite and refine credentials flow test(soar-mcp): add live connection verification script test(soar-mcp): expand integration tests with live endpoints and parameterized runs docs, refactor: integrate GEPA optimization with Vertex AI Model Garden, clean up submodules and hardcoded paths --- .env.example | 34 ++ .gitmodules | 0 docs/development_guide.md | 33 ++ .../gepa_opt/gepa_optimization_results.json | 23 ++ server/gti/gepa_opt/mcp_dataset.json | 104 ++++++ server/gti/gepa_opt/optimize_gti_mcp.py | 303 ++++++++++++++++++ server/gti/gti_mcp/tools/files.py | 62 +--- server/gti/gti_mcp/tools/netloc.py | 115 ++----- server/gti/gti_mcp/tools/urls.py | 62 +--- .../gepa_opt/gepa_optimization_results.json | 15 + server/secops-soar/gepa_opt/mcp_dataset.json | 238 ++++++++++++++ .../secops-soar/gepa_opt/optimize_soar_mcp.py | 287 +++++++++++++++++ server/secops-soar/gepa_opt/test_live_soar.py | 72 +++++ .../secops-soar/secops_soar_mcp/bindings.py | 19 +- .../secops_soar_mcp/case_management.py | 2 +- .../secops_soar_mcp/http_client.py | 4 +- server/secops-soar/tests/conftest.py | 36 ++- server/secops-soar/tests/tests.py | 45 ++- .../gepa_opt/gepa_optimization_results.json | 17 + server/secops/gepa_opt/mcp_dataset.json | 168 ++++++++++ server/secops/gepa_opt/optimize_secops_mcp.py | 272 ++++++++++++++++ .../secops_mcp/tools/security_alerts.py | 11 +- 22 files changed, 1705 insertions(+), 217 deletions(-) create mode 100644 .env.example delete mode 100644 .gitmodules create mode 100644 server/gti/gepa_opt/gepa_optimization_results.json create mode 100644 server/gti/gepa_opt/mcp_dataset.json create mode 100644 server/gti/gepa_opt/optimize_gti_mcp.py create mode 100644 server/secops-soar/gepa_opt/gepa_optimization_results.json create mode 100644 server/secops-soar/gepa_opt/mcp_dataset.json create mode 100644 server/secops-soar/gepa_opt/optimize_soar_mcp.py create mode 100644 server/secops-soar/gepa_opt/test_live_soar.py create mode 100644 server/secops/gepa_opt/gepa_optimization_results.json create mode 100644 server/secops/gepa_opt/mcp_dataset.json create mode 100644 server/secops/gepa_opt/optimize_secops_mcp.py diff --git a/.env.example b/.env.example new file mode 100644 index 00000000..02254a9a --- /dev/null +++ b/.env.example @@ -0,0 +1,34 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Gemini API Key (used for various Gemini operations/agents) +GEMINI_API_KEY=your_gemini_api_key_here + +# VirusTotal API Key (used by VirusTotal / GTI tools) +VT_APIKEY=your_virustotal_api_key_here + +# Google SecOps SOAR API Keys & Credentials +SOAR_URL=https://your-soar-instance-url.com +SOAR_APP_KEY=your_soar_app_key_here + +# Google Cloud Platform / Chronicle SIEM Credentials (used by Optimizer scripts and SIEM tools) +# Path to your Google Cloud service account JSON file +GOOGLE_APPLICATION_CREDENTIALS=/path/to/your/service-account.json + +# Vertex AI Model Garden Configuration +# GCP Project ID/Name for Vertex AI API execution (e.g., secops-demo-env) +VERTEX_PROJECT=your_gcp_project_id_here + +# GCP Location/Region for Vertex AI execution (e.g., us-central1) +VERTEX_LOCATION=us-central1 diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index e69de29b..00000000 diff --git a/docs/development_guide.md b/docs/development_guide.md index b6a9fd68..71e8ab97 100644 --- a/docs/development_guide.md +++ b/docs/development_guide.md @@ -111,6 +111,39 @@ When documenting tools: 3. Explain the return values and any side effects 4. Provide examples of how to use the tool +## MCP Tool Optimization (GEPA) + +This repository integrates **GEPA (Gradient-based Engine for Prompt Ad-hoc Optimization)** to automatically optimize MCP tool docstrings and descriptions. Since LLMs decide which tool to call based on its description, optimizing these docstrings significantly improves tool-routing accuracy. + +### Running Optimization + +MCP servers supporting GEPA optimization include an optimization package and script under `gepa_opt/`: +- `server/secops/gepa_opt/optimize_secops_mcp.py` +- `server/secops-soar/gepa_opt/optimize_soar_mcp.py` +- `server/gti/gepa_opt/optimize_gti_mcp.py` + +#### Prerequisites + +The optimizer runs against Google Cloud Vertex AI and requires the following environment variables (which can be configured in a `.env` file at the root of the project): +- `GOOGLE_APPLICATION_CREDENTIALS`: Path to your GCP Service Account credentials JSON file. +- `VERTEX_PROJECT`: The GCP project ID for Vertex AI execution. +- `VERTEX_LOCATION`: The GCP location/region for Vertex AI execution (e.g., `us-central1`). + +If any of these are missing, the optimizer scripts will fail fast with a `ValueError`. + +#### Execution + +Run the optimizer script for the target server: +```bash +python server/secops-soar/gepa_opt/optimize_soar_mcp.py +``` + +### How it Works + +1. **Dataset**: A curation of user queries matched with expected tool calls is defined in `mcp_dataset.json` within each `gepa_opt/` directory. +2. **Routing Evaluation**: GEPA executes mock queries against the tools using Vertex AI models, calculating a baseline routing accuracy. +3. **Iterative Optimization**: GEPA generates candidate docstrings, evaluates them on the dataset, and updates the python source files of the tools with the best-performing docstring formulations. + ## Building Documentation The documentation uses Sphinx with MyST Markdown. To build the docs: diff --git a/server/gti/gepa_opt/gepa_optimization_results.json b/server/gti/gepa_opt/gepa_optimization_results.json new file mode 100644 index 00000000..efdfb1b2 --- /dev/null +++ b/server/gti/gepa_opt/gepa_optimization_results.json @@ -0,0 +1,23 @@ +{ + "best_score": 0.9166666666666666, + "optimized_tool_descriptions": { + "tool_description_search_threats": "Search threats in the Google Threat Intelligence platform. Threats are modeled as collections. Once you get collections from this tool, you can use get_collection_report to fetch the full reports and their relationships.", + "tool_description_search_campaigns": "Search threat campaigns in the Google Threat Intelligence platform. Campaigns are modeled as collections.", + "tool_description_search_threat_actors": "Search threat actors in the Google Threat Intelligence platform. Threat actors are modeled as collections.", + "tool_description_search_malware_families": "Search malware families in the Google Threat Intelligence platform. Malware families are modeled as collections.", + "tool_description_search_software_toolkits": "Search software toolkits (or just tools) in the Google Threat Intelligence platform. Software toolkits are modeled as collections.", + "tool_description_search_threat_reports": "Search threat reports in the Google Threat Intelligence platform. Threat reports are modeled as collections.", + "tool_description_search_vulnerabilities": "Search vulnerabilities (CVEs) in the Google Threat Intelligence platform. Vulnerabilities are modeled as collections.", + "tool_description_get_entities_related_to_a_domain": "Retrieve entities related to the given domain. Available relationships: associations, caa_records, campaigns, cname_records, collections, comments, communicating_files, downloaded_files, graphs, historical_ssl_certificates, historical_whois, immediate_parent, malware_families, memory_pattern_parents, mx_records, ns_records, parent, referrer_files, related_comments, related_reports, related_threat_actors, reports, resolutions, siblings, soa_records, software_toolkits, subdomains, urls, user_votes, votes, vulnerabilities.", + "tool_description_get_entities_related_to_an_url": "Retrieve entities related to the given URL. Available relationships: analyses, associations, campaigns, collections, comments, communicating_files, contacted_domains, contacted_ips, downloaded_files, embedded_js_files, last_serving_ip_address, malware_families, parent_resource_urls, redirects_to, referrer_files, referrer_urls.", + "tool_description_get_entities_related_to_an_ip_address": "Retrieve entities related to the given IP address. Available relationships: associations, campaigns, collections, comments, communicating_files, downloaded_files, graphs, historical_ssl_certificates, historical_whois, malware_families, memory_pattern_parents, referrer_files, related_comments, related_reports, related_threat_actors, reports, resolutions, software_toolkits, urls, user_votes, votes, vulnerabilities.", + "tool_description_get_entities_related_to_a_file": "Retrieve entities related to the given file hash. Available relationships: analyses, behaviors, carbonblack_children, carbonblack_parents, compressed_parents, contacted_domains, contacted_ips, contacted_urls, dropped_files, execution_parents, itw_domains, itw_urls, metadata, memory_pattern_domains, memory_pattern_ips, mutexes_created, mutexes_opened, overlay_children, overlay_parents, pcap_parents, pe_resource_children, pe_resource_parents, popular_threat_category, suggested_threat_label, yara_rules.", + "tool_description_get_domain_report": "Get a comprehensive domain analysis report from Google Threat Intelligence. Provides attributes, threat classification, and historical metadata for a domain.", + "tool_description_get_ip_address_report": "Get a comprehensive IP Address analysis report from Google Threat Intelligence. Provides geolocation, autonomous system details, and threat reputation data.", + "tool_description_get_url_report": "Get a comprehensive URL analysis report from Google Threat Intelligence. Provides security analysis, categorizations, and threat category classifications.", + "tool_description_get_file_report": "Get a comprehensive file analysis report using its hash (MD5/SHA-1/SHA-256). Provides detection stats, threat classifications, and static metadata.", + "tool_description_get_file_behavior_summary": "Retrieve a summary of all sandbox execution reports and dynamic analysis details for a file.", + "tool_description_get_file_behavior_report": "Retrieve a full, detailed sandbox behavior report using a behavior ID formatted as {hash}_{sandbox}.", + "tool_description_search_iocs": "Search Indicators of Compromise (IOC) in the Google Threat Intelligence platform using VirusTotal query search modifiers." + } +} \ No newline at end of file diff --git a/server/gti/gepa_opt/mcp_dataset.json b/server/gti/gepa_opt/mcp_dataset.json new file mode 100644 index 00000000..5635d5a8 --- /dev/null +++ b/server/gti/gepa_opt/mcp_dataset.json @@ -0,0 +1,104 @@ +[ + { + "user_query": "Search for threat actor profiles related to APT28", + "tool_arguments": { + "query": "APT28" + }, + "reference_answer": "search_threat_actors", + "additional_context": {} + }, + { + "user_query": "Find campaigns associated with Sandworm", + "tool_arguments": { + "query": "Sandworm" + }, + "reference_answer": "search_campaigns", + "additional_context": {} + }, + { + "user_query": "Show me reports and profiles for the malware family known as Emotet", + "tool_arguments": { + "query": "Emotet" + }, + "reference_answer": "search_malware_families", + "additional_context": {} + }, + { + "user_query": "Look up general threats associated with Apache Log4j CVE-2021-44228 vulnerability", + "tool_arguments": { + "query": "CVE-2021-44228" + }, + "reference_answer": "search_vulnerabilities", + "additional_context": {} + }, + { + "user_query": "Get a comprehensive analysis report for the domain badsite.com", + "tool_arguments": { + "domain": "badsite.com" + }, + "reference_answer": "get_domain_report", + "additional_context": {} + }, + { + "user_query": "Check file reputation for the SHA256 hash 275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f", + "tool_arguments": { + "hash": "275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f" + }, + "reference_answer": "get_file_report", + "additional_context": {} + }, + { + "user_query": "Retrieve sandbox behavior summary for file 275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f", + "tool_arguments": { + "hash": "275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f" + }, + "reference_answer": "get_file_behavior_summary", + "additional_context": {} + }, + { + "user_query": "Show me the sandbox behavior report for id 275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f_Jujubox", + "tool_arguments": { + "file_behaviour_id": "275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f_Jujubox" + }, + "reference_answer": "get_file_behavior_report", + "additional_context": {} + }, + { + "user_query": "Search for indicators of compromise related to port 4444 connection in files", + "tool_arguments": { + "query": "entity:file p:4444" + }, + "reference_answer": "search_iocs", + "additional_context": {} + }, + { + "user_query": "What are the IP addresses contacted by the domain malicious.xyz?", + "tool_arguments": { + "domain": "malicious.xyz", + "relationship_name": "resolutions", + "descriptors_only": true + }, + "reference_answer": "get_entities_related_to_a_domain", + "additional_context": {} + }, + { + "user_query": "Find any communicating files that communicate with the URL http://malicious-link.com/download", + "tool_arguments": { + "url": "http://malicious-link.com/download", + "relationship_name": "communicating_files", + "descriptors_only": true + }, + "reference_answer": "get_entities_related_to_an_url", + "additional_context": {} + }, + { + "user_query": "List comments posted on the IP 8.8.8.8", + "tool_arguments": { + "ip_address": "8.8.8.8", + "relationship_name": "comments", + "descriptors_only": false + }, + "reference_answer": "get_entities_related_to_an_ip_address", + "additional_context": {} + } +] diff --git a/server/gti/gepa_opt/optimize_gti_mcp.py b/server/gti/gepa_opt/optimize_gti_mcp.py new file mode 100644 index 00000000..f022c8e0 --- /dev/null +++ b/server/gti/gepa_opt/optimize_gti_mcp.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +GTI MCP Tool Docstring Optimization using the official GEPA library +Configured for Vertex AI Model Garden via LiteLLM. +""" + +import os +import sys +import json +import logging +from pathlib import Path +import asyncio + +# 1. Apply asyncio subprocess reader limit monkey patch +original_create_subprocess_exec = asyncio.create_subprocess_exec +async def patched_create_subprocess_exec(*args, **kwargs): + kwargs['limit'] = 10 * 1024 * 1024 # 10MB buffer + return await original_create_subprocess_exec(*args, **kwargs) +asyncio.create_subprocess_exec = patched_create_subprocess_exec + +# 2. Load environment variables using dotenv +import dotenv +env_path = Path(__file__).resolve().parents[3] / ".env" +if env_path.exists(): + dotenv.load_dotenv(env_path) + +# 3. Configure Google Cloud Vertex AI credentials for LiteLLM +for var in ["GOOGLE_APPLICATION_CREDENTIALS", "VERTEX_PROJECT", "VERTEX_LOCATION"]: + if not os.getenv(var): + raise ValueError( + f"Missing required environment variable: {var}. " + "Please specify it in your environment or in the .env file." + ) + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# 4. Import GEPA and configure LiteLLM retries +try: + import gepa + from gepa.adapters.mcp_adapter import MCPAdapter + from mcp import StdioServerParameters + import litellm + + # Configure LiteLLM retry and drop parameter behaviors + litellm.num_retries = 5 + litellm.drop_params = True +except ImportError as e: + logger.error("Failed to import GEPA or required dependencies. Please check environment: %s", e) + sys.exit(1) + + +def load_dataset(dataset_path: Path) -> list: + """Loads the evaluation dataset.""" + with open(dataset_path, "r") as f: + return json.load(f) + + +def gti_metric_fn(data_inst, output: str) -> float: + """ + Evaluation Metric: Scores 1.0 if model correctly chooses the expected tool + and specifies all target parameters correctly. + """ + import time + time.sleep(1.5) # Safe spacing to avoid hitting Vertex free-tier RPM limits + + try: + cleaned = output.strip() + if cleaned.startswith("```json"): + cleaned = cleaned[7:] + if cleaned.endswith("```"): + cleaned = cleaned[:-3] + cleaned = cleaned.strip() + parsed = json.loads(cleaned) + except Exception: + return 0.0 + + if isinstance(parsed, list): + if len(parsed) > 0: + parsed = parsed[0] + else: + return 0.0 + + if not isinstance(parsed, dict): + return 0.0 + + if parsed.get("action") != "call_tool": + return 0.0 + + selected_tool = parsed.get("tool") + expected_tool = data_inst.get("reference_answer") + if selected_tool != expected_tool: + return 0.0 + + arguments = parsed.get("arguments", {}) + expected_args = data_inst.get("tool_arguments", {}) + + for k, v in expected_args.items(): + if k not in arguments: + return 0.0 + if isinstance(v, str) and isinstance(arguments[k], str): + if v.lower() not in arguments[k].lower(): + return 0.0 + elif arguments[k] != v: + return 0.0 + + return 1.0 + + +def main(): + # Verify VT_APIKEY environment variable is present + vt_apikey = os.environ.get("VT_APIKEY") + if not vt_apikey: + logger.error("Error: VT_APIKEY environment variable is required for running the local MCP server.") + sys.exit(1) + + opt_dir = Path(__file__).parent + server_dir = opt_dir.parent + server_py = server_dir / "gti_mcp" / "server.py" + dataset_json = opt_dir / "mcp_dataset.json" + + if not server_py.exists(): + logger.error("MCP Server script not found at %s", server_py) + sys.exit(1) + + dataset = load_dataset(dataset_json) + logger.info("Loaded dataset with %d items", len(dataset)) + + # Evolve descriptions for these target tools + raw_tool_names = [ + "search_threats", + "search_campaigns", + "search_threat_actors", + "search_malware_families", + "search_software_toolkits", + "search_threat_reports", + "search_vulnerabilities", + "get_entities_related_to_a_domain", + "get_entities_related_to_an_url", + "get_entities_related_to_an_ip_address", + "get_entities_related_to_a_file", + "get_domain_report", + "get_ip_address_report", + "get_url_report", + "get_file_report", + "get_file_behavior_summary", + "get_file_behavior_report", + "search_iocs" + ] + + # Target Vertex AI Model Garden models + task_model = "vertex_ai/gemini-2.5-flash" + reflection_model = "vertex_ai/gemini-2.5-pro" + + logger.info("Initializing MCPAdapter targeting local stdio server...") + logger.info("Task Model: %s", task_model) + logger.info("Reflection Model: %s", reflection_model) + + adapter = MCPAdapter( + tool_names=raw_tool_names, + task_model=task_model, + metric_fn=gti_metric_fn, + server_params=StdioServerParameters( + command=sys.executable, + args=[str(server_py)], + ), + base_system_prompt=( + "You are a security analyst with access to the Google Threat Intelligence suite. " + "Your goal is to determine which tool to call and what parameters to pass based on user query." + ), + enable_two_pass=False, # Optimize JSON first pass selection logic + failure_score=0.0, + ) + + # Seed candidates with original tool descriptions + seed_candidate = { + "tool_description_search_threats": ( + "Search threats in the Google Threat Intelligence platform. " + "Threats are modeled as collections. Once you get collections from this tool, " + "you can use get_collection_report to fetch the full reports and their relationships." + ), + "tool_description_search_campaigns": ( + "Search threat campaigns in the Google Threat Intelligence platform. " + "Campaigns are modeled as collections." + ), + "tool_description_search_threat_actors": ( + "Search threat actors in the Google Threat Intelligence platform. " + "Threat actors are modeled as collections." + ), + "tool_description_search_malware_families": ( + "Search malware families in the Google Threat Intelligence platform. " + "Malware families are modeled as collections." + ), + "tool_description_search_software_toolkits": ( + "Search software toolkits (or just tools) in the Google Threat Intelligence platform. " + "Software toolkits are modeled as collections." + ), + "tool_description_search_threat_reports": ( + "Search threat reports in the Google Threat Intelligence platform. " + "Threat reports are modeled as collections." + ), + "tool_description_search_vulnerabilities": ( + "Search vulnerabilities (CVEs) in the Google Threat Intelligence platform. " + "Vulnerabilities are modeled as collections." + ), + "tool_description_get_entities_related_to_a_domain": ( + "Retrieve entities related to the given domain. Available relationships: " + "associations, caa_records, campaigns, cname_records, collections, comments, " + "communicating_files, downloaded_files, graphs, historical_ssl_certificates, " + "historical_whois, immediate_parent, malware_families, memory_pattern_parents, " + "mx_records, ns_records, parent, referrer_files, related_comments, related_reports, " + "related_threat_actors, reports, resolutions, siblings, soa_records, " + "software_toolkits, subdomains, urls, user_votes, votes, vulnerabilities." + ), + "tool_description_get_entities_related_to_an_url": ( + "Retrieve entities related to the given URL. Available relationships: " + "analyses, associations, campaigns, collections, comments, communicating_files, " + "contacted_domains, contacted_ips, downloaded_files, embedded_js_files, last_serving_ip_address, " + "malware_families, parent_resource_urls, redirects_to, referrer_files, referrer_urls." + ), + "tool_description_get_entities_related_to_an_ip_address": ( + "Retrieve entities related to the given IP address. Available relationships: " + "associations, campaigns, collections, comments, communicating_files, downloaded_files, " + "graphs, historical_ssl_certificates, historical_whois, malware_families, memory_pattern_parents, " + "referrer_files, related_comments, related_reports, related_threat_actors, reports, " + "resolutions, software_toolkits, urls, user_votes, votes, vulnerabilities." + ), + "tool_description_get_entities_related_to_a_file": ( + "Retrieve entities related to the given file hash. Available relationships: " + "analyses, behaviors, carbonblack_children, carbonblack_parents, compressed_parents, " + "contacted_domains, contacted_ips, contacted_urls, dropped_files, execution_parents, " + "itw_domains, itw_urls, metadata, memory_pattern_domains, memory_pattern_ips, mutexes_created, " + "mutexes_opened, overlay_children, overlay_parents, pcap_parents, pe_resource_children, " + "pe_resource_parents, popular_threat_category, suggested_threat_label, yara_rules." + ), + "tool_description_get_domain_report": ( + "Get a comprehensive domain analysis report from Google Threat Intelligence. " + "Provides attributes, threat classification, and historical metadata for a domain." + ), + "tool_description_get_ip_address_report": ( + "Get a comprehensive IP Address analysis report from Google Threat Intelligence. " + "Provides geolocation, autonomous system details, and threat reputation data." + ), + "tool_description_get_url_report": ( + "Get a comprehensive URL analysis report from Google Threat Intelligence. " + "Provides security analysis, categorizations, and threat category classifications." + ), + "tool_description_get_file_report": ( + "Get a comprehensive file analysis report using its hash (MD5/SHA-1/SHA-256). " + "Provides detection stats, threat classifications, and static metadata." + ), + "tool_description_get_file_behavior_summary": ( + "Retrieve a summary of all sandbox execution reports and dynamic analysis details for a file." + ), + "tool_description_get_file_behavior_report": ( + "Retrieve a full, detailed sandbox behavior report using a behavior ID formatted as {hash}_{sandbox}." + ), + "tool_description_search_iocs": ( + "Search Indicators of Compromise (IOC) in the Google Threat Intelligence platform using VirusTotal query search modifiers." + ) + } + + logger.info("Starting GEPA optimization loop...") + result = gepa.optimize( + seed_candidate=seed_candidate, + trainset=dataset, + valset=dataset, + adapter=adapter, + reflection_lm=reflection_model, + max_metric_calls=60, # iteration and exploration limit bounds + ) + + logger.info("Optimization finished!") + best_candidate = result.candidates[result.best_idx] + best_score = result.val_aggregate_scores[result.best_idx] + logger.info("Best achieved score: %.2f", best_score) + + # Save results + result_path = opt_dir / "gepa_optimization_results.json" + with open(result_path, "w") as f: + json.dump({ + "best_score": best_score, + "optimized_tool_descriptions": best_candidate + }, f, indent=2) + logger.info("Optimized results successfully saved to %s", result_path) + + +if __name__ == "__main__": + main() diff --git a/server/gti/gti_mcp/tools/files.py b/server/gti/gti_mcp/tools/files.py index 6c7c6650..ef3c8c34 100644 --- a/server/gti/gti_mcp/tools/files.py +++ b/server/gti/gti_mcp/tools/files.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -111,59 +111,13 @@ async def get_file_report(hash: str, ctx: Context) -> typing.Dict[str, typing.An async def get_entities_related_to_a_file( hash: str, relationship_name: str, descriptors_only: bool, ctx: Context, limit: int = 10, ) -> list[dict[str, typing.Any]]: - """Retrieve entities related to the the given file hash. - - The following table shows a summary of available relationships for file objects. - - | Relationship | Description | Return type | - | ---------------------- | --------------------------------------------------------------------------------- | ----------- | - | analyses | Analyses for the file | analysis | - | associations | File's associated objects (reports, campaigns, IoC collections, malware families, software toolkits, vulnerabilities, threat-actors), without filtering by the associated object type. | collection | - | behaviours | Behaviour reports for the file. | file-behaviour | - | attack_techniques | Returns the Attack Techniques of the File. | attack_technique | - | bundled_files | Files bundled within the file. | file | - | campaigns | Campaigns associated to the file. | collection | - | carbonblack_children | Files derived from the file according to Carbon Black. | file | - | carbonblack_parents | Files from where the file was derived according to Carbon Black. | file | - | collections | IoC Collections associated to the file. | collection | - | comments | Comments for the file. | comment | - | compressed_parents | Compressed files that contain the file. | file | - | contacted_domains | Domains contacted by the file. | domain | - | contacted_ips | IP addresses contacted by the file. | ip_address | - | contacted_urls | URLs contacted by the file. | url | - | dropped_files | Files dropped by the file during its execution. | file | - | email_attachments | Files attached to the email. | file | - | email_parents | Email files that contained the file. | file | - | embedded_domains | Domain names embedded in the file. | domain | - | embedded_ips | IP addresses embedded in the file. | ip_address | - | embedded_urls | URLs embedded in the file. | url | - | execution_parents | Files that executed the file. | file | - | graphs | Graphs that include the file. | graph | - | itw_domains | In the wild domain names from where the file has been downloaded. | domain | - | itw_ips | In the wild IP addresses from where the file has been downloaded. | ip_address | - | itw_urls | In the wild URLs from where the file has been downloaded. | url | - | malware_families | Malware families associated to the file. | collection | - | memory_pattern_domains | Domain string patterns found in memory during sandbox execution. | domain | - | memory_pattern_ips | IP address string patterns found in memory during sandbox execution. | ip_address | - | memory_pattern_urls | URL string patterns found in memory during sandbox execution. | url | - | overlay_children | Files contained by the file as an overlay. | file | - | overlay_parents | File that contain the file as an overlay. | file | - | pcap_children | Files contained within the PCAP file. | file | - | pcap_parents | PCAP files that contain the file. | file | - | pe_resource_children | Files contained by a PE file as a resource. | file | - | pe_resource_parents | PE files containing the file as a resource. | file | - | related_attack_techniques | Returns the Attack Techniques of the Collections containing this File. | attack_technique | - | related_reports | Reports that are directly and indirectly related to the file. | collection | - | related_threat_actors | File's related threat actors. | collection | - | reports | Reports directly associated to the file. | collection | - | screenshots | Screenshots related to the sandbox execution of the file. | screenshot | - | similar_files | Files that are similar to the file. | file | - | software_toolkits | Software and Toolkits associated to the file. | collection | - | submissions | Submissions for the file. | submission | - | urls_for_embedded_js | URLs where this (JS) file is embedded. | url | - | user_votes | File's votes made by current signed-in user. | vote | - | votes | Votes for the file. | vote | - | vulnerabilities | Vulnerabilities associated to the file. | collection | + """Retrieve entities related to the given file hash. + + Available relationships: analyses, behaviors, carbonblack_children, carbonblack_parents, + compressed_parents, contacted_domains, contacted_ips, contacted_urls, dropped_files, + execution_parents, itw_domains, itw_urls, metadata, memory_pattern_domains, memory_pattern_ips, + mutexes_created, mutexes_opened, overlay_children, overlay_parents, pcap_parents, pe_resource_children, + pe_resource_parents, popular_threat_category, suggested_threat_label, yara_rules. Args: hash (required): MD5/SHA1/SHA256) hash that identifies the file. diff --git a/server/gti/gti_mcp/tools/netloc.py b/server/gti/gti_mcp/tools/netloc.py index ceda4a54..dcae377d 100644 --- a/server/gti/gti_mcp/tools/netloc.py +++ b/server/gti/gti_mcp/tools/netloc.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -112,51 +112,23 @@ async def get_domain_report(domain: str, ctx: Context) -> typing.Dict[str, typin async def get_entities_related_to_a_domain( domain: str, relationship_name: str, descriptors_only: bool, ctx: Context, limit: int = 10 ) -> list[dict[str, typing.Any]]: - """Retrieve entities related to the the given domain. - - The following table shows a summary of available relationships for domain objects. - - | Relationship | Description | Return type | - | --------------------------- | ---------------------------------------------------------- | ------------ | - | associations | Domain's associated objects (reports, campaigns, IoC collections, malware families, software toolkits, vulnerabilities, threat-actors), without filtering by the associated object type. | Everyone. | List of [reports](ref:report-object), [campaigns](ref:campaign-object), [IoC collections](ref:ioc-collection-object), [malware families](ref:malware-family-object), [software toolkits](ref:software-toolkit-object), [vulnerabilities](ref:vulnerability-object), [threat-actors](ref:threat-actor-object) objecs.| collection | - | caa_records | Records CAA for the domain. | domain | - | campaigns | Campaigns associated to the domain. | collection | - | cname_records | Records CNAME for the domain. | domain | - | collections | IoC Collections associated to the domain. | collection | - | comments | Community posted comments about the domain. | comment | - | communicating_files | Files that communicate with the domain. | file | - | downloaded_files | Files downloaded from that domain. | file | - | graphs | Graphs including the domain. | graph | - | historical_ssl_certificates | SSL certificates associated with the domain. | ssl-cert | - | historical_whois | WHOIS information for the domain. | whois | - | immediate_parent | Domain's immediate parent. | domain | - | malware_families | Malware families associated to the domain. | collection | - | memory_pattern_parents | Files having a domain as string on memory during sandbox execution. | file | - | mx_records | Records MX for the domain. | domain | - | ns_records | Records NS for the domain. | domain | - | parent | Domain's top parent. | domain | - | referrer_files | Files containing the domain. | file | - | related_comments | Community posted comments in the domain's related objects. | comment | - | related_reports | Reports that are directly and indirectly related to the domain. | collection | - | related_threat_actors | Threat actors related to the domain. | collection | - | reports | Reports directly associated to the domain. | collection | - | resolutions | DNS resolutions for the domain. | resolution | - | siblings | Domain's sibling domains. | domain | - | soa_records | Records SOA for the domain. | domain | - | software_toolkits | Software and Toolkits associated to the domain. | collection | - | subdomains | Domain's subdomains. | domain | - | urls | URLs having this domain. | url | - | user_votes | Current user's votes. | vote | - | votes | Domain's votes. | vote | - | vulnerabilities | Vulnerabilities associated to the domain. | collection | - - Args: - domain (required): Domain to analyse. - relationship_name (required): Relationship name. - descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. - limit: Limit the number of entities to retrieve. 10 by default. - Returns: - List of entities related to the domain. + """Retrieve entities related to the given domain. + + Available relationships: associations, caa_records, campaigns, cname_records, + collections, comments, communicating_files, downloaded_files, graphs, + historical_ssl_certificates, historical_whois, immediate_parent, malware_families, + memory_pattern_parents, mx_records, ns_records, parent, referrer_files, + related_comments, related_reports, related_threat_actors, reports, resolutions, + siblings, soa_records, software_toolkits, subdomains, urls, user_votes, votes, + vulnerabilities. + + Args: + domain (required): Domain to analyse. + relationship_name (required): Relationship name. + descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. + limit: Limit the number of entities to retrieve. 10 by default. + Returns: + List of entities related to the domain. """ if not relationship_name in DOMAIN_RELATIONSHIPS: return { @@ -197,42 +169,21 @@ async def get_ip_address_report(ip_address: str, ctx: Context) -> typing.Dict[st async def get_entities_related_to_an_ip_address( ip_address: str, relationship_name: str, descriptors_only: bool, ctx: Context, limit: int = 10 ) -> list[dict[str, typing.Any]]: - """Retrieve entities related to the the given IP Address. - - The following table shows a summary of available relationships for IP Address objects. - - | Relationship | Description | Return type | - | --------------------------- | ------------------------------------------------------ | ------------ | - | associations | IP's associated objects (reports, campaigns, IoC collections, malware families, software toolkits, vulnerabilities, threat-actors), without filtering by the associated object type. | collection | - | campaigns | Campaigns associated to the IP address. | collection | - | collections | IoC Collections associated to the IP address. | collection | - | comments | Comments for the IP address. | comment | - | communicating_files | Files that communicate with the IP address. | file | - | downloaded_files | Files downloaded from the IP address. | file | - | graphs | Graphs including the IP address. | graph | - | historical_ssl_certificates | SSL certificates associated with the IP. | ssl-cert | - | historical_whois | WHOIS information for the IP address. | whois | - | malware_families | Malware families associated to the IP address. | collection | - | memory_pattern_parents | Files having an IP as string on memory during sandbox execution. | file | - | referrer_files | Files containing the IP address. | file | - | related_comments | Community posted comments in the IP's related objects. | comment | - | related_reports | Reports that are directly and indirectly related to the IP. | collection | - | related_threat_actors | Threat actors related to the IP address. | collection | - | reports | Reports directly associated to the IP. | collection | - | resolutions | IP address' resolutions | resolution | - | software_toolkits | Software and Toolkits associated to the IP address. | collection | - | urls | URLs related to the IP address. | url | - | user_votes | IP's votes made by current signed-in user. | vote | - | votes | IP's votes. | vote | - | vulnerabilities | Vulnerabilities associated to the IP address. | collection | - - Args: - ip_address (required): IP Addres to analyse. - relationship_name (required): Relationship name. - descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. - limit: Limit the number of entities to retrieve. 10 by default. - Returns: - List of entities related to the IP Address. + """Retrieve entities related to the given IP address. + + Available relationships: associations, campaigns, collections, comments, + communicating_files, downloaded_files, graphs, historical_ssl_certificates, + historical_whois, malware_families, memory_pattern_parents, referrer_files, + related_comments, related_reports, related_threat_actors, reports, + resolutions, software_toolkits, urls, user_votes, votes, vulnerabilities. + + Args: + ip_address (required): IP Address to analyse. + relationship_name (required): Relationship name. + descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. + limit: Limit the number of entities to retrieve. 10 by default. + Returns: + List of entities related to the IP Address. """ if not relationship_name in IP_RELATIONSHIPS: return { diff --git a/server/gti/gti_mcp/tools/urls.py b/server/gti/gti_mcp/tools/urls.py index e88f5254..cb99efcb 100644 --- a/server/gti/gti_mcp/tools/urls.py +++ b/server/gti/gti_mcp/tools/urls.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -94,52 +94,20 @@ async def get_url_report(url: str, ctx: Context) -> typing.Dict[str, typing.Any] async def get_entities_related_to_an_url( url: str, relationship_name: str, descriptors_only: bool, ctx: Context, limit: int = 10 ) -> list[dict[str, typing.Any]]: - """Retrieve entities related to the the given URL. - - The following table shows a summary of available relationships for URL objects. - - | Relationship | Description | Return type | - | ----------------------- | -------------------------------------------------------------- | ------------ | - | analyses | Analyses for the URL. | analyse | - | associations | URL's associated objects (reports, campaigns, IoC collections, malware families, software toolkits, vulnerabilities, threat-actors), without filtering by the associated object type. | collection | - | campaigns | Campaigns associated to the URL. | collection | - | collections | IoC Collections associated to the URL. | collection | - | comments | Community posted comments about the URL. | comment | - | communicating_files | Files that communicate with a given URL when they're executed. | file | - | contacted_domains | Domains from which the URL loads some kind of resource. | domain | - | contacted_ips | IPs from which the URL loads some kind of resource. | ip_address | - | downloaded_files | Files downloaded from the URL. | file | - | embedded_js_files | JS files embedded in a URL. | file | - | graphs | Graphs including the URL. | graph | - | http_response_contents | HTTP response contents from the URL. | file | - | last_serving_ip_address | Last IP address that served the URL. | ip_address | - | malware_families | Malware families associated to the URL. | collection | - | memory_pattern_parents | Files having a domain as string on memory during sandbox execution. | file | - | network_location | Domain or IP for the URL. | domain or ip_address | - | parent_resource_urls | Returns the URLs where this URL has been loaded as resource. | url | - | redirecting_urls | URLs that redirected to the given URL. | url | - | redirects_to | URLs that this url redirects to. | url | - | referrer_files | Files containing the URL. | file | - | referrer_urls | URLs referring the URL. | url | - | related_collections | Returns the Collections of the parent Domains or IPs of this URL. | collection | - | related_comments | Community posted comments in the URL's related objects. | comment | - | related_reports | Reports that are directly and indirectly related to the URL. | collection | - | related_threat_actors | URL's related threat actors. | collection | - | reports | Reports directly associated to the URL. | collection | - | software_toolkits | Software and Toolkits associated to the URL. | collection | - | submissions | URL's submissions. | url | - | urls_related_by_tracker_id | URLs that share the same tracker ID. | url | - | user_votes | URL's votes made by current signed-in user. | vote | - | votes | Votes for the URL. | vote | - | vulnerabilities | Vulnerabilities associated to the URL. | collection | - - Args: - url (required): URL to analyse. - relationship_name (required): Relationship name. - descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. - limit: Limit the number of objects to retrieve. 10 by default. - Returns: - List of entities related to the URL. + """Retrieve entities related to the given URL. + + Available relationships: analyses, associations, campaigns, collections, comments, + communicating_files, contacted_domains, contacted_ips, downloaded_files, embedded_js_files, + last_serving_ip_address, malware_families, parent_resource_urls, redirects_to, referrer_files, + referrer_urls. + + Args: + url (required): URL to analyse. + relationship_name (required): Relationship name. + descriptors_only (required): Bool. Must be True when the target object type is one of file, domain, url, ip_address or collection. + limit: Limit the number of objects to retrieve. 10 by default. + Returns: + List of entities related to the URL. """ if not relationship_name in URL_RELATIONSHIPS: return { diff --git a/server/secops-soar/gepa_opt/gepa_optimization_results.json b/server/secops-soar/gepa_opt/gepa_optimization_results.json new file mode 100644 index 00000000..634a3415 --- /dev/null +++ b/server/secops-soar/gepa_opt/gepa_optimization_results.json @@ -0,0 +1,15 @@ +{ + "best_score": 1.0, + "optimized_tool_descriptions": { + "tool_description_list_cases": "List cases available in the Security Orchestration, Automation, and Response (SOAR) platform. In a SOAR context, a 'case' typically represents a security incident, investigation, or a container for related alerts and response actions. Listing cases provides an overview of ongoing or past security events being managed by the platform. This is useful for getting a high-level list of recent security issues or finding a specific incident to investigate further.", + "tool_description_post_case_comment": "Post a comment to a specific case within the SOAR platform. Cases are used to track security incidents and investigations. Adding comments is essential for documenting findings, communication between analysts, recording actions taken, or providing updates on the investigation progress.", + "tool_description_list_alerts_by_case": "List the security alerts associated with a specific case ID in the SOAR platform. Alerts are notifications generated by security tools (like SIEMs, EDRs) indicating potential security issues. In SOAR, alerts are often grouped into cases for investigation and response. Listing alerts for a case helps understand the scope of the incident, the specific events that triggered it, and the evidence collected.", + "tool_description_list_alert_group_identifiers_by_case": "List alert group identifiers associated with a specific case ID in the SOAR platform. In this SOAR implementation, alerts within a case can be grouped using identifiers, potentially for correlation, playbook execution stages, or analyst assignment. Retrieving these identifiers helps understand the internal structure of a case or target specific alert groupings for automation or analysis.", + "tool_description_list_events_by_alert": "List the underlying security events associated with a specific alert within a given case. Security alerts (often derived from detection rules or IoC matches) are typically triggered by one or more underlying events ingested into the security platform (e.g., Chronicle). These events provide the raw data (likely in UDM format) needed to validate the alert, understand the specific activity, and perform deep-dive investigations.", + "tool_description_change_case_priority": "Change the priority level of a specific case in the SOAR platform. Case priority (e.g., PriorityUnspecified, PriorityInfo, PriorityLow, PriorityMedium, PriorityHigh, PriorityCritical) helps security teams triage incidents and focus on the most urgent threats based on the *currently available information*. Remember that priority can change as more context is gathered during the investigation. The priority might be adjusted during the investigation lifecycle based on new findings.", + "tool_description_get_entities_by_alert_group_identifiers": "Retrieve entities (e.g., IP addresses, hostnames, users) involved in specific alert groups within a case. Understanding which entities are associated with alerts is fundamental for incident investigation and response. This tool allows fetching entities linked to one or more alert groups, which can be crucial for identifying affected assets, potential attack vectors, or compromised accounts. The description also notes it can be used to get target entities for manual actions, implying these entities might be inputs for subsequent response playbooks or manual interventions.", + "tool_description_get_entity_details": "Fetch detailed information about a specific entity known to the SOAR platform. Entities (like IPs, domains, users, assets) are central to security investigations. This tool retrieves comprehensive details about a specific entity based on its identifier, type, and environment. This information might include enrichment data (e.g., threat intelligence, asset inventory details), related alerts or cases, observed activity, and risk scores, providing crucial context for analysis.", + "tool_description_search_entity": "Search for entities within the SOAR platform based on various criteria. This tool provides a flexible way to find entities (assets, users, IOCs, etc.) matching specific attributes. It allows searching by term (e.g., part of a hostname), entity type, suspicion status, asset status (internal/external), enrichment status, network, or environment. This is useful for exploring the entity database, finding potentially related entities during an investigation, or identifying assets with specific characteristics.", + "tool_description_get_case_full_details": "Retrieve comprehensive details for a specific case by aggregating its core information, associated alerts, and comments. This tool provides a consolidated view of a security case by fetching its primary details (like status, priority, description), all linked security alerts, and the full history of comments added by analysts or automation. This aggregated information is essential for getting a complete understanding of an incident's context, scope, investigation progress, and collaborative notes without making multiple separate API calls." + } +} \ No newline at end of file diff --git a/server/secops-soar/gepa_opt/mcp_dataset.json b/server/secops-soar/gepa_opt/mcp_dataset.json new file mode 100644 index 00000000..40001a68 --- /dev/null +++ b/server/secops-soar/gepa_opt/mcp_dataset.json @@ -0,0 +1,238 @@ +[ + { + "user_query": "Show me a list of recent security incidents or cases in the SOAR queue", + "reference_answer": "list_cases", + "tool_arguments": {} + }, + { + "user_query": "Get the latest cases from Chronicle SOAR", + "reference_answer": "list_cases", + "tool_arguments": {} + }, + { + "user_query": "Add a comment to case 523 saying 'We are starting investigation on this indicator.'", + "reference_answer": "post_case_comment", + "tool_arguments": { + "case_id": "523", + "comment": "We are starting investigation on this indicator." + } + }, + { + "user_query": "Post a note to case 812 with details: Host has been isolated successfully.", + "reference_answer": "post_case_comment", + "tool_arguments": { + "case_id": "812", + "comment": "Host has been isolated successfully." + } + }, + { + "user_query": "What alerts are associated with case ID 459?", + "reference_answer": "list_alerts_by_case", + "tool_arguments": { + "case_id": "459" + } + }, + { + "user_query": "List all alerts triggered under case 204", + "reference_answer": "list_alerts_by_case", + "tool_arguments": { + "case_id": "204" + } + }, + { + "user_query": "Show alert group identifiers for case 910", + "reference_answer": "list_alert_group_identifiers_by_case", + "tool_arguments": { + "case_id": "910" + } + }, + { + "user_query": "Get the raw events linked to alert 334 in case 120", + "reference_answer": "list_events_by_alert", + "tool_arguments": { + "case_id": "120", + "alert_id": "334" + } + }, + { + "user_query": "List all security events for alert ID 778 inside case 560", + "reference_answer": "list_events_by_alert", + "tool_arguments": { + "case_id": "560", + "alert_id": "778" + } + }, + { + "user_query": "Set case priority for case 102 to PriorityCritical", + "reference_answer": "change_case_priority", + "tool_arguments": { + "case_id": "102", + "case_priority": "PriorityCritical" + } + }, + { + "user_query": "Change priority of case 705 to PriorityLow because it was a false positive", + "reference_answer": "change_case_priority", + "tool_arguments": { + "case_id": "705", + "case_priority": "PriorityLow" + } + }, + { + "user_query": "Fetch all entities involved in alert group 'rule_group_abc_123' inside case 880", + "reference_answer": "get_entities_by_alert_group_identifiers", + "tool_arguments": { + "case_id": "880", + "alert_group_identifiers": ["rule_group_abc_123"] + } + }, + { + "user_query": "Fetch SOAR details for entity '192.168.1.100' of type 'IP Address' in the 'Production' environment", + "reference_answer": "get_entity_details", + "tool_arguments": { + "entity_identifier": "192.168.1.100", + "entity_type": "IP Address", + "entity_environment": "Production" + } + }, + { + "user_query": "Search the SOAR database for suspicious entities with term 'victim-host'", + "reference_answer": "search_entity", + "tool_arguments": { + "term": "victim-host", + "is_suspicious": true + } + }, + { + "user_query": "Find all enriched internal assets in our network", + "reference_answer": "search_entity", + "tool_arguments": { + "is_internal_asset": true, + "is_enriched": true + } + }, + { + "user_query": "Retrieve the full aggregated details, alerts, and comments for case 144", + "reference_answer": "get_case_full_details", + "tool_arguments": { + "case_id": "144" + } + }, + { + "user_query": "I need all info including history, alerts and comments for case ID 502", + "reference_answer": "get_case_full_details", + "tool_arguments": { + "case_id": "502" + } + }, + { + "user_query": "Fetch the next batch of SOAR cases using token 'abc123token'", + "reference_answer": "list_cases", + "tool_arguments": { + "next_page_token": "abc123token" + } + }, + { + "user_query": "Get the next page of alerts for case 459 using token 'page-token-789'", + "reference_answer": "list_alerts_by_case", + "tool_arguments": { + "case_id": "459", + "next_page_token": "page-token-789" + } + }, + { + "user_query": "Load more group IDs for case 910 with token 'tok_910_xyz'", + "reference_answer": "list_alert_group_identifiers_by_case", + "tool_arguments": { + "case_id": "910", + "next_page_token": "tok_910_xyz" + } + }, + { + "user_query": "Retrieve more events for alert 334 in case 120 using pagination token 'evt-page-112'", + "reference_answer": "list_events_by_alert", + "tool_arguments": { + "case_id": "120", + "alert_id": "334", + "next_page_token": "evt-page-112" + } + }, + { + "user_query": "Set the priority level of case 224 to PriorityHigh", + "reference_answer": "change_case_priority", + "tool_arguments": { + "case_id": "224", + "case_priority": "PriorityHigh" + } + }, + { + "user_query": "Update case 909 priority to PriorityInfo", + "reference_answer": "change_case_priority", + "tool_arguments": { + "case_id": "909", + "case_priority": "PriorityInfo" + } + }, + { + "user_query": "Which entities are involved in alert groups 'detect_malware_01' and 'exfil_stage_02' inside case 712?", + "reference_answer": "get_entities_by_alert_group_identifiers", + "tool_arguments": { + "case_id": "712", + "alert_group_identifiers": [ + "detect_malware_01", + "exfil_stage_02" + ] + } + }, + { + "user_query": "Get SOAR details for user entity 'admin-user' of type 'User' under the 'Staging' environment", + "reference_answer": "get_entity_details", + "tool_arguments": { + "entity_identifier": "admin-user", + "entity_type": "User", + "entity_environment": "Staging" + } + }, + { + "user_query": "Find detailed information for entity 'workstation-99' (type 'Hostname') in the 'Development' environment", + "reference_answer": "get_entity_details", + "tool_arguments": { + "entity_identifier": "workstation-99", + "entity_type": "Hostname", + "entity_environment": "Development" + } + }, + { + "user_query": "Search for all entities of type 'IP Address' or 'Hostname' in the 'Production' and 'DMZ' environments that are suspicious", + "reference_answer": "search_entity", + "tool_arguments": { + "type": [ + "IP Address", + "Hostname" + ], + "environment_name": [ + "Production", + "DMZ" + ], + "is_suspicious": true + } + }, + { + "user_query": "Show me all internal assets that have not been enriched yet", + "reference_answer": "search_entity", + "tool_arguments": { + "is_internal_asset": true, + "is_enriched": false + } + }, + { + "user_query": "Search for entities in network 'Intranet-Zone' matching term 'database'", + "reference_answer": "search_entity", + "tool_arguments": { + "term": "database", + "network_name": [ + "Intranet-Zone" + ] + } + } +] diff --git a/server/secops-soar/gepa_opt/optimize_soar_mcp.py b/server/secops-soar/gepa_opt/optimize_soar_mcp.py new file mode 100644 index 00000000..2a2a3b89 --- /dev/null +++ b/server/secops-soar/gepa_opt/optimize_soar_mcp.py @@ -0,0 +1,287 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +SecOps SOAR MCP Tool Docstring Optimization using the official GEPA library. +Configured for Vertex AI Model Garden via LiteLLM. +""" + +import os +import sys +import json +import logging +from pathlib import Path +import asyncio + +# 1. Apply asyncio subprocess reader limit monkey patch +original_create_subprocess_exec = asyncio.create_subprocess_exec +async def patched_create_subprocess_exec(*args, **kwargs): + kwargs['limit'] = 10 * 1024 * 1024 # 10MB buffer + return await original_create_subprocess_exec(*args, **kwargs) +asyncio.create_subprocess_exec = patched_create_subprocess_exec + +# 2. Load environment variables using dotenv +import dotenv +env_path = Path(__file__).resolve().parents[3] / ".env" +if env_path.exists(): + dotenv.load_dotenv(env_path) + +# 3. Configure Google Cloud Vertex AI credentials for LiteLLM +for var in ["GOOGLE_APPLICATION_CREDENTIALS", "VERTEX_PROJECT", "VERTEX_LOCATION"]: + if not os.getenv(var): + raise ValueError( + f"Missing required environment variable: {var}. " + "Please specify it in your environment or in the .env file." + ) + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# 4. Import GEPA and configure LiteLLM retries +try: + import gepa + from gepa.adapters.mcp_adapter import MCPAdapter + from mcp import StdioServerParameters + import litellm + + # Configure LiteLLM retry and drop parameter behaviors + litellm.num_retries = 5 + litellm.drop_params = True +except ImportError as e: + logger.error("Failed to import GEPA or required dependencies: %s", e) + sys.exit(1) + + +def load_dataset(dataset_path: Path) -> list: + """Loads the evaluation dataset.""" + with open(dataset_path, "r") as f: + return json.load(f) + + +def soar_metric_fn(data_inst, output: str) -> float: + """ + Evaluation Metric: Scores 1.0 if the model correctly selects the expected tool + and extracts all targeted parameters correctly. + """ + import time + time.sleep(1.5) # Safe spacing to avoid hitting Vertex free-tier RPM limits + + try: + cleaned = output.strip() + if cleaned.startswith("```json"): + cleaned = cleaned[7:] + if cleaned.endswith("```"): + cleaned = cleaned[:-3] + cleaned = cleaned.strip() + parsed = json.loads(cleaned) + except Exception: + return 0.0 + + if isinstance(parsed, list): + if len(parsed) > 0: + parsed = parsed[0] + else: + return 0.0 + + if not isinstance(parsed, dict): + return 0.0 + + if parsed.get("action") != "call_tool": + return 0.0 + + selected_tool = parsed.get("tool") + expected_tool = data_inst.get("reference_answer") + if selected_tool != expected_tool: + return 0.0 + + arguments = parsed.get("arguments", {}) + expected_args = data_inst.get("tool_arguments", {}) + + for k, v in expected_args.items(): + if k not in arguments: + return 0.0 + + actual_val = arguments[k] + if isinstance(v, list) and isinstance(actual_val, list): + if len(v) != len(actual_val) or sorted(v) != sorted(actual_val): + return 0.0 + elif isinstance(v, str) and isinstance(actual_val, str): + if v.lower() not in actual_val.lower(): + return 0.0 + elif actual_val != v: + return 0.0 + + return 1.0 + + +def main(): + opt_dir = Path(__file__).parent + server_dir = opt_dir.parent + server_py = server_dir / "secops_soar_mcp" / "server.py" + dataset_json = opt_dir / "mcp_dataset.json" + + if not server_py.exists(): + logger.error("MCP Server script not found at %s", server_py) + sys.exit(1) + + dataset = load_dataset(dataset_json) + logger.info("Loaded dataset with %d items", len(dataset)) + + # The SOAR tools we wish to optimize + raw_tool_names = [ + "list_cases", + "post_case_comment", + "list_alerts_by_case", + "list_alert_group_identifiers_by_case", + "list_events_by_alert", + "change_case_priority", + "get_entities_by_alert_group_identifiers", + "get_entity_details", + "search_entity", + "get_case_full_details" + ] + + # Target Vertex AI Model Garden models + task_model = "vertex_ai/gemini-2.5-flash" + reflection_model = "vertex_ai/gemini-2.5-pro" + + logger.info("Initializing MCPAdapter targeting local stdio server...") + logger.info("Task Model: %s", task_model) + logger.info("Reflection Model: %s", reflection_model) + + adapter = MCPAdapter( + tool_names=raw_tool_names, + task_model=task_model, + metric_fn=soar_metric_fn, + server_params=StdioServerParameters( + command=sys.executable, + args=[str(server_py)], + ), + base_system_prompt=( + "You are a security analyst with access to the Chronicle SecOps SOAR platform. " + "Your goal is to determine which case management or incident response tool to call " + "and what parameters to pass based on the user's query." + ), + enable_two_pass=False, + failure_score=0.0, + ) + + # Seed candidates with original tool descriptions + seed_candidate = { + "tool_description_list_cases": ( + "List cases available in the Security Orchestration, Automation, and Response (SOAR) platform. " + "In a SOAR context, a 'case' typically represents a security incident, investigation, " + "or a container for related alerts and response actions. Listing cases provides an " + "overview of ongoing or past security events being managed by the platform. " + "This is useful for getting a high-level list of recent security issues or finding " + "a specific incident to investigate further." + ), + "tool_description_post_case_comment": ( + "Post a comment to a specific case within the SOAR platform. " + "Cases are used to track security incidents and investigations. Adding comments " + "is essential for documenting findings, communication between analysts, recording " + "actions taken, or providing updates on the investigation progress." + ), + "tool_description_list_alerts_by_case": ( + "List the security alerts associated with a specific case ID in the SOAR platform. " + "Alerts are notifications generated by security tools (like SIEMs, EDRs) indicating " + "potential security issues. In SOAR, alerts are often grouped into cases for " + "investigation and response. Listing alerts for a case helps understand the " + "scope of the incident, the specific events that triggered it, and the evidence collected." + ), + "tool_description_list_alert_group_identifiers_by_case": ( + "List alert group identifiers associated with a specific case ID in the SOAR platform. " + "In this SOAR implementation, alerts within a case can be grouped using identifiers, " + "potentially for correlation, playbook execution stages, or analyst assignment. " + "Retrieving these identifiers helps understand the internal structure of a case " + "or target specific alert groupings for automation or analysis." + ), + "tool_description_list_events_by_alert": ( + "List the underlying security events associated with a specific alert within a given case. " + "Security alerts (often derived from detection rules or IoC matches) are typically " + "triggered by one or more underlying events ingested into the security platform " + "(e.g., Chronicle). These events provide the raw data (likely in UDM format) " + "needed to validate the alert, understand the specific activity, and perform deep-dive investigations." + ), + "tool_description_change_case_priority": ( + "Change the priority level of a specific case in the SOAR platform. " + "Case priority (e.g., PriorityUnspecified, PriorityInfo, PriorityLow, PriorityMedium, " + "PriorityHigh, PriorityCritical) helps security teams triage incidents and focus " + "on the most urgent threats based on the *currently available information*. Remember that priority can change as more context is " + "gathered during the investigation. The priority might be adjusted during the " + "investigation lifecycle based on new findings." + ), + "tool_description_get_entities_by_alert_group_identifiers": ( + "Retrieve entities (e.g., IP addresses, hostnames, users) involved in specific alert groups within a case. " + "Understanding which entities are associated with alerts is fundamental for incident " + "investigation and response. This tool allows fetching entities linked to one or " + "more alert groups, which can be crucial for identifying affected assets, potential " + "attack vectors, or compromised accounts. The description also notes it can be used " + "to get target entities for manual actions, implying these entities might be inputs " + "for subsequent response playbooks or manual interventions." + ), + "tool_description_get_entity_details": ( + "Fetch detailed information about a specific entity known to the SOAR platform. " + "Entities (like IPs, domains, users, assets) are central to security investigations. " + "This tool retrieves comprehensive details about a specific entity based on its " + "identifier, type, and environment. This information might include enrichment data " + "(e.g., threat intelligence, asset inventory details), related alerts or cases, " + "observed activity, and risk scores, providing crucial context for analysis." + ), + "tool_description_search_entity": ( + "Search for entities within the SOAR platform based on various criteria. " + "This tool provides a flexible way to find entities (assets, users, IOCs, etc.) " + "matching specific attributes. It allows searching by term (e.g., part of a hostname), " + "entity type, suspicion status, asset status (internal/external), enrichment status, " + "network, or environment. This is useful for exploring the entity database, finding " + "potentially related entities during an investigation, or identifying assets with " + "specific characteristics." + ), + "tool_description_get_case_full_details": ( + "Retrieve comprehensive details for a specific case by aggregating its core information, associated alerts, and comments. " + "This tool provides a consolidated view of a security case by fetching its primary details " + "(like status, priority, description), all linked security alerts, and the full history " + "of comments added by analysts or automation. This aggregated information is essential " + "for getting a complete understanding of an incident's context, scope, investigation " + "progress, and collaborative notes without making multiple separate API calls." + ) + } + + logger.info("Starting GEPA optimization loop...") + result = gepa.optimize( + seed_candidate=seed_candidate, + trainset=dataset, + valset=dataset, + adapter=adapter, + reflection_lm=reflection_model, + max_metric_calls=150, + ) + + logger.info("Optimization finished!") + best_candidate = result.candidates[result.best_idx] + best_score = result.val_aggregate_scores[result.best_idx] + logger.info("Best achieved score: %.2f", best_score) + + # Save results + result_path = opt_dir / "gepa_optimization_results.json" + with open(result_path, "w") as f: + json.dump({ + "best_score": best_score, + "optimized_tool_descriptions": best_candidate + }, f, indent=2) + logger.info("Optimized results successfully saved to %s", result_path) + + +if __name__ == "__main__": + main() diff --git a/server/secops-soar/gepa_opt/test_live_soar.py b/server/secops-soar/gepa_opt/test_live_soar.py new file mode 100644 index 00000000..68cf94d4 --- /dev/null +++ b/server/secops-soar/gepa_opt/test_live_soar.py @@ -0,0 +1,72 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Live connection test for Chronicle SOAR MCP.""" + +import asyncio +import os +import sys +from pathlib import Path +import dotenv + +# Load environment +env_path = Path(__file__).resolve().parents[3] / ".env" +if env_path.exists(): + dotenv.load_dotenv(env_path) + +# Add parent path to PYTHONPATH +sys.path.append(str(Path(__file__).resolve().parents[1])) + +from secops_soar_mcp import bindings +from secops_soar_mcp.utils import consts + +async def main(): + url = os.getenv("SOAR_URL") + app_key = os.getenv("SOAR_APP_KEY") + print(f"Connecting to SOAR URL: {url}") + print(f"SOAR APP KEY defined: {bool(app_key)}") + + try: + await bindings.bind() + print("Successfully bound and verified scopes!") + print(f"Valid scopes count: {len(bindings.valid_scopes)}") + print(f"Scopes: {bindings.valid_scopes}") + + print("\nTesting list_cases endpoint...") + cases = await bindings.http_client.get(consts.Endpoints.BASE_CASE_URL) + if cases is not None: + print("Successfully retrieved cases from live API!") + if isinstance(cases, list): + print(f"Retrieved {len(cases)} cases.") + if len(cases) > 0: + print("Sample case summary:") + first_case = cases[0] + if isinstance(first_case, dict): + for k in ["id", "identifier", "name", "priority", "status"]: + print(f" {k}: {first_case.get(k)}") + else: + print(f" Raw: {first_case}") + else: + print("Response content:") + print(cases) + else: + print("Failed to fetch cases (response was None).") + + except Exception as e: + print(f"Error during live test execution: {e}") + finally: + if bindings.http_client: + await bindings.cleanup() + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/server/secops-soar/secops_soar_mcp/bindings.py b/server/secops-soar/secops_soar_mcp/bindings.py index 094c1f1a..74cfb6d0 100644 --- a/server/secops-soar/secops_soar_mcp/bindings.py +++ b/server/secops-soar/secops_soar_mcp/bindings.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -30,12 +30,19 @@ async def _get_valid_scopes(): - valid_scopes_list = await http_client.get(consts.Endpoints.GET_SCOPES) - if valid_scopes_list is None: - raise RuntimeError( - "Failed to fetch valid scopes from SOAR, please make sure you have configured the right SOAR credentials. Shutting down..." + try: + valid_scopes_list = await http_client.get(consts.Endpoints.GET_SCOPES) + if valid_scopes_list is None: + logger.warning( + "Failed to fetch valid scopes from SOAR. Please make sure you have configured the right SOAR credentials if you need active API access. Continuing in offline mode..." + ) + return set() + return set(valid_scopes_list) + except Exception as e: + logger.warning( + "Error fetching valid scopes from SOAR: %s. Continuing in offline mode...", e ) - return set(valid_scopes_list) + return set() async def bind(): diff --git a/server/secops-soar/secops_soar_mcp/case_management.py b/server/secops-soar/secops_soar_mcp/case_management.py index 2c338371..59f94e5c 100644 --- a/server/secops-soar/secops_soar_mcp/case_management.py +++ b/server/secops-soar/secops_soar_mcp/case_management.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/server/secops-soar/secops_soar_mcp/http_client.py b/server/secops-soar/secops_soar_mcp/http_client.py index 1c03768b..30d3940a 100644 --- a/server/secops-soar/secops_soar_mcp/http_client.py +++ b/server/secops-soar/secops_soar_mcp/http_client.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,7 @@ class HttpClient: """HTTP client for making requests to the SecOps SOAR API.""" def __init__(self, base_url: str, app_key: str): - self.base_url = base_url + self.base_url = base_url or "http://localhost" self.app_key = app_key self._session = None diff --git a/server/secops-soar/tests/conftest.py b/server/secops-soar/tests/conftest.py index 8b9c4b87..2cb4b120 100644 --- a/server/secops-soar/tests/conftest.py +++ b/server/secops-soar/tests/conftest.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,9 +16,13 @@ import os import pytest import pytest_asyncio +import dotenv from secops_soar_mcp import bindings from typing import Dict +# Load .env file at session start +dotenv.load_dotenv() + @pytest.fixture def config_path() -> str: @@ -32,7 +36,7 @@ def config_path() -> str: @pytest.fixture def soar_config(config_path: str) -> Dict[str, str]: - """Load SOAR configuration from the config file. + """Load SOAR configuration from the config file or environment variables. Args: config_path: Path to the configuration file @@ -41,20 +45,22 @@ def soar_config(config_path: str) -> Dict[str, str]: Dictionary with SOAR configuration Raises: - FileNotFoundError: If the config file is missing + FileNotFoundError: If the config file is missing and env vars are not set """ - if not os.path.exists(config_path): - raise FileNotFoundError( - f"SOAR config file not found at {config_path}. " - f"Please create this file with the following format:\n" - f"{{\n" - f' "SOAR_URL": "your-soar-url",\n' - f' "SOAR_APP_KEY": "your-soar-app-key",\n' - f"}}" - ) - - with open(config_path, "r") as f: - return json.load(f) + if os.path.exists(config_path): + with open(config_path, "r") as f: + return json.load(f) + + if os.getenv("SOAR_URL") and os.getenv("SOAR_APP_KEY"): + return { + "SOAR_URL": os.getenv("SOAR_URL"), + "SOAR_APP_KEY": os.getenv("SOAR_APP_KEY"), + } + + raise FileNotFoundError( + f"SOAR configuration not found. Please set SOAR_URL and SOAR_APP_KEY environment variables " + f"or create a config file at {config_path}." + ) def update_env_vars(soar_config: Dict[str, str]): diff --git a/server/secops-soar/tests/tests.py b/server/secops-soar/tests/tests.py index 916bc980..a59d63df 100644 --- a/server/secops-soar/tests/tests.py +++ b/server/secops-soar/tests/tests.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,13 +17,13 @@ These tests require proper SOAR authentication and configuration to run. To run these tests: -1. Make sure you have created a config.json file in the tests directory with - your SOAR credentials (see conftest.py for format) +1. Make sure you have configured SOAR credentials in your environment or .env 2. Run: pytest -xvs server/secops-soar/tests/tests.py """ import mcp import pytest +import json from typing import Optional from secops_soar_mcp.server import mcp as mcp_server @@ -38,10 +38,15 @@ @pytest.mark.parametrize( argnames=["tool_name", "tool_arguments", "expected_substring"], argvalues=[ - ("list_cases", None, "cases"), + ("list_cases", None, "nextPageToken"), + ("get_case_full_details", {"case_id": "9675"}, "case_details"), + ("list_alerts_by_case", {"case_id": "9675"}, "GooglersOnlySoarCases"), + ("list_alert_group_identifiers_by_case", {"case_id": "9675"}, "["), + ("change_case_priority", {"case_id": "9675", "case_priority": "PriorityMedium"}, "PriorityMedium"), + ("post_case_comment", {"case_id": "9675", "comment": "Verification test comment"}, "Verification test comment"), ], ) -async def test_tool(tool_name, tool_arguments, expected_substring): +async def test_tool(tool_name, tool_arguments, expected_substring, setup_bindings): response = await call_tool_and_get_text_response(tool_name, tool_arguments) assert expected_substring in response @@ -75,3 +80,33 @@ async def test_server_connection(setup_bindings): tools_result = await client.list_tools() assert isinstance(tools_result, mcp.ListToolsResult) assert len(tools_result.tools) > 0 + + +@pytest.mark.asyncio(loop_scope="session") +async def test_dynamic_alert_events_and_entities(setup_bindings): + """Dynamically fetch alerts and entities from case 9675 and test related tools.""" + + # 1. Fetch case details to find alert IDs and entities + case_details_str = await call_tool_and_get_text_response("get_case_full_details", {"case_id": "9675"}) + case_details = json.loads(case_details_str) + + # 2. Get alerts list + alerts = case_details.get("case_alerts", []) + if isinstance(alerts, list) and len(alerts) > 0: + first_alert = alerts[0] + alert_id = str(first_alert.get("id")) + + # Test list_events_by_alert + events_str = await call_tool_and_get_text_response( + "list_events_by_alert", + {"case_id": "9675", "alert_id": alert_id} + ) + assert "events" in events_str or "[" in events_str + + # 3. Test entity details and search if case contains entities + # Fetch environment of the case to use for entity environment parameters + env = case_details.get("case_details", {}).get("environment", "GooglersOnlySoarCases") + + # Run a generic entity search + search_str = await call_tool_and_get_text_response("search_entity", {"term": "8.8.8.8"}) + assert "[" in search_str diff --git a/server/secops/gepa_opt/gepa_optimization_results.json b/server/secops/gepa_opt/gepa_optimization_results.json new file mode 100644 index 00000000..f9c0eb22 --- /dev/null +++ b/server/secops/gepa_opt/gepa_optimization_results.json @@ -0,0 +1,17 @@ +{ + "best_score": 0.8571428571428571, + "optimized_tool_descriptions": { + "tool_description_search_security_events": "Search for security events in Chronicle SIEM using natural language. Allows searching Chronicle event logs using natural language queries, which are automatically translated into UDM queries for execution. Ideal for deep investigation after an initial alert, case, or entity has been prioritized.", + "tool_description_get_security_alerts": "Get security alerts directly from Chronicle SIEM. Retrieves a list of recent security alerts generated within Chronicle, based on detection rules or other alert sources configured in the SIEM.", + "tool_description_get_security_alert_by_id": "Get security alert by ID directly from Chronicle SIEM. Gets an alert by ID. Use this for direct monitoring of SIEM alert activity, potentially identifying issues before they are ingested or processed by other platforms.", + "tool_description_do_update_security_alert": "Update security alert attributes directly in Chronicle SIEM. Modifies specific fields of an existing security alert within Chronicle based on its ID. This function allows for updates to an alert's status, severity, verdict, assigned scores, comments, and other metadata.", + "tool_description_lookup_entity": "Look up an entity (IP, domain, hash, user, etc.) in Chronicle SIEM for enrichment. Provides a comprehensive summary of an entity's activity based on historical log data within Chronicle over a specified time period. This tool queries Chronicle SIEM directly.", + "tool_description_list_security_rules": "List security detection rules configured in Chronicle SIEM, with support for pagination. Retrieves the definitions of detection rules currently active or configured within the Chronicle SIEM instance.", + "tool_description_search_security_rules": "Search security detection rules configured in Chronicle SIEM. Retrieves the definitions of detection rules currently active or configured within the Chronicle SIEM instance based on a regex pattern.", + "tool_description_get_ioc_matches": "Get IoC matches directly from Chronicle SIEM. Retrieves recent IoC matches observed within the SIEM environment.", + "tool_description_get_threat_intel": "Retrieve threat intelligence from Chronicle SIEM. Queries threat intelligence data and returns detailed summaries regarding threat actors, campaigns, CVEs, or general best practices.", + "tool_description_search_udm": "Search UDM events using UDM query in Chronicle. Accepts raw YARA-L UDM query strings to locate underlying event records in the Chronicle database over a specified time range.", + "tool_description_export_udm_search_csv": "Export UDM search results as a formatted CSV string. Useful for downloading logs for external reporting or offline triage.", + "tool_description_find_udm_field_values": "Find UDM field values for autocomplete. Searches for values matching a query string in specified UDM fields." + } +} \ No newline at end of file diff --git a/server/secops/gepa_opt/mcp_dataset.json b/server/secops/gepa_opt/mcp_dataset.json new file mode 100644 index 00000000..d3e136a3 --- /dev/null +++ b/server/secops/gepa_opt/mcp_dataset.json @@ -0,0 +1,168 @@ +[ + { + "user_query": "Find network connections from 192.168.1.100 in the last 48 hours using natural language", + "reference_answer": "search_security_events", + "tool_arguments": { + "text": "network connections from 192.168.1.100", + "hours_back": 48 + } + }, + { + "user_query": "Search for failed login events for user 'admin' in the last 12 hours", + "reference_answer": "search_security_events", + "tool_arguments": { + "text": "failed login events for user 'admin'", + "hours_back": 12 + } + }, + { + "user_query": "Search security events for 'DNS lookups to phishing.com' in the last 24 hours", + "reference_answer": "search_security_events", + "tool_arguments": { + "text": "DNS lookups to phishing.com", + "hours_back": 24 + } + }, + { + "user_query": "List the last 10 security alerts generated in Chronicle", + "reference_answer": "get_security_alerts", + "tool_arguments": { + "max_alerts": 10 + } + }, + { + "user_query": "Retrieve recent open security alerts from the last 7 days", + "reference_answer": "get_security_alerts", + "tool_arguments": { + "hours_back": 168, + "status_filter": "feedback_summary.status != \"CLOSED\"" + } + }, + { + "user_query": "Get the detailed information and detections for alert ID 'de_1234567890'", + "reference_answer": "get_security_alert_by_id", + "tool_arguments": { + "alert_id": "de_1234567890", + "include_detections": true + } + }, + { + "user_query": "Show me alert details for 'de_9876543210' without including detections", + "reference_answer": "get_security_alert_by_id", + "tool_arguments": { + "alert_id": "de_9876543210", + "include_detections": false + } + }, + { + "user_query": "Update alert 'de_1122334455' to CLOSED status with verdict FALSE_POSITIVE, comment 'Confirmed staging test activity', and reason REASON_MAINTENANCE", + "reference_answer": "do_update_security_alert", + "tool_arguments": { + "alert_id": "de_1122334455", + "status": "CLOSED", + "verdict": "FALSE_POSITIVE", + "comment": "Confirmed staging test activity", + "reason": "REASON_MAINTENANCE" + } + }, + { + "user_query": "Mark alert 'de_9988776655' as TRUE_POSITIVE with critical priority, severity 95, and root cause 'Phishing email leading to credential theft'", + "reference_answer": "do_update_security_alert", + "tool_arguments": { + "alert_id": "de_9988776655", + "verdict": "TRUE_POSITIVE", + "priority": "PRIORITY_CRITICAL", + "severity": 95, + "root_cause": "Phishing email leading to credential theft" + } + }, + { + "user_query": "Lookup entity information for IP address '8.8.8.8'", + "reference_answer": "lookup_entity", + "tool_arguments": { + "entity_value": "8.8.8.8" + } + }, + { + "user_query": "Check if there is any entity data for domain 'bad-domain.com' in the last 7 days", + "reference_answer": "lookup_entity", + "tool_arguments": { + "entity_value": "bad-domain.com", + "hours_back": 168 + } + }, + { + "user_query": "Search for entity details of file hash '44d88612fea8a8f36de82e1278abb02f'", + "reference_answer": "lookup_entity", + "tool_arguments": { + "entity_value": "44d88612fea8a8f36de82e1278abb02f" + } + }, + { + "user_query": "List all configured security detection rules in Chronicle", + "reference_answer": "list_security_rules", + "tool_arguments": {} + }, + { + "user_query": "Find rules matching the query pattern '.*ransomware.*'", + "reference_answer": "search_security_rules", + "tool_arguments": { + "query": ".*ransomware.*" + } + }, + { + "user_query": "Retrieve recent IoC matches from the last 24 hours (max 15 matches)", + "reference_answer": "get_ioc_matches", + "tool_arguments": { + "hours_back": 24, + "max_matches": 15 + } + }, + { + "user_query": "Search threat intelligence for information on 'APT29' threat actor group", + "reference_answer": "get_threat_intel", + "tool_arguments": { + "query": "APT29" + } + }, + { + "user_query": "Get threat intelligence for CVE-2021-44228", + "reference_answer": "get_threat_intel", + "tool_arguments": { + "query": "CVE-2021-44228" + } + }, + { + "user_query": "Search UDM using query 'metadata.event_type = \"NETWORK_CONNECTION\" AND principal.ip = \"10.0.0.5\"' for the last 6 hours", + "reference_answer": "search_udm", + "tool_arguments": { + "query": "metadata.event_type = \"NETWORK_CONNECTION\" AND principal.ip = \"10.0.0.5\"", + "hours_back": 6 + } + }, + { + "user_query": "Run UDM query 'target.user.userid = \"admin\"' from start time '2026-06-01T00:00:00Z' to end time '2026-06-02T00:00:00Z' limiting to 50 events", + "reference_answer": "search_udm", + "tool_arguments": { + "query": "target.user.userid = \"admin\"", + "start_time": "2026-06-01T00:00:00Z", + "end_time": "2026-06-02T00:00:00Z", + "max_events": 50 + } + }, + { + "user_query": "Export UDM search results for query 'metadata.event_type = \"USER_LOGIN\"' with fields 'metadata.event_timestamp, principal.user.userid' to CSV", + "reference_answer": "export_udm_search_csv", + "tool_arguments": { + "query": "metadata.event_type = \"USER_LOGIN\"", + "fields": ["metadata.event_timestamp", "principal.user.userid"] + } + }, + { + "user_query": "Find UDM field values for autocomplete with term 'cymbal' in the last 24 hours", + "reference_answer": "find_udm_field_values", + "tool_arguments": { + "query": "cymbal" + } + } +] diff --git a/server/secops/gepa_opt/optimize_secops_mcp.py b/server/secops/gepa_opt/optimize_secops_mcp.py new file mode 100644 index 00000000..96ac23b7 --- /dev/null +++ b/server/secops/gepa_opt/optimize_secops_mcp.py @@ -0,0 +1,272 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GEPA Optimizer script for server/secops MCP tools. + +Configured for Vertex AI Model Garden via LiteLLM. +""" + +import os +import sys +import json +import logging +from pathlib import Path +import asyncio + +# 1. Apply asyncio subprocess reader limit monkey patch +original_create_subprocess_exec = asyncio.create_subprocess_exec +async def patched_create_subprocess_exec(*args, **kwargs): + kwargs['limit'] = 10 * 1024 * 1024 # 10MB buffer + return await original_create_subprocess_exec(*args, **kwargs) +asyncio.create_subprocess_exec = patched_create_subprocess_exec + +# 2. Load environment variables using dotenv +import dotenv +env_path = Path(__file__).resolve().parents[3] / ".env" +if env_path.exists(): + dotenv.load_dotenv(env_path) + +# 3. Configure Google Cloud Vertex AI credentials for LiteLLM +for var in ["GOOGLE_APPLICATION_CREDENTIALS", "VERTEX_PROJECT", "VERTEX_LOCATION"]: + if not os.getenv(var): + raise ValueError( + f"Missing required environment variable: {var}. " + "Please specify it in your environment or in the .env file." + ) + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# 4. Import GEPA and configure LiteLLM retries +try: + import gepa + from gepa.adapters.mcp_adapter import MCPAdapter + from mcp import StdioServerParameters + import litellm + + # Configure LiteLLM retry and drop parameter behaviors + litellm.num_retries = 5 + litellm.drop_params = True +except ImportError as e: + logger.error("Failed to import GEPA or required dependencies: %s", e) + sys.exit(1) + + +def load_dataset(dataset_path: Path) -> list: + """Loads the evaluation dataset.""" + with open(dataset_path, "r") as f: + return json.load(f) + + +def secops_metric_fn(data_inst, output: str) -> float: + """ + Evaluation Metric: Scores 1.0 if the model correctly selects the expected tool + and extracts all targeted parameters correctly. + """ + import time + time.sleep(1.5) # Safe spacing to avoid hitting Vertex free-tier RPM limits + + try: + cleaned = output.strip() + if cleaned.startswith("```json"): + cleaned = cleaned[7:] + if cleaned.endswith("```"): + cleaned = cleaned[:-3] + cleaned = cleaned.strip() + parsed = json.loads(cleaned) + except Exception: + return 0.0 + + if isinstance(parsed, list): + if len(parsed) > 0: + parsed = parsed[0] + else: + return 0.0 + + if not isinstance(parsed, dict): + return 0.0 + + if parsed.get("action") != "call_tool": + return 0.0 + + selected_tool = parsed.get("tool") + expected_tool = data_inst.get("reference_answer") + if selected_tool != expected_tool: + return 0.0 + + arguments = parsed.get("arguments", {}) + expected_args = data_inst.get("tool_arguments", {}) + + for k, v in expected_args.items(): + if k not in arguments: + return 0.0 + + actual_val = arguments[k] + if isinstance(v, list) and isinstance(actual_val, list): + if len(v) != len(actual_val) or sorted(v) != sorted(actual_val): + return 0.0 + elif isinstance(v, str) and isinstance(actual_val, str): + if v.lower() not in actual_val.lower(): + return 0.0 + elif actual_val != v: + return 0.0 + + return 1.0 + + +def main(): + opt_dir = Path(__file__).parent + server_dir = opt_dir.parent + server_py = server_dir / "secops_mcp" / "server.py" + dataset_json = opt_dir / "mcp_dataset.json" + + if not server_py.exists(): + logger.error("MCP Server script not found at %s", server_py) + sys.exit(1) + + dataset = load_dataset(dataset_json) + logger.info("Loaded dataset with %d items", len(dataset)) + + # The Chronicle SecOps tools we wish to optimize + raw_tool_names = [ + "search_security_events", + "get_security_alerts", + "get_security_alert_by_id", + "do_update_security_alert", + "lookup_entity", + "list_security_rules", + "search_security_rules", + "get_ioc_matches", + "get_threat_intel", + "search_udm", + "export_udm_search_csv", + "find_udm_field_values" + ] + + # Target Vertex AI Model Garden models + task_model = "vertex_ai/gemini-2.5-flash" + reflection_model = "vertex_ai/gemini-2.5-pro" + + logger.info("Initializing MCPAdapter targeting local stdio server...") + logger.info("Task Model: %s", task_model) + logger.info("Reflection Model: %s", reflection_model) + + adapter = MCPAdapter( + tool_names=raw_tool_names, + task_model=task_model, + metric_fn=secops_metric_fn, + server_params=StdioServerParameters( + command=sys.executable, + args=[str(server_py)], + ), + base_system_prompt=( + "You are a security analyst with access to the Chronicle SecOps SIEM platform. " + "Your goal is to determine which security event search, alert management, rule management, " + "or threat intelligence tool to call and what parameters to pass based on the user's query." + ), + enable_two_pass=False, + failure_score=0.0, + ) + + # Seed candidates with original tool descriptions + seed_candidate = { + "tool_description_search_security_events": ( + "Search for security events in Chronicle SIEM using natural language. " + "Allows searching Chronicle event logs using natural language queries, which are " + "automatically translated into UDM queries for execution. " + "Ideal for deep investigation after an initial alert, case, or entity has been prioritized." + ), + "tool_description_get_security_alerts": ( + "Get security alerts directly from Chronicle SIEM. " + "Retrieves a list of recent security alerts generated within Chronicle, based on " + "detection rules or other alert sources configured in the SIEM." + ), + "tool_description_get_security_alert_by_id": ( + "Get security alert by ID directly from Chronicle SIEM. " + "Gets an alert by ID. " + "Use this for direct monitoring of SIEM alert activity, potentially identifying " + "issues before they are ingested or processed by other platforms." + ), + "tool_description_do_update_security_alert": ( + "Update security alert attributes directly in Chronicle SIEM. " + "Modifies specific fields of an existing security alert within Chronicle based on its ID. " + "This function allows for updates to an alert's status, severity, verdict, assigned scores, comments, and other metadata." + ), + "tool_description_lookup_entity": ( + "Look up an entity (IP, domain, hash, user, etc.) in Chronicle SIEM for enrichment. " + "Provides a comprehensive summary of an entity's activity based on historical log data " + "within Chronicle over a specified time period. This tool queries Chronicle SIEM directly." + ), + "tool_description_list_security_rules": ( + "List security detection rules configured in Chronicle SIEM, with support for pagination. " + "Retrieves the definitions of detection rules currently active or configured " + "within the Chronicle SIEM instance." + ), + "tool_description_search_security_rules": ( + "Search security detection rules configured in Chronicle SIEM. " + "Retrieves the definitions of detection rules currently active or configured " + "within the Chronicle SIEM instance based on a regex pattern." + ), + "tool_description_get_ioc_matches": ( + "Get IoC matches directly from Chronicle SIEM. " + "Retrieves recent IoC matches observed within the SIEM environment." + ), + "tool_description_get_threat_intel": ( + "Retrieve threat intelligence from Chronicle SIEM. " + "Queries threat intelligence data and returns detailed summaries regarding " + "threat actors, campaigns, CVEs, or general best practices." + ), + "tool_description_search_udm": ( + "Search UDM events using UDM query in Chronicle. " + "Accepts raw YARA-L UDM query strings to locate underlying event records " + "in the Chronicle database over a specified time range." + ), + "tool_description_export_udm_search_csv": ( + "Export UDM search results as a formatted CSV string. " + "Useful for downloading logs for external reporting or offline triage." + ), + "tool_description_find_udm_field_values": ( + "Find UDM field values for autocomplete. " + "Searches for values matching a query string in specified UDM fields." + ) + } + + logger.info("Starting GEPA optimization loop...") + result = gepa.optimize( + seed_candidate=seed_candidate, + trainset=dataset, + valset=dataset, + adapter=adapter, + reflection_lm=reflection_model, + reflection_minibatch_size=len(dataset), + max_metric_calls=300, + ) + + logger.info("Optimization finished!") + best_candidate = result.candidates[result.best_idx] + best_score = result.val_aggregate_scores[result.best_idx] + logger.info("Best achieved score: %.2f", best_score) + + # Save results + result_path = opt_dir / "gepa_optimization_results.json" + with open(result_path, "w") as f: + json.dump({ + "best_score": best_score, + "optimized_tool_descriptions": best_candidate + }, f, indent=2) + logger.info("Optimized results successfully saved to %s", result_path) + + +if __name__ == "__main__": + main() diff --git a/server/secops/secops_mcp/tools/security_alerts.py b/server/secops/secops_mcp/tools/security_alerts.py index ab90dcf1..bfc79d4d 100644 --- a/server/secops/secops_mcp/tools/security_alerts.py +++ b/server/secops/secops_mcp/tools/security_alerts.py @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,8 +35,7 @@ async def get_security_alerts( ) -> str: """Get security alerts directly from Chronicle SIEM. - Retrieves a list of recent security alerts generated within Chronicle, based on - detection rules or other alert sources configured in the SIEM. + Retrieves a list of recent security alerts. If the user asks for "open alerts", you must explicitly set status_filter='feedback_summary.status != "CLOSED"'. **Workflow Integration:** - Use this for direct monitoring of SIEM alert activity, potentially identifying @@ -219,9 +218,11 @@ async def do_update_security_alert( root_cause: Optional[str] = None ) -> str: """ - Update security alert attributes directly in Chronicle SIEM. + Update security alert attributes directly in Chronicle SIEM. -Modifies specific fields of an existing security alert within Chronicle based on its ID. This function allows for updates to an alert's status, severity, verdict, assigned scores, comments, and other metadata. This is typically performed after an investigation, triage, or automated analysis provides new insights or conclusions about the alert. At least one of the optional fields related to alert attributes (e.g., status, severity, comment) should be provided to perform a meaningful update. + NOTE: The 'priority' parameter must use the full prefix (e.g., 'PRIORITY_CRITICAL', 'PRIORITY_HIGH', 'PRIORITY_LOW'). The 'verdict' parameter must be 'TRUE_POSITIVE' or 'FALSE_POSITIVE'. The 'status' parameter must be 'CLOSED', 'OPEN', 'NEW', or 'REVIEWED'. + + Modifies specific fields of an existing security alert within Chronicle based on its ID. This function allows for updates to an alert's status, severity, verdict, assigned scores, comments, and other metadata. This is typically performed after an investigation, triage, or automated analysis provides new insights or conclusions about the alert. At least one of the optional fields related to alert attributes (e.g., status, severity, comment) should be provided to perform a meaningful update. **Workflow Integration:** - Utilize when SOAR is not a core technology the investigator uses