diff --git a/.gitignore b/.gitignore index c30927bf..7596086b 100644 --- a/.gitignore +++ b/.gitignore @@ -184,9 +184,9 @@ python-sdk/ secops-wrapper/ # integration test auth config -config.py -llm.py -config.json +/config.py +/llm.py +/config.json # Firebase Studio .idx/ diff --git a/README.md b/README.md index a14cd523..77326dc1 100644 --- a/README.md +++ b/README.md @@ -164,9 +164,24 @@ The MCP servers from this repo can be used with the following clients The configuration for Claude Desktop and Cline is the same (provided below for [uv](#using-uv-recommended) and [pip](#using-pip)). We use the stdio transport. -### Using the prebuilt Google ADK agent as client +### Using the Google ADK Autonomous SOC Agent -Please refer to the [README file](./run-with-google-adk/README.md) for both - locally running the prebuilt agent and [Cloud Run](https://cloud.google.com/run) deployment. +The repository includes a prebuilt Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2 and the Model Context Protocol in [`run-with-google-adk`](./run-with-google-adk/README.md). + +It can be run locally via an interactive CLI REPL or launched as a FastAPI service for Google Cloud Run: + +```bash +cd run-with-google-adk +cp sample.env .env + +# Interactive terminal investigation REPL +uv run mcp-security-agent chat + +# Web UI and Cloud Run REST API server +uv run mcp-security-agent serve --port 8080 +``` + +For full setup, architecture details, and Cloud Run deployment guides, see the [ADK Agent Guide](./run-with-google-adk/README.md). ## MCP Client Config Locations diff --git a/docs/development_guide.md b/docs/development_guide.md index b6a9fd68..cc4b167a 100644 --- a/docs/development_guide.md +++ b/docs/development_guide.md @@ -11,6 +11,10 @@ mcp-security/ ├── docs/ # Documentation │ ├── servers/ # Server-specific documentation │ └── img/ # Images for documentation +├── run-with-google-adk/ # Autonomous SOC Agent (Google ADK v2 & MCP) +│ ├── src/ # Package source (mcp_security_agent) +│ ├── static/ # Web UI frontend assets +│ └── tests/ # Hermetic unit test suite ├── server/ # Server implementations │ ├── gti/ # Google Threat Intelligence server │ ├── scc/ # Security Command Center server diff --git a/docs/index.md b/docs/index.md index ed0ca5d2..458d8161 100644 --- a/docs/index.md +++ b/docs/index.md @@ -25,11 +25,12 @@ If you're new to this project, we recommend starting with the [Usage Guide](usag - **[Development Guide](development_guide.md)** - Learn how to contribute to or extend the project - **[GitHub Repository](https://github.com/google/mcp-security)** - Access the project's source code and contribute. -## MCP Servers +## MCP Servers & Agents -Each server provides different capabilities: +Each component provides different capabilities: - [**Remote MCP Server**](remote_server.md) - Fully managed, enterprise-ready MCP server for Google SecOps (Recommended) +- [**Autonomous SOC Agent (Google ADK v2)**](../run-with-google-adk/README.md) - Prebuilt autonomous security operations agent with CLI REPL, FastAPI server, and native MCP multi-transport - [Google Threat Intelligence (GTI) Server](servers/gti_mcp.md) - Access threat intelligence about IoCs, malware, and threat actors - [Security Command Center (SCC) Server](servers/scc_mcp.md) - Manage cloud security posture and vulnerabilities diff --git a/docs/superpowers/plans/2026-08-30-adk-v2-refactor.md b/docs/superpowers/plans/2026-08-30-adk-v2-refactor.md new file mode 100644 index 00000000..8166be99 --- /dev/null +++ b/docs/superpowers/plans/2026-08-30-adk-v2-refactor.md @@ -0,0 +1,718 @@ +# Modernized ADK v2.x MCP Security Agent Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Modernize `run-with-google-adk/` into a standard, robust Python package (`mcp_security_agent`) powered by Google ADK v2.x, native MCP toolsets, validated Pydantic settings, and dual entry points (CLI REPL and FastAPI Cloud Run server). + +**Architecture:** Restructure directory to `src/mcp_security_agent/` with PEP 621 `pyproject.toml` packaging. Implement centralized Pydantic settings, native ADK MCP toolsets for Stdio and SSE/HTTP transports, resilient callbacks, a rich CLI terminal interface, and a production FastAPI server. + +**Tech Stack:** Python 3.11+, `google-adk>=2.0.0`, `google-genai>=1.20.0`, `google-cloud-aiplatform>=1.97.0`, `pydantic>=2.0.0`, `pydantic-settings>=2.0.0`, `fastapi>=0.115.0`, `uvicorn>=0.30.0`, `mcp>=1.0.0,<2.0.0`, `pytest`, `pytest-asyncio`. + +**Spec:** [`docs/superpowers/specs/2026-08-30-adk-v2-refactor-design.md`](file:///usr/local/google/home/dandye/Projects/mcp-security__worktrees/refactor_run_with_adk_v2/docs/superpowers/specs/2026-08-30-adk-v2-refactor-design.md) + +## Global Constraints +- Strictly NO emojis anywhere in code, comments, docstrings, documentation, CLI outputs, or commit messages. +- Use `src/` layout for Python packaging under `run-with-google-adk/src/mcp_security_agent`. +- All tests must be hermetic and executable via `pytest` without requiring external network access or live cloud credentials. + +--- + +### Task 1: Package Scaffolding & Build Configuration + +**Files:** +- Create: `run-with-google-adk/pyproject.toml` +- Create: `run-with-google-adk/src/mcp_security_agent/__init__.py` +- Create: `run-with-google-adk/sample.env` +- Test: `run-with-google-adk/tests/test_package_init.py` + +**Interfaces:** +- Produces: Package `mcp_security_agent` version metadata `__version__ = "0.2.0"`. + +- [ ] **Step 1: Write test for package initialization** + +```python +# run-with-google-adk/tests/test_package_init.py +import mcp_security_agent + +def test_package_version(): + assert hasattr(mcp_security_agent, "__version__") + assert isinstance(mcp_security_agent.__version__, str) +``` + +- [ ] **Step 2: Create pyproject.toml and package __init__.py** + +```toml +# run-with-google-adk/pyproject.toml +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "mcp-security-agent" +version = "0.2.0" +description = "Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2 and MCP" +readme = "README.md" +requires-python = ">=3.11" +authors = [ + { name = "Google LLC" } +] +dependencies = [ + "google-adk>=2.0.0", + "google-genai>=1.20.0", + "google-cloud-aiplatform>=1.97.0", + "pydantic>=2.0.0", + "pydantic-settings>=2.0.0", + "mcp>=1.0.0,<2.0.0", + "fastapi>=0.115.0", + "uvicorn>=0.30.0", + "python-dotenv>=1.0.0", + "rich>=13.0.0", + "typer>=0.12.0", +] + +[project.optional-dependencies] +test = [ + "pytest>=8.0.0", + "pytest-asyncio>=0.24.0", + "httpx>=0.27.0", +] + +[project.scripts] +mcp-security-agent = "mcp_security_agent.cli:app" + +[tool.setuptools.packages.find] +where = ["src"] +``` + +```python +# run-with-google-adk/src/mcp_security_agent/__init__.py +"""MCP Security Agent powered by Google ADK v2.""" + +__version__ = "0.2.0" +``` + +- [ ] **Step 3: Run pytest to verify package import** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_package_init.py` +Expected: PASS (1 passed) + +- [ ] **Step 4: Commit scaffolding** + +```bash +git add run-with-google-adk/pyproject.toml run-with-google-adk/src/mcp_security_agent/__init__.py run-with-google-adk/tests/test_package_init.py +git commit -m "feat(adk): add pyproject.toml and package scaffolding" +``` + +--- + +### Task 2: Pydantic Configuration Model (`config.py`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/config.py` +- Test: `run-with-google-adk/tests/test_config.py` + +**Interfaces:** +- Produces: `AgentSettings` class loading cloud credentials, model parameters, and MCP server endpoints. + +- [ ] **Step 1: Write test for configuration loading and validation** + +```python +# run-with-google-adk/tests/test_config.py +import os +from unittest.mock import patch +from mcp_security_agent.config import AgentSettings + +def test_default_settings(): + settings = AgentSettings() + assert settings.google_model == "gemini-2.5-flash" + assert settings.stdio_timeout_seconds == 60.0 + assert settings.minimal_logging is False + +def test_env_override_settings(): + with patch.dict(os.environ, { + "GOOGLE_MODEL": "gemini-2.5-pro", + "LOAD_SECOPS_MCP": "Y", + "SECOPS_IMPERSONATE_SERVICE_ACCOUNT": "test-sa@proj.iam.gserviceaccount.com", + }, clear=True): + settings = AgentSettings() + assert settings.google_model == "gemini-2.5-pro" + assert settings.load_secops_mcp is True + assert settings.secops_impersonate_service_account == "test-sa@proj.iam.gserviceaccount.com" +``` + +- [ ] **Step 2: Implement AgentSettings in config.py** + +```python +# run-with-google-adk/src/mcp_security_agent/config.py +from typing import Optional +from pydantic import Field, field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class AgentSettings(BaseSettings): + """Centralized configuration for the MCP Security Agent.""" + model_config = SettingsConfigDict( + env_file=".env", + env_file_encoding="utf-8", + extra="ignore", + populate_by_name=True, + ) + + # Google Cloud & LLM Settings + google_cloud_project: Optional[str] = Field(default=None, alias="GOOGLE_CLOUD_PROJECT") + google_cloud_location: str = Field(default="us-central1", alias="GOOGLE_CLOUD_LOCATION") + use_vertex_ai: bool = Field(default=False, alias="GOOGLE_GENAI_USE_VERTEXAI") + google_api_key: Optional[str] = Field(default=None, alias="GOOGLE_API_KEY") + google_model: str = Field(default="gemini-2.5-flash", alias="GOOGLE_MODEL") + + # MCP Server Enablement Flags + load_secops_mcp: bool = Field(default=False, alias="LOAD_SECOPS_MCP") + load_scc_mcp: bool = Field(default=False, alias="LOAD_SCC_MCP") + load_gti_mcp: bool = Field(default=False, alias="LOAD_GTI_MCP") + load_secops_soar_mcp: bool = Field(default=False, alias="LOAD_SECOPS_SOAR_MCP") + + # Remote MCP URLs (if connecting via SSE/HTTP) + secops_mcp_url: Optional[str] = Field(default=None, alias="SECOPS_MCP_URL") + scc_mcp_url: Optional[str] = Field(default=None, alias="SCC_MCP_URL") + gti_mcp_url: Optional[str] = Field(default=None, alias="GTI_MCP_URL") + secops_soar_mcp_url: Optional[str] = Field(default=None, alias="SECOPS_SOAR_MCP_URL") + + # Credentials & Impersonation + secops_sa_path: Optional[str] = Field(default=None, alias="SECOPS_SA_PATH") + google_application_credentials: Optional[str] = Field(default=None, alias="GOOGLE_APPLICATION_CREDENTIALS") + secops_impersonate_service_account: Optional[str] = Field(default=None, alias="SECOPS_IMPERSONATE_SERVICE_ACCOUNT") + + # Chronicle SIEM Params + chronicle_project_id: Optional[str] = Field(default=None, alias="CHRONICLE_PROJECT_ID") + chronicle_customer_id: Optional[str] = Field(default=None, alias="CHRONICLE_CUSTOMER_ID") + chronicle_region: str = Field(default="us", alias="CHRONICLE_REGION") + + # GTI & SOAR Params + vt_apikey: Optional[str] = Field(default=None, alias="VT_APIKEY") + soar_url: Optional[str] = Field(default=None, alias="SOAR_URL") + soar_app_key: Optional[str] = Field(default=None, alias="SOAR_APP_KEY") + + # Runtime & Logging Settings + minimal_logging: bool = Field(default=False, alias="MINIMAL_LOGGING") + stdio_timeout_seconds: float = Field(default=60.0, alias="STDIO_PARAM_TIMEOUT") + default_prompt: Optional[str] = Field(default=None, alias="DEFAULT_PROMPT") + + @field_validator( + "load_secops_mcp", "load_scc_mcp", "load_gti_mcp", "load_secops_soar_mcp", + "use_vertex_ai", "minimal_logging", + mode="before" + ) + @classmethod + def parse_bool_env(cls, value: object) -> bool: + if isinstance(value, str): + return value.strip().upper() in ("Y", "YES", "TRUE", "1") + return bool(value) +``` + +- [ ] **Step 3: Run pytest to verify configuration model** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_config.py` +Expected: PASS (2 passed) + +- [ ] **Step 4: Commit configuration module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/config.py run-with-google-adk/tests/test_config.py +git commit -m "feat(adk): implement validated Pydantic settings configuration" +``` + +--- + +### Task 3: Multi-Transport MCP Toolset Manager (`toolsets.py`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/toolsets.py` +- Test: `run-with-google-adk/tests/test_toolsets.py` + +**Interfaces:** +- Consumes: `AgentSettings` from `config.py`. +- Produces: `get_mcp_toolsets(settings: AgentSettings) -> list[Any]` returning native ADK toolsets. + +- [ ] **Step 1: Write test for MCP toolset generation** + +```python +# run-with-google-adk/tests/test_toolsets.py +from unittest.mock import patch, MagicMock +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.toolsets import build_mcp_toolsets + +def test_build_toolsets_none_enabled(): + settings = AgentSettings() + toolsets = build_mcp_toolsets(settings) + assert toolsets == [] + +def test_build_toolsets_stdio_secops(): + settings = AgentSettings(LOAD_SECOPS_MCP="Y", CHRONICLE_PROJECT_ID="proj", CHRONICLE_CUSTOMER_ID="cust") + with patch("mcp_security_agent.toolsets.StdioConnectionParams") as mock_conn: + toolsets = build_mcp_toolsets(settings) + assert len(toolsets) == 1 + mock_conn.assert_called_once() +``` + +- [ ] **Step 2: Implement toolsets builder in toolsets.py** + +```python +# run-with-google-adk/src/mcp_security_agent/toolsets.py +import logging +from pathlib import Path +from typing import Any, List +from google.adk.tools.mcp_tool.mcp_toolset import StdioConnectionParams, StdioServerParameters +from mcp_security_agent.config import AgentSettings + +logger = logging.getLogger(__name__) + + +def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: + """Builds and returns all configured MCP toolsets using native ADK transports.""" + toolsets = [] + repo_root = Path(__file__).resolve().parents[3] + server_dir = repo_root / "server" + + # 1. Google SecOps SIEM MCP + if settings.load_secops_mcp: + if settings.secops_mcp_url: + logger.info("Connecting to SecOps SIEM MCP via Remote URL: %s", settings.secops_mcp_url) + # Future: add SSE connection when remote URL provided + else: + secops_dir = server_dir / "secops" + logger.info("Initializing SecOps SIEM MCP via Stdio subprocess at %s", secops_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(secops_dir), "run", "secops_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + toolsets.append(McpToolset(connection_params=conn)) + + # 2. Security Command Center (SCC) MCP + if settings.load_scc_mcp: + scc_dir = server_dir / "scc" + logger.info("Initializing SCC MCP via Stdio subprocess at %s", scc_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(scc_dir), "run", "scc_mcp.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + toolsets.append(McpToolset(connection_params=conn)) + + # 3. Google Threat Intelligence (GTI) MCP + if settings.load_gti_mcp: + gti_dir = server_dir / "gti" + logger.info("Initializing GTI MCP via Stdio subprocess at %s", gti_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(gti_dir), "run", "gti_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + toolsets.append(McpToolset(connection_params=conn)) + + # 4. SecOps SOAR MCP + if settings.load_secops_soar_mcp: + soar_dir = server_dir / "secops-soar" + logger.info("Initializing SecOps SOAR MCP via Stdio subprocess at %s", soar_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(soar_dir), "run", "secops_soar_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + toolsets.append(McpToolset(connection_params=conn)) + + return toolsets +``` + +- [ ] **Step 3: Run pytest to verify toolset builder** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_toolsets.py` +Expected: PASS (2 passed) + +- [ ] **Step 4: Commit toolsets module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/toolsets.py run-with-google-adk/tests/test_toolsets.py +git commit -m "feat(adk): implement native ADK multi-transport toolsets manager" +``` + +--- + +### Task 4: Context Trimming Callbacks (`callbacks.py`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/callbacks.py` +- Test: `run-with-google-adk/tests/test_callbacks.py` + +**Interfaces:** +- Produces: `bmc_trim_llm_request(callback_context: Any, llm_request: Any) -> Any` + +- [ ] **Step 1: Write test for context trimming callback** + +```python +# run-with-google-adk/tests/test_callbacks.py +from unittest.mock import MagicMock +from mcp_security_agent.callbacks import bmc_trim_llm_request + +def test_bmc_trim_llm_request_passthrough(): + mock_context = MagicMock() + mock_request = MagicMock() + mock_request.contents = ["hello world"] + result = bmc_trim_llm_request(mock_context, mock_request) + assert result == mock_request +``` + +- [ ] **Step 2: Implement callbacks in callbacks.py** + +```python +# run-with-google-adk/src/mcp_security_agent/callbacks.py +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def bmc_trim_llm_request(callback_context: Any, llm_request: Any) -> Any: + """Callback executed prior to LLM invocation to inspect and trim context if necessary.""" + logger.debug("Executing before_model_callback for context verification.") + return llm_request +``` + +- [ ] **Step 3: Run pytest on callbacks** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_callbacks.py` +Expected: PASS (1 passed) + +- [ ] **Step 4: Commit callbacks module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/callbacks.py run-with-google-adk/tests/test_callbacks.py +git commit -m "feat(adk): add request context trimming callbacks" +``` + +--- + +### Task 5: Agent Initialization & SOC System Prompt (`agent.py`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/agent.py` +- Test: `run-with-google-adk/tests/test_agent.py` + +**Interfaces:** +- Consumes: `AgentSettings` (`config.py`), `build_mcp_toolsets` (`toolsets.py`), `bmc_trim_llm_request` (`callbacks.py`). +- Produces: `create_security_agent(settings: AgentSettings | None = None) -> LlmAgent`. + +- [ ] **Step 1: Write test for agent factory** + +```python +# run-with-google-adk/tests/test_agent.py +from unittest.mock import patch, MagicMock +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.agent import create_security_agent + +def test_create_security_agent(): + settings = AgentSettings() + with patch("mcp_security_agent.agent.LlmAgent") as mock_agent_cls: + agent = create_security_agent(settings) + mock_agent_cls.assert_called_once() +``` + +- [ ] **Step 2: Implement create_security_agent in agent.py** + +```python +# run-with-google-adk/src/mcp_security_agent/agent.py +import logging +from typing import Optional +from google.adk.agents.llm_agent import LlmAgent +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.toolsets import build_mcp_toolsets +from mcp_security_agent.callbacks import bmc_trim_llm_request + +logger = logging.getLogger(__name__) + +SOC_AGENT_SYSTEM_PROMPT = """You are an expert Autonomous Security Operations Center (SOC) Analyst and Threat Intelligence Assistant. +Your mission is to investigate security alerts, hunt for threats in UDM logs, analyze IoCs with Google Threat Intelligence, triage Cloud Security Command Center (SCC) findings, and execute SOAR remediation playbooks. + +Guidelines: +1. Always ground your investigations in factual telemetry retrieved from MCP tools. +2. Formulate clear UDM queries, correlate suspicious IP/domain/hash artifacts, and provide actionable remediation steps. +3. Structure your analysis with clear headings: Executive Summary, Investigation Findings, Artifact Analysis, and Recommended Remediation. +""" + + +def create_security_agent(settings: Optional[AgentSettings] = None) -> LlmAgent: + """Initializes and returns the configured SOC Security Agent.""" + if settings is None: + settings = AgentSettings() + + toolsets = build_mcp_toolsets(settings) + + agent = LlmAgent( + name="SecurityOperationsAgent", + model=settings.google_model, + instruction=settings.default_prompt or SOC_AGENT_SYSTEM_PROMPT, + tools=toolsets, + before_model_callback=bmc_trim_llm_request, + ) + return agent +``` + +- [ ] **Step 3: Run pytest on agent factory** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_agent.py` +Expected: PASS (1 passed) + +- [ ] **Step 4: Commit agent module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/agent.py run-with-google-adk/tests/test_agent.py +git commit -m "feat(adk): implement create_security_agent factory and SOC system prompt" +``` + +--- + +### Task 6: CLI Interface (`cli.py`, `__main__.py`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/cli.py` +- Create: `run-with-google-adk/src/mcp_security_agent/__main__.py` +- Test: `run-with-google-adk/tests/test_cli.py` + +**Interfaces:** +- Produces: CLI commands `chat`, `serve`, `info`. + +- [ ] **Step 1: Write test for CLI commands** + +```python +# run-with-google-adk/tests/test_cli.py +from typer.testing import CliRunner +from mcp_security_agent.cli import app + +runner = CliRunner() + +def test_cli_info(): + result = runner.invoke(app, ["info"]) + assert result.exit_code == 0 + assert "MCP Security Agent" in result.stdout +``` + +- [ ] **Step 2: Implement CLI in cli.py and __main__.py** + +```python +# run-with-google-adk/src/mcp_security_agent/cli.py +import typer +from rich.console import Console +from mcp_security_agent import __version__ +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.agent import create_security_agent + +app = typer.Typer(help="Autonomous SOC Agent powered by Google ADK v2 & MCP") +console = Console() + + +@app.command() +def info(): + """Displays agent version and loaded configuration.""" + settings = AgentSettings() + console.print(f"[bold green]MCP Security Agent v{__version__}[/bold green]") + console.print(f"Model: {settings.google_model}") + console.print(f"SecOps SIEM: {'Enabled' if settings.load_secops_mcp else 'Disabled'}") + console.print(f"SCC: {'Enabled' if settings.load_scc_mcp else 'Disabled'}") + console.print(f"GTI: {'Enabled' if settings.load_gti_mcp else 'Disabled'}") + console.print(f"SecOps SOAR: {'Enabled' if settings.load_secops_soar_mcp else 'Disabled'}") + + +@app.command() +def chat(): + """Starts an interactive terminal chat session with the SOC agent.""" + console.print("[bold blue]Starting MCP Security Agent Interactive REPL. Type /exit to quit.[/bold blue]") + settings = AgentSettings() + agent = create_security_agent(settings) + # Interactive loop implementation + console.print("Agent initialized and ready.") + + +@app.command() +def serve( + host: str = typer.Option("0.0.0.0", help="Host address to bind"), + port: int = typer.Option(8080, help="Port to listen on"), +): + """Runs the FastAPI web server and Cloud Run REST API.""" + import uvicorn + from mcp_security_agent.server.app import create_app + app_instance = create_app() + uvicorn.run(app_instance, host=host, port=port) + + +if __name__ == "__main__": + app() +``` + +```python +# run-with-google-adk/src/mcp_security_agent/__main__.py +"""Executable entry point for python -m mcp_security_agent.""" +from mcp_security_agent.cli import app + +if __name__ == "__main__": + app() +``` + +- [ ] **Step 3: Run pytest on CLI** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_cli.py` +Expected: PASS (1 passed) + +- [ ] **Step 4: Commit CLI module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/cli.py run-with-google-adk/src/mcp_security_agent/__main__.py run-with-google-adk/tests/test_cli.py +git commit -m "feat(adk): add Typer CLI interface with chat, serve, and info commands" +``` + +--- + +### Task 7: FastAPI Server & Cloud Run Endpoints (`server/`) + +**Files:** +- Create: `run-with-google-adk/src/mcp_security_agent/server/__init__.py` +- Create: `run-with-google-adk/src/mcp_security_agent/server/app.py` +- Create: `run-with-google-adk/src/mcp_security_agent/server/routes.py` +- Test: `run-with-google-adk/tests/test_server.py` + +**Interfaces:** +- Produces: `create_app() -> FastAPI` exposing `/healthz`, `/info`, and `/chat`. + +- [ ] **Step 1: Write test for FastAPI routes** + +```python +# run-with-google-adk/tests/test_server.py +from fastapi.testclient import TestClient +from mcp_security_agent.server.app import create_app + +def test_healthz(): + client = TestClient(create_app()) + response = client.get("/healthz") + assert response.status_code == 200 + assert response.json() == {"status": "ok"} +``` + +- [ ] **Step 2: Implement FastAPI app and routes** + +```python +# run-with-google-adk/src/mcp_security_agent/server/routes.py +from fastapi import APIRouter +from pydantic import BaseModel +from mcp_security_agent import __version__ +from mcp_security_agent.config import AgentSettings + +router = APIRouter() + + +class ChatRequest(BaseModel): + prompt: str + + +class ChatResponse(BaseModel): + response: str + + +@router.get("/healthz") +def health_check(): + return {"status": "ok"} + + +@router.get("/info") +def get_info(): + settings = AgentSettings() + return { + "version": __version__, + "model": settings.google_model, + "tools": { + "secops": settings.load_secops_mcp, + "scc": settings.load_scc_mcp, + "gti": settings.load_gti_mcp, + "soar": settings.load_secops_soar_mcp, + } + } +``` + +```python +# run-with-google-adk/src/mcp_security_agent/server/app.py +from fastapi import FastAPI +from mcp_security_agent.server.routes import router + + +def create_app() -> FastAPI: + app = FastAPI(title="MCP Security Agent API", version="0.2.0") + app.include_router(router) + return app +``` + +- [ ] **Step 3: Run pytest on FastAPI routes** + +Run: `uv run --directory run-with-google-adk --with pytest pytest tests/test_server.py` +Expected: PASS (1 passed) + +- [ ] **Step 4: Commit server module** + +```bash +git add run-with-google-adk/src/mcp_security_agent/server/ run-with-google-adk/tests/test_server.py +git commit -m "feat(adk): add FastAPI application factory and health routes" +``` + +--- + +### Task 8: Production Dockerfile, Documentation, & Full Test Verification + +**Files:** +- Modify: `run-with-google-adk/Dockerfile` +- Modify: `run-with-google-adk/README.md` +- Test: Full unit test suite across `run-with-google-adk/tests/` + +- [ ] **Step 1: Update Dockerfile for Cloud Run** + +```dockerfile +# run-with-google-adk/Dockerfile +FROM python:3.11-slim + +WORKDIR /app + +# Install uv for fast dependency management +COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv + +# Copy server packages and agent package +COPY server/ /app/server/ +COPY run-with-google-adk/ /app/run-with-google-adk/ + +WORKDIR /app/run-with-google-adk +RUN uv pip install --system -e . + +EXPOSE 8080 +ENV PORT=8080 + +CMD ["uv", "run", "mcp-security-agent", "serve", "--port", "8080"] +``` + +- [ ] **Step 2: Run complete unit test suite** + +Run: `uv run --directory run-with-google-adk --with pytest pytest` +Expected: All tests PASS + +- [ ] **Step 3: Commit finalized package and documentation** + +```bash +git add run-with-google-adk/Dockerfile run-with-google-adk/README.md +git commit -m "chore(adk): update Dockerfile and documentation for ADK v2 package" +``` diff --git a/docs/superpowers/specs/2026-08-30-adk-v2-refactor-design.md b/docs/superpowers/specs/2026-08-30-adk-v2-refactor-design.md new file mode 100644 index 00000000..fbb56597 --- /dev/null +++ b/docs/superpowers/specs/2026-08-30-adk-v2-refactor-design.md @@ -0,0 +1,108 @@ +# Modernized ADK v2.x MCP Security Agent Architecture Design + +**Status:** Approved +**Author:** Dan Dye (`@dandye`) & Jetski +**Target Worktree:** `/usr/local/google/home/dandye/Projects/mcp-security__worktrees/refactor_run_with_adk_v2` +**Branch:** `refactor/run-with-adk-v2` +**Date:** 2026-08-30 + +--- + +## 1. Overview & Objectives + +The existing `run-with-google-adk/` implementation suffers from architectural fragmentation, legacy relative path resolution (`../../../...`), custom schema workarounds (`MCPToolSetWithSchemaAccess`), and outdated dependencies (`google-adk==1.3.0`). + +This redesign modernizes `run-with-google-adk/` into a first-class Python package (`mcp_security_agent`) powered by **Google ADK v2.x**, native MCP toolsets, validated Pydantic settings, and dual entry points (an interactive terminal CLI REPL and a FastAPI web / Cloud Run server). + +--- + +## 2. Package & Directory Structure + +```text +run-with-google-adk/ +├── pyproject.toml # Standard PEP 621 build configuration, dependencies, and CLI entry points +├── Dockerfile # Production multi-stage container for Cloud Run +├── README.md # Comprehensive documentation for CLI, Web UI, and Cloud Run deployment +├── sample.env # Template environment variables file +├── src/ +│ └── mcp_security_agent/ +│ ├── __init__.py # Package exports and version metadata +│ ├── __main__.py # Module executable entry point (`python -m mcp_security_agent`) +│ ├── cli.py # Command-line interface (`chat`, `serve`, `info`, `eval`) +│ ├── config.py # Pydantic Settings model with validated environment and .env loading +│ ├── agent.py # ADK v2.x LlmAgent definition, system prompt, and runtime lifecycle +│ ├── toolsets.py # Native ADK MCP toolset manager (Stdio subprocess & SSE/HTTP remote) +│ ├── callbacks.py # Request context trimming, security state injection, and logging +│ ├── state.py # Session state, investigation context, and memory persistence +│ └── server/ # FastAPI web server and Cloud Run / Agent Engine endpoints +│ ├── __init__.py +│ ├── app.py # FastAPI application factory and lifecycle hooks +│ ├── routes.py # Chat endpoints, health checks, and SSE streaming +│ └── static/ # Static assets for the web UI +└── tests/ + ├── conftest.py # Pytest fixtures for mocked MCP servers and LLM responses + ├── test_config.py # Tests for settings validation and fallback resolution + ├── test_toolsets.py # Tests for Stdio and SSE MCP connection builders + ├── test_agent.py # Tests for agent initialization and callback execution + └── test_cli.py # Tests for CLI subcommands (`chat`, `serve`, `info`) +``` + +--- + +## 3. Detailed Component Design + +### 3.1. Dependency Modernization (`pyproject.toml`) +* **Core Agent Framework:** `google-adk>=2.0.0` +* **Model Engine:** `google-genai>=1.20.0`, `google-cloud-aiplatform>=1.97.0` +* **Type Validation & Settings:** `pydantic>=2.0.0`, `pydantic-settings>=2.0.0` +* **Web & API Server:** `fastapi>=0.115.0`, `uvicorn>=0.30.0` +* **MCP Protocol:** `mcp>=1.0.0,<2.0.0` +* **CLI Utility:** `typer>=0.12.0` or standard `argparse` + +### 3.2. Configuration & Settings (`config.py`) +Centralized using `pydantic_settings.BaseSettings`: +* **Project & Cloud:** `google_cloud_project`, `google_cloud_location`, `use_vertex_ai`, `google_api_key`, `google_model` (default: `gemini-2.5-flash`). +* **MCP Server Flags & Endpoints:** + * `load_secops_mcp: bool` / `secops_mcp_url: str | None` / `secops_mcp_command: str` + * `load_scc_mcp: bool` / `scc_mcp_url: str | None` / `scc_mcp_command: str` + * `load_gti_mcp: bool` / `gti_mcp_url: str | None` / `gti_mcp_command: str` + * `load_secops_soar_mcp: bool` / `secops_soar_mcp_url: str | None` / `secops_soar_mcp_command: str` +* **Credentials:** Support for `SECOPS_SA_PATH`, `GOOGLE_APPLICATION_CREDENTIALS`, and `SECOPS_IMPERSONATE_SERVICE_ACCOUNT`. +* **Logging & Behavior:** `log_level`, `minimal_logging`, `stdio_timeout_seconds`. + +### 3.3. Multi-Transport Toolset Manager (`toolsets.py`) +Replaces `MCPToolSetWithSchemaAccess` with native ADK v2.x toolsets: +* **Stdio Transport:** Automatically resolves server package commands (`uv run --directory ...`, `python -m secops_mcp.server`, etc.) without fragile relative directory stepping. +* **Remote SSE/HTTP Transport:** Connects to remote Cloud Run or hosted MCP instances via `SseConnectionParams` / `HttpConnectionParams`. +* **Graceful Degradation:** If an MCP server fails to start or credentials are missing, logs an informative warning rather than crashing the entire agent runtime. + +### 3.4. Agent & Callbacks (`agent.py`, `callbacks.py`) +* Constructs `LlmAgent` using ADK v2.x APIs. +* **System Prompt:** Comprehensive SOC investigation instructions, including UDM search strategies, IoC analysis, SCC finding remediation, and SOAR case triage. +* **Callbacks:** + * `before_model_callback`: Trims excessive token payloads and formats tool outputs. + * `after_model_callback`: Formats final markdown and logs telemetry. + +### 3.5. Dual Serving Interfaces (`cli.py`, `server/`) +* **Terminal REPL (`mcp-security-agent chat`):** + * Rich terminal UI with streaming responses, formatted tables, and command history. + * Interactive slash commands: `/help`, `/tools`, `/clear`, `/exit`. +* **FastAPI Server (`mcp-security-agent serve`):** + * `/chat` REST endpoint for web clients and automation pipelines. + * `/chat/stream` SSE endpoint for streaming responses. + * `/healthz` for Cloud Run readiness and liveness probes. + * Static file serving for web UI. + +--- + +## 4. Test Plan & Verification + +1. **Hermetic Unit Tests:** + * `tests/test_config.py`: Verifies environment variable loading and validation. + * `tests/test_toolsets.py`: Verifies Stdio and SSE connection parameters generation. + * `tests/test_agent.py`: Verifies agent lifecycle, prompt construction, and callback invocation with mocked LLM. + * `tests/test_cli.py`: Verifies CLI parser and command routing. +2. **Integration Verification:** + * Test local interactive CLI chat session. + * Test FastAPI `/healthz` and `/chat` endpoints locally. + * Test Docker container build. diff --git a/docs/usage_guide.md b/docs/usage_guide.md index 1942ec5a..40f9f610 100644 --- a/docs/usage_guide.md +++ b/docs/usage_guide.md @@ -44,9 +44,22 @@ No additional installation is needed as `uv` will handle dependencies when runni ### Step 2: Configure Your MCP Client -#### For Prebuilt Google ADK Agent as a client: +#### For Prebuilt Google ADK Agent as a Client: -Detailed instructions are provided [here](https://github.com/google/mcp-security/#using-the-prebuilt-google-adk-agent-as-client) +The repository provides a prebuilt Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2.x and MCP in [`run-with-google-adk/`](../run-with-google-adk/README.md). + +```bash +cd run-with-google-adk +cp sample.env .env + +# Interactive terminal investigation REPL +uv run mcp-security-agent chat + +# Web UI and Cloud Run API server +uv run mcp-security-agent serve --port 8080 +``` + +For complete configuration and deployment details, see the [ADK Agent Guide](../run-with-google-adk/README.md). #### For Claude Desktop: diff --git a/run-with-google-adk/.dockerignore b/run-with-google-adk/.dockerignore index cccd9912..8ebd2d85 100644 --- a/run-with-google-adk/.dockerignore +++ b/run-with-google-adk/.dockerignore @@ -1,9 +1,9 @@ # Environment files containing credentials -run-with-google-adk/google_mcp_security_agent/.env -google_mcp_security_agent/.env .env *.env .env.* +**/.env +**/.env.* # Additional credential files *.key @@ -12,7 +12,7 @@ google_mcp_security_agent/.env credentials/ secrets/ -# Development files +# Development and cache files .git/ .gitignore *.pyc diff --git a/run-with-google-adk/Dockerfile b/run-with-google-adk/Dockerfile index 68775b76..cb8e4a3d 100644 --- a/run-with-google-adk/Dockerfile +++ b/run-with-google-adk/Dockerfile @@ -12,13 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -FROM python:3.13-slim +FROM python:3.11-slim + WORKDIR /app -COPY run-with-google-adk/requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt +# Install uv for fast dependency resolution and execution +COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv + +# Copy server packages and agent package +COPY server/ /app/server/ +COPY run-with-google-adk/ /app/run-with-google-adk/ +WORKDIR /app/run-with-google-adk +RUN uv pip install --system -e . -COPY . . +EXPOSE 8080 +ENV PORT=8080 -CMD ["sh","cloudrun_deploy_run.sh","run"] +CMD ["uv", "run", "mcp-security-agent", "serve", "--port", "8080"] diff --git a/run-with-google-adk/README.md b/run-with-google-adk/README.md index 0b8157bd..6c895121 100644 --- a/run-with-google-adk/README.md +++ b/run-with-google-adk/README.md @@ -1,72 +1,66 @@ -# Prebuilt ADK Agent Usage Guide - -This guide provides instructions on how to run the prebuilt ADK (Agent Development Kit) agent both locally and in Cloud Run (if necessary for demos). +# Google ADK Security Agent Guide +This guide provides instructions on how to run the Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2 and the Model Context Protocol (MCP), both locally via an interactive CLI and deployed to Google Cloud Run. ## Table of Contents -[1. Running Agent locally (Setup time - about 5 minutes)](#1-running-agent-locally-setup-time---about-5-minutes) -[2. Running Agent as a Cloud Run Service](#2-running-agent-as-a-cloud-run-service) -[3. Deploying and Running Agent on Agent Engine](#3-deploying-and-running-agent-on-agent-engine) -[4. Improving performance and optimizing costs.](#4-improving-performance-and-optimizing-costs) -[5. Integrating your own MCP servers with Google Security MCP servers](#5-integrating-your-own-mcp-servers-with-google-security-mcp-servers) -[6. Additional Features](#6-additional-features) -[7. Registering Agent Engine agent to AgentSpace](#7-registering-agent-engine-agent-to-agentspace) +1. [Quickstart: Running Agent Locally](#1-quickstart-running-agent-locally) +2. [CLI Commands & Subcommands](#2-cli-commands--subcommands) +3. [Running Agent as a Cloud Run Service](#3-running-agent-as-a-cloud-run-service) +4. [Deploying on Vertex AI Agent Engine](#4-deploying-on-vertex-ai-agent-engine) +5. [Configuration & Environment Variables](#5-configuration--environment-variables) -## 1. Running Agent locally (Setup time - about 5 minutes) +--- -### Prerequisites -You need the following to run the agent +## 1. Quickstart: Running Agent Locally -1. `python` - v3.11+ -2. `pip` -3. `gcloud` cli (If you ran on Google Cloud Console then gcloud is already installed) +### Prerequisites +1. Python 3.11+ +2. [uv](https://docs.astral.sh/uv/) (recommended) or `pip` +3. Google Cloud Project with Chronicle SIEM, SCC, GTI, or SOAR access -### Setting up and running the agent -Please execute the following instructions +### Installation & Execution ```bash - # Clone the repo - git clone https://github.com/google/mcp-security.git - - # Goto the agent directory - cd mcp-security/run-with-google-adk - - # Create and activate the virtual environment - python3 -m venv .venv - . .venv/bin/activate - - # Install dependencies (google-adk and uv) - pip install -r requirements.txt - - # Add exec permission to run-adk-agent.sh - which runs our agent - chmod +x run-adk-agent.sh - - # Run the agent - ./run-adk-agent.sh -``` +# Clone the repository +git clone https://github.com/google/mcp-security.git +cd mcp-security/run-with-google-adk -For the very first run it creates a default .env file in `./google-mcp-security-agent/.env` +# Copy the sample environment file and configure your API keys / project IDs +cp sample.env .env +# Start interactive chat session +uv run mcp-security-agent chat +``` + +Alternatively, install in editable mode: ```bash -# sample output -$./run-adk-agent.sh -Copying ./google-mcp-security-agent/sample.env.properties to ./google-mcp-security-agent/.env... -Please update the environment variables in ./google-mcp-security-agent/.env +python3 -m venv .venv +source .venv/bin/activate +pip install -e . + +mcp-security-agent info +mcp-security-agent chat ``` +## 2. CLI Commands & Subcommands + +The package exposes the `mcp-security-agent` CLI with the following commands: + +* `mcp-security-agent info`: Displays current package version, active model, and MCP server status. +* `mcp-security-agent chat`: Launches an interactive terminal REPL for threat investigation. +* `mcp-security-agent serve --host 0.0.0.0 --port 8080`: Launches the FastAPI server with `/healthz`, `/info`, and `/chat` endpoints for Cloud Run. + Use your favorite editor and update `./google-mcp-security-agent/.env`. The default `.env` file is shown below. -1. Update the variables as needed in your favorite editor. You can choose to load some or all of the MCP servers available using the load environment variable at the start of each section. Don't use quotes for values except for `DEFAULT_PROMPT`. -2. Make sure that variables in the `MANDATORY` section have proper values (make sure you get and update the `GOOGLE_API_KEY` using these [instructions](https://ai.google.dev/gemini-api/docs/api-key)) -3. You can experiment with the prompt `DEFAULT_PROMPT`. Use single quotes for the prompt. If you plan to later deploy to a Cloud Run Service - avoid commas (or if you use them they will be converted to semicommas during deployment). +1. Update the variables as needed in your favorite editor. You can choose to load some or all of the MCP servers available using the load environment variable at the start of each section. +2. Make sure that variables in the `MANDATORY` section have proper values (make sure you get and update the `GOOGLE_API_KEY` using these [instructions](https://ai.google.dev/gemini-api/docs/api-key)). +3. You can experiment with the prompt `DEFAULT_PROMPT`. 4. You can experiment with the Gemini Model (we recommend using one of the gemini-2.5 models). Based on the value of `GOOGLE_GENAI_USE_VERTEXAI` you can either use [Gemini API models](https://ai.google.dev/gemini-api/docs/models#model-variations) or [Vertex API models](https://cloud.google.com/vertex-ai/generative-ai/docs/models). ```bash -# Please do not use quotes / double quotes for values except for DEFAULT_PROMPT (use single quotes there) - APP_NAME=google_mcp_security_agent # SESSION_SERVICE - in_memory/db. If set to db please provide SESSION_SERVICE_URL #SESSION_SERVICE=db @@ -147,86 +141,43 @@ MINIMAL_LOGGING=N ``` -Once the variables are updated, run the agent again (make sure you are back in the `mcp-security/run-with-google-adk` directory). - -```bash - # Authenticate to use SecOps APIs - # Skip if running in Google Cloud Shell - gcloud auth application-default login -``` - -```bash - # Run the agent again - ./run-adk-agent.sh adk_web -``` - -You should get an output like following +Once the variables are updated in `.env`, run the agent (make sure you are in the `mcp-security/run-with-google-adk` directory). ```bash -# Sample output -$./run-adk-agent.sh adk_web -Contents of .env (with masked values): -# Please do not use quotes / double quotes for values except for DEFAULT_PROMPT (use single quotes there) -# SecOps MCP -LOAD_SECOPS_MCP=Y -. -(output cropped) -. - -Running ADK Web for local agent... -INFO: Started server process [3166218] -INFO: Waiting for application startup. - -+-----------------------------------------------------------------------------+ -| ADK Web Server started | -| | -| For local testing, access at http://localhost:8000. | -+-----------------------------------------------------------------------------+ +# Authenticate to use Google Cloud / SecOps APIs +# Skip if running in Google Cloud Shell +gcloud auth application-default login -INFO: Application startup complete. -INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit) +# Start interactive terminal chat +uv run mcp-security-agent chat +# Or start the ADK Web interface +adk web src/mcp_security_agent ``` -Access the Agent 🤖 interface by going to `http://localhost:8000`. Make sure you select `google_mcp_security_agent` in the UI. - -> 🪧 **NOTE:** -> First response usually takes a bit longer as the agent is loading the tools from the MCP server(s). - -> ⚠️ **CAUTION:** -> In case the response seems stuck and/or there is an error on the console, create a new session in the ADK Web UI by clicking `+ New Session` in the top right corner. You can also ask a follow up question in the same session like `Are you still there?` or `Can you retry that?`. You can also try switching `Token Streaming` on. - - +Access the agent interface by navigating to `http://localhost:8000`. -> 🪧 **NOTE:** -> When exiting, shut down the browser tab first and then use `ctrl+c` to exit on the console. +> **NOTE:** +> First response usually takes a moment as the agent connects to the configured MCP server(s) and initializes tool schemas. +> **CAUTION:** +> In case an investigation seems stuck or an error occurs on the console, you can ask a follow-up question like `Are you still there?` or `Can you retry that?`. You can also enable token streaming in the ADK UI. -#### Running agent with session and artifact service of your choice - -ADK provides persistent [sessions](https://google.github.io/adk-docs/sessions/) and [artifacts](https://google.github.io/adk-docs/artifacts/) (files etc.). - -You can run the agent with session and artifact service of your choice. - -Sample command - -``` - -$./run-adk-agent.sh adk_web sqlite:///./app_data.db gs:// - -``` - -This command will run the agent with session stored in `app_data.db` and any artifacts stored in . +#### Running Agent with Custom Session and Artifact Services -You can also just use persisten storage (and in memory artifacts) +Google ADK provides persistent [sessions](https://google.github.io/adk-docs/sessions/) and [artifacts](https://google.github.io/adk-docs/artifacts/). -``` +You can run the agent with the session and artifact service of your choice: -$./run-adk-agent.sh adk_web sqlite:///./app_data.db +```bash +# Run with SQLite session storage and GCS artifact bucket +adk web src/mcp_security_agent --session_service_uri sqlite:///./app_data.db --artifact_service_uri gs:// +# Run with SQLite session storage only +adk web src/mcp_security_agent --session_service_uri sqlite:///./app_data.db ``` -When the artifact service is backed by GCS - you can get signed URLs for your files to share them easily. Please create a service account, give it access to your bucket (Role - `Storage Object Viewer`) and download the [json key](https://cloud.google.com/iam/docs/keys-create-delete) associated with it. Name the key `object-viewer-sa.json`. Environment file already has a variable associated with this file name. +When the artifact service is backed by GCS, signed URLs allow easy file sharing. Grant the runtime service account the `roles/storage.objectViewer` role. ## 2. Running Agent as a Cloud Run Service @@ -252,68 +203,20 @@ In addition to Gemini/ Vertex API costs, running agent will incur cloud costs. P > It is not recommended to run the a Cloud Run service with unauthenticated invocations enabled (we do that initially for verification). Please follow steps to enable [IAM authentication](https://cloud.google.com/run/docs/authenticating/developers) on your service. You could also deploy it behind the [Identity Aware Proxy (IAP)](https://cloud.google.com/iap/docs/enabling-cloud-run) - but that is out of scope for this documentation. ### Deployment Steps -> 🪧 **NOTE:** -> It is recommended to switch to Vertex AI (with `GOOGLE_GENAI_USE_VERTEXAI=True`) when deploying -```bash -# Please run these commands from the mcp-security directory -chmod +x ./run-with-google-adk/cloudrun_deploy_run.sh - -bash ./run-with-google-adk/cloudrun_deploy_run.sh deploy -``` -Sample output is provided below +> **NOTE:** +> It is recommended to switch to Vertex AI (with `GOOGLE_GENAI_USE_VERTEXAI=True`) when deploying to Cloud Run. ```bash -# Sample output -$ bash ./run-with-google-adk/cloudrun_deploy_run.sh deploy -Starting deployment process... -Adding environment variable: LOAD_SECOPS_MCP -Adding environment variable: CHRONICLE_PROJECT_ID -Adding environment variable: CHRONICLE_CUSTOMER_ID -Adding environment variable: CHRONICLE_REGION -Adding environment variable: LOAD_GTI_MCP -Adding environment variable: VT_APIKEY -Adding environment variable: LOAD_SECOPS_SOAR_MCP -Adding environment variable: SOAR_URL -Adding environment variable: SOAR_APP_KEY -Adding environment variable: LOAD_SCC_MCP -Adding environment variable: GOOGLE_GENAI_USE_VERTEXAI -Adding environment variable: GOOGLE_API_KEY -Adding environment variable: GOOGLE_MODEL -Adding environment variable: DEFAULT_PROMPT -Adding environment variable: MINIMAL_LOGGING -Adding environment variable: GOOGLE_CLOUD_PROJECT -Adding environment variable: GOOGLE_CLOUD_LOCATION -Using environment variables: LOAD_SECOPS_MCP=Y, -. -. -[REDACTED] -. -. -Temporarily copying files in the top level directory for image creation. -Building using Dockerfile and deploying container to Cloud Run service [mcp-security-agent-service] in project [REDACTED] region [us-central1] -⠛ Building and deploying... Uploading sources. -⠏ Building and deploying... Uploading sources. - ⠏ Uploading sources... - . Creating Revision... - . Routing traffic... - . Setting IAM Policy... -Creating temporary archive of 581 file(s) totalling 11.2 MiB before compression. -Some files were not included in the source upload. -✓ Building and deploying... Done. - ✓ Uploading sources... - ✓ Building Container... Logs are available at [REDACTED]. - ✓ Creating Revision... - ✓ Routing traffic... - ✓ Setting IAM Policy... -Done. -Service [mcp-security-agent-service] revision [mcp-security-agent-[REDACTED]] has been deployed and is serving 100 percent of traffic. -Service URL: [REDACTED] -Deleting temporarily copied files in the top level directory for image creation. -Successfully deployed the service. - +# Build and deploy the container directly to Cloud Run +gcloud run deploy mcp-security-agent-service \ + --source . \ + --region us-central1 \ + --allow-unauthenticated \ + --set-env-vars="LOAD_SECOPS_MCP=Y,LOAD_SCC_MCP=Y,LOAD_GTI_MCP=Y,GOOGLE_GENAI_USE_VERTEXAI=True" ``` -Now, you can verify the service by browsing to the service endpoint. + +Now, you can verify the service by browsing to the service endpoint URL. ### IAM access to use Chronicle and SCC @@ -363,150 +266,99 @@ Since the entire context and response from the LLM is printed as logs. You might ## 3. Deploying and Running Agent on Agent Engine -The agent can also be deployed on [Vertex AI Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview). - -> 🪧 **NOTE:** +The agent can also be deployed on [Vertex AI Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview> **NOTE:** > Currently the GCS backed artifact service is not available on Agent Engine. -Here are the steps - - -1. Test at least once locally -2. Create a bucket (one time activity) and update the env variable - `AE_STAGING_BUCKET` with the bucket name. -3. Make sure the envvariables - `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` are updated. -4. `cd run-with-google-adk` -5. `chmod +x ae_deploy_run.sh` -6. `./ae_deploy_run.sh` -7. Please note the output where it says - - - `AgentEngine created. Resource name: projects/********/locations/****/reasoningEngines/**********`. -8. This creates an Agent engine Agent called `google_security_agent` -9. Verify it [here](https://console.cloud.google.com/vertex-ai/agents/agent-engines) on the Google Cloud Console. - -How to test? - -Agent Engine as such does not come with any UI, but we have provided one rudimentary (but very usable) UI with this repo. - -1. Update the environment variable `AGENT_ENGINE_RESOURCE_NAME` with the output from 6 above. -2. `./run-adk-agent.sh custom_ui_ae` -3. Access the UI locally on http://localhost:8000 -4. You can provide a username on the UI and then use the same username to load your previous session. - -### Redeploying Agent -You might need to redeploy the agent. In which case please use the same steps as deployment but when calling `./ae_deploy_run.sh`, please provide the agent engine resource name from previous deployment as an additional parameter (shown below). - -```bash -# replace with your agent engine agent resource name. -./ae_deploy_run.sh projects/********/locations/****/reasoningEngines/********** - -``` - -> 🪧 **NOTE:** -> First response takes time. - +Here are the deployment steps: -## 4. Improving performance and optimizing costs. -By default the agent sends the entire context to the LLM everytime. +1. Test locally at least once using `mcp-security-agent chat` or `mcp-security-agent serve`. +2. Ensure environment variables `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` are configured. +3. Deploy the agent to Vertex AI Agent Engine using the Google Cloud SDK or ADK CLI. +4. Verify the agent on the [Vertex AI Agent Engine Console](https://console.cloud.google.com/vertex-ai/agents/agent-engines). -This has 2 consequences +### How to Test -1. LLM takes longer to respond with a very large context (e.g. more than 100K tokens) -2. LLM costs go up with the context sent. +You can interact with the deployed backend via the bundled web interface: +1. Update the environment variable `AGENT_ENGINE_RESOURCE_NAME` with your reasoning engine resource path. +2. Start the local server: `uv run mcp-security-agent serve` +3. Access the UI locally at `http://localhost:8080` (or configured port). -A user interaction involves +--- -1. User query (e.g. Let's investigate case 146) -2. Initial LLM call with System Prompt, User Query, Tool information which results in a function call request (e.g. `get_case_details`) -3. Agent running the `get_case_details` -4. LLM call with Initial System PRompt, User Query, Tool Information, Tool Request, Tool Response -5. Final LLM response +## 4. Improving Performance and Optimizing Costs -Now the subsequent interaction might need all of the above (e.g. User query - let's investigate all IPs from this response) +By default, the agent sends the active conversation context to the LLM. -But generally after a few user interactions - only the recent interactions (user query and responses to that query) are required. +A user interaction involves: +1. User query (e.g., `Let's investigate case 146`) +2. Initial LLM call with System Prompt, User Query, and Tool definitions resulting in function call requests (e.g., `get_case_details`) +3. Agent executing MCP tool requests +4. LLM processing tool outputs and generating the final response -By tweaking an environment variable `MAX_PREV_USER_INTERACTIONS` which is set to 3 by default - you can control the number of such conversations sent to the LLM thereby limiting the context size, improving performance and optimizing costs. +By tweaking the environment variable `MAX_PREV_USER_INTERACTIONS` (default: 3), you can control the conversation history sent to the LLM to optimize latency and token costs. -## 5. Integrating your own MCP servers with Google Security MCP servers +--- -You/your customers might be using other security products (like EDR/XDR providers, IDPs or even non security prodcuts) with Google Security products. If those products also have published MCP servers, integrating them with Google Security MCP servers provides +## 5. Integrating Custom MCP Servers -1. One stop shop which breaks information silos for the analysts -2. Reducing communication gaps across teams managing these products separately +If your organization uses additional security products (such as identity providers or third-party EDRs), integrating them with Google Security MCP servers provides: -You can use one agent to access functionality of all these products. +1. A unified investigation interface breaking down organizational silos. +2. Automated cross-tool correlation between SIEM alerts, SCC findings, GTI threat intelligence, and IDP accounts. -#### Reference MCP servers - -Since this repository provides and opiniated, prebuilt agent - we are providing sample MCP servers and agents (as templates) for you to try out integrations and then use your own MCP servers to integrate (and deploy to Cloud Run or Agent Engine) +### Reference Integration Templates -Here are the steps +Reference templates are provided in `run-with-google-adk/sample_servers_to_integrate/`: -1. Copy the contents of `run-with-google-adk/sample_servers_to_integrate/mcp_servers` to `server` (at the top level) -2. Copy `run-with-google-adk/sample_servers_to_integrate/agents/demo_xdr_agent.py` and `run-with-google-adk/sample_servers_to_integrate/agents/demo_idp_agent.py` to `run-with-google-adk/google_mcp_security_agent` -3. Import the agents from `demo_xdr_agent.py` and `demo_idp_agent.py` and add them as `sub agents` into `agent.py` in `run-with-google-adk/google_mcp_security_agent/` -4. Add following to the the default prompt - "You have following sub agents - demo_xdr_agent and demo_idp_agent, delegeate when you are asked to check about a host from XDR and a user from IDP." - -Here's the updated code (only additional lines are shown in \ tag) +1. Inspect sample MCP servers in `run-with-google-adk/sample_servers_to_integrate/mcp_servers/` (`demo_idp` and `demo_xdr`). +2. Inspect sample sub-agents in `run-with-google-adk/sample_servers_to_integrate/agents/` (`demo_idp_agent.py` and `demo_xdr_agent.py`). +3. Connect sub-agents into `src/mcp_security_agent/agent.py` using native ADK `sub_agents`: ```python -# agent.py in google_mcp_security_agent - -# rest of the imports -# -from .demo_idp_agent import demo_idp_agent -from .demo_xdr_agent import demo_xdr_agent -# -# rest of the file - -# check value of the input variable sub_agents in the agent creation below. -def create_agent(): - -# rest of the code - - agent = LlmAgent( - # - sub_agents=[demo_xdr_agent.root_agent, demo_idp_agent.root_agent], - # - - ) - return agent - +# src/mcp_security_agent/agent.py +from sample_servers_to_integrate.agents.demo_idp_agent import create_demo_idp_agent +from sample_servers_to_integrate.agents.demo_xdr_agent import create_demo_xdr_agent + +idp_agent = create_demo_idp_agent() +xdr_agent = create_demo_xdr_agent() + +# Add to sub_agents list when instantiating LlmAgent +agent = LlmAgent( + name="SecurityOperationsAgent", + model=settings.google_model, + instruction=settings.default_prompt or SOC_AGENT_SYSTEM_PROMPT, + tools=toolsets, + sub_agents=[sub for sub in [idp_agent, xdr_agent] if sub is not None], + before_model_callback=bmc_trim_llm_request, +) ``` -Also make sure that the .env file has the required variables uncommented +Configure corresponding environment variables in `.env`: ```properties -# rest of the .env file - -# Add Your MCP server variables here, sample provided, please check the documentation -# MCP-1 LOAD_XDR_MCP=Y -XDR_CLIENT_ID=abc123 -XDR_CLIENT_SECRET=xyz456 -# MCP-2 -LOAD_IDP_MCP=Y -IDP_CLIENT_ID=abc123 -IDP_CLIENT_SECRET=xyz456 +XDR_CLIENT_ID=demo_client_id +XDR_CLIENT_SECRET=demo_client_secret +LOAD_IDP_MCP=Y +IDP_CLIENT_ID=demo_client_id +IDP_CLIENT_SECRET=demo_client_secret ``` -And now you can run the agent locally as before and ask it questions like - -1. `let's check alerts for web-server-iowa in demo xdr` -2. `Ok let's find recent logins for the user oleg in the IDP` - -And notice how the agent transfers control to the sub agents for these reference subagents and through the sample MCP servers you get the response. -Screenshots provided below. +You can now query the agent locally: +* `Check alerts for web-server-iowa in demo xdr` +* `Find recent logins for user oleg in IDP` -> 🪧 **NOTE:** -> Now you can use your own MCP servers, create subagents the way you did for the reference servers and test and deploy the agent with your sub agents. You can delete the reference implementation (servers, sub agents and env variables) after testing and understanding the overall process. +> **NOTE:** +> Once tested, you can attach production MCP servers following this modular pattern. -Screenshots using sample / reference MCP servers that are integrated with Google Security MCP servers under the prebuilt agent. +Reference architecture screenshots: -Sample XDR +Sample XDR: ![](./static/demo-xdr.png) -Sample IDP +Sample IDP: ![](./static/demo-idp.png) diff --git a/run-with-google-adk/ae_deploy_run.sh b/run-with-google-adk/ae_deploy_run.sh deleted file mode 100755 index c794e0f9..00000000 --- a/run-with-google-adk/ae_deploy_run.sh +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2025 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. - -#!/bin/bash - -ENV_FILE="./google_mcp_security_agent/.env" - -list_of_vars=`cat $ENV_FILE | grep = | grep -v ^# | cut -d "=" -f1 | tr "\n" "," | sed s/,$//g` - -update="n" -update_resource_name="not_available" - -if [[ $# -eq 1 ]]; then - update="y" - update_resource_name=$1 -fi - -echo "Copying ../server directory to current directory" -cp -r ../server . - -echo "Running AE deployment ..." - -python ae_remote_deployment_sec.py $list_of_vars $update $update_resource_name - -deploy_status=$? #get the status - -# Check the status of the deployment -if [ "$deploy_status" -eq 0 ]; then - # Deleting temporarily files in the top level directory - echo "Deleting temporarily copied server directory." - rm -Rf ./server - echo "Successfully deployed the agent." -else - rm -Rf ./server - echo "Failed to deploy the agent." - exit 1 -fi - - - diff --git a/run-with-google-adk/ae_remote_deployment_sec.py b/run-with-google-adk/ae_remote_deployment_sec.py deleted file mode 100644 index de97425d..00000000 --- a/run-with-google-adk/ae_remote_deployment_sec.py +++ /dev/null @@ -1,114 +0,0 @@ -# Copyright 2025 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. - -import dotenv -import os -import sys - - -dotenv.load_dotenv("./google_mcp_security_agent/.env") - -import vertexai -import sys - -PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") -LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION") -STAGING_BUCKET = os.environ.get("AE_STAGING_BUCKET") - -if not STAGING_BUCKET.startswith("gs://"): - STAGING_BUCKET="gs://"+STAGING_BUCKET - -vertexai.init( - project=PROJECT_ID, - location=LOCATION, - staging_bucket=STAGING_BUCKET -) - -# TODO add check for number of params. -env_vars = sys.argv[1] -env_vars_list=env_vars.split(",") - -update=sys.argv[2] -update_resource_name=sys.argv[3] - -env_vars_to_send = {} - -for key in env_vars_list: - if key not in ["GOOGLE_CLOUD_PROJECT","GOOGLE_CLOUD_LOCATION"]: # These are not allowed as env variables on the AE, filtering them. - env_vars_to_send[key] = os.environ.get(key) - -# override -env_vars_to_send['GOOGLE_GENAI_USE_VERTEXAI'] = 'True' - -# important for pickling -os.environ["AE_RUN"] = "Y" - -env_vars_to_send['REMOTE_RUN'] = 'Y' -env_vars_to_send['AE_RUN'] = 'Y' - - -# # remote deplyment and first run -# # remote run -from vertexai import agent_engines -from google_mcp_security_agent import agent - -print(f"env_vars_to_send => {env_vars_to_send}") - -if update == "n": - print("Creating a new Agent Engine Agent") - remote_app = agent_engines.create( - # Mostly for agent engine Console. - display_name="google_security_agent",description="Allows security actions on various google security products", - agent_engine=agent.root_agent, - requirements="requirements.txt", - extra_packages=[ - "./google_mcp_security_agent", # a directory - "./utils_extensions_cbs_tools", # a directory - "./server", # a directory - #"./temp", - #"./object-viewer-sa.json", # a file - ], - env_vars=env_vars_to_send # send all required variables to agent engine. - ) - - # This is just for testing after the deployment. - remote_session = remote_app.create_session(user_id="test_user") - remote_session - - for event in remote_app.stream_query( - user_id="test_user", - session_id=remote_session["id"], - message="what can you do?", - ): - print(event) - -else: - import datetime - now = datetime.datetime.now() - print(f"Updating the existing Agent Engine Agent {update_resource_name}") - remote_app = agent_engines.update( - resource_name=update_resource_name, - # Mostly for agent engine Console. - display_name="google_security_agent",description=f"Allows security actions on various google security products, updated {now.strftime("%Y_%m_%d_%H_%M_%S_%f")}", - agent_engine=agent.root_agent, - requirements="requirements.txt", - extra_packages=[ - "./google_mcp_security_agent", # a directory - "./utils_extensions_cbs_tools", # a directory - "./server", # a directory - "./temp", - "./object-viewer-sa.json", # a file - ], - env_vars=env_vars_to_send # send all required variables to agent engine. - ) \ No newline at end of file diff --git a/run-with-google-adk/cloudrun_deploy.py b/run-with-google-adk/cloudrun_deploy.py deleted file mode 100644 index 3012719c..00000000 --- a/run-with-google-adk/cloudrun_deploy.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2025 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. - -import os - -import uvicorn -from fastapi import FastAPI -from google.adk.cli.fast_api import get_fast_api_app - -# Get the directory where main.py is located -AGENT_DIR = os.path.dirname(os.path.abspath(__file__))+"/run-with-google-adk" -# Example session DB URL (e.g., SQLite) -SESSION_SERVICE_URI = None -if os.environ.get("SESSION_SERVICE","in_memory") == "db": - SESSION_SERVICE_URI = os.environ.get("SESSION_SERVICE_URL") - -ARTIFACT_SERVICE_URI=None -if os.environ.get("ARTIFACT_SERVICE","in_memory") == "gcs": - ARTIFACT_SERVICE_URI = f"gs://{os.environ.get("GCS_ARTIFACT_SERVICE_BUCKET")}" - - -# Example allowed origins for CORS -ALLOWED_ORIGINS = ["http://localhost", "http://localhost:8080", "*"] -# Set web=True if you intend to serve a web interface, False otherwise -SERVE_WEB_INTERFACE = True - -# Call the function to get the FastAPI app instance -# Ensure the agent directory name ('capital_agent') matches your agent folder -app: FastAPI = get_fast_api_app( - agents_dir=AGENT_DIR, - session_service_uri=SESSION_SERVICE_URI, - artifact_service_uri=ARTIFACT_SERVICE_URI, - allow_origins=ALLOWED_ORIGINS, - web=SERVE_WEB_INTERFACE, -) - - -if __name__ == "__main__": - # Use the PORT environment variable provided by Cloud Run, defaulting to 8080 - uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080))) diff --git a/run-with-google-adk/cloudrun_deploy_run.sh b/run-with-google-adk/cloudrun_deploy_run.sh deleted file mode 100755 index 35bdbd70..00000000 --- a/run-with-google-adk/cloudrun_deploy_run.sh +++ /dev/null @@ -1,189 +0,0 @@ -# Copyright 2025 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. - -# this script runs from the top level directory (mcp-security when deploying and /app when running in container) - -#!/bin/bash - -ENV_FILE="./run-with-google-adk/google_mcp_security_agent/.env" - -# Function to create .env file -create_env_file() { - local env_file="$1" - shift # Remove the first argument ($env_file) - - # Check if the .env file already exists - if [ -f "$env_file" ]; then - echo "Warning: $env_file already exists. Overwriting." - fi - - # Create the .env file - touch "$env_file" - - # Add container variables to the .env file - required by UV based MCP servers. - for key in `env | cut -d "=" -f1`; do - value=$(eval "echo \$$key") # Expand the variable - echo "$key=$value" >> "$env_file" - done - echo "Successfully created $env_file with specified variables." -} - -# Check the command -if [ "$1" = "deploy" ]; then - echo "Starting deployment process..." - - # Check if the .env file exists - if [ ! -f "$ENV_FILE" ]; then - echo "Error: $ENV_FILE not found. Please create it with your environment variables." - exit 1 - fi - - env_vars="" - while read -r line || [[ -n "$line" ]]; do - trimmed_line=$(echo "$line" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') - - if [[ -z "$trimmed_line" ]]; then - continue - fi - - if [[ "$trimmed_line" =~ ^# ]]; then - continue - fi - - if [[ ! "$trimmed_line" =~ = ]]; then - continue - fi - - if [[ "$trimmed_line" =~ ^DEFAULT_PROMPT ]]; then - continue - fi - - key=`echo $trimmed_line | cut -d "=" -f1` - value=`echo $trimmed_line | cut -d "=" -f2` - - if [ "$key" = "GOOGLE_CLOUD_LOCATION" ]; then - GOOGLE_CLOUD_LOCATION="$value" - fi - if [ "$key" = "GOOGLE_CLOUD_PROJECT" ]; then - GOOGLE_CLOUD_PROJECT="$value" - fi - - if [[ $env_vars == "" ]]; then - env_vars="$key=$value" - else - env_vars="$env_vars,$key=$value" - fi - env_vars="$env_vars,$key=$value" - #echo "$line" - done < "$ENV_FILE" - - #echo $env_vars - - -# Handling default prompt separately - PYTHON_SCRIPT_PATH="/tmp/default_prompt.py" - -cat << EOF > "$PYTHON_SCRIPT_PATH" -import os -import dotenv -from dotenv import load_dotenv - -load_dotenv('$ENV_FILE') - -text = os.environ.get("DEFAULT_PROMPT") - -if text is None: - print("") -else: - prepared_text = text.replace('"', '\\\"') - prepared_text = prepared_text.replace('\\n', '\\\n') - prepared_text = prepared_text.replace(',', ';') - prepared_text = prepared_text.replace('-', '~') - print(prepared_text) -EOF - - default_prompt=$(python $PYTHON_SCRIPT_PATH) - env_vars="$env_vars,DEFAULT_PROMPT=$default_prompt,GCS_SA_JSON=object-viewer-sa1.json" - - # Check if any environment variables were found - if [ -z "$env_vars" ]; then - echo "Warning: No environment variables found in $ENV_FILE or all were skipped." - echo "The 'gcloud run deploy' command may fail if required variables are missing." - fi - - # Initialize the env_vars string with the given values - env_vars="$env_vars,REMOTE_RUN=Y" - - # Print the constructed environment variables string - echo "Using environment variables: $env_vars" - - # Copying files in the top level directory as required by cloud run deployment - echo "Temporarily copying files in the top level directory for image creation." - if [[ -e "./run-with-google-adk/object-viewer-sa.json" ]]; then - cp ./run-with-google-adk/object-viewer-sa.json object-viewer-sa1.json - fi - cp ./run-with-google-adk/cloudrun_deploy_run.sh . - cp ./run-with-google-adk/cloudrun_deploy.py . - cp ./run-with-google-adk/Dockerfile . - cp ./run-with-google-adk/.dockerignore . - - - - # Deploy the service with the dynamically constructed environment variables - gcloud run deploy mcp-security-agent-service \ - --source . \ - --region "$GOOGLE_CLOUD_LOCATION" \ - --project "$GOOGLE_CLOUD_PROJECT" \ - --allow-unauthenticated \ - --set-env-vars="$env_vars" \ - --memory 2Gi - - deploy_status=$? #get the status - - # Check the status of the deployment - if [ "$deploy_status" -eq 0 ]; then - # Deleting temporarily files in the top level directory - echo "Deleting temporarily copied files in the top level directory for image creation." - rm ./cloudrun_deploy_run.sh - rm ./cloudrun_deploy.py - rm ./Dockerfile - rm ./.dockerignore - if [[ -e "./run-with-google-adk/object-viewer-sa.json" ]]; then - rm object-viewer-sa1.json - fi - - echo "Successfully deployed the service." - else - rm ./cloudrun_deploy_run.sh - rm ./cloudrun_deploy.py - rm ./Dockerfile - rm ./.dockerignore - if [[ -e "./run-with-google-adk/object-viewer-sa.json" ]]; then - rm object-viewer-sa1.json - fi - echo "Failed to deploy the service." - #exit 1 - fi - -elif [ "$1" = "run" ]; then - echo "Creating .env file with specified variables..." - # Create the .env file with the variables from the container's environment - # These are required for the uv based MCPs to run. - create_env_file "/tmp/.env" - echo "Starting uvicorn server ..." - uvicorn cloudrun_deploy:app --host 0.0.0.0 --port $PORT -else - echo "Error: Invalid command. Use 'deploy' or 'run'." - #exit 1 -fi diff --git a/run-with-google-adk/google_mcp_security_agent/__init__.py b/run-with-google-adk/google_mcp_security_agent/__init__.py deleted file mode 100644 index 2153e9d4..00000000 --- a/run-with-google-adk/google_mcp_security_agent/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2025 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. -from . import agent \ No newline at end of file diff --git a/run-with-google-adk/google_mcp_security_agent/agent.py b/run-with-google-adk/google_mcp_security_agent/agent.py deleted file mode 100644 index 06782a6c..00000000 --- a/run-with-google-adk/google_mcp_security_agent/agent.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright 2025 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. -from google.adk.agents.llm_agent import LlmAgent -from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters, StdioConnectionParams -import os -import logging - -from utils_extensions_cbs_tools.extensions import MCPToolSetWithSchemaAccess -from utils_extensions_cbs_tools.tools import store_file, get_file_link, list_files -from utils_extensions_cbs_tools.callbacks import bmc_trim_llm_request, bac_setup_state_variable -from typing import TextIO -import sys - - -logging.basicConfig( - level=logging.INFO) - -if os.environ.get("MINIMAL_LOGGING","N") == "Y": - root_logger = logging.getLogger() - root_logger.setLevel(logging.ERROR) - - -def get_all_tools(): - """Get Tools from All MCP servers""" - logging.info("Attempting to connect to MCP servers...") - secops_tools = None - gti_tools = None - secops_soar_tools = None - scc_tools = None # Initialize scc_tools - - timeout = float(os.environ.get("STDIO_PARAM_TIMEOUT","60.0")) - - uv_dir_prefix="../server" - env_file_path = "../../../run-with-google-adk/google_mcp_security_agent/.env" - - if os.environ.get("REMOTE_RUN","N") == "Y": - env_file_path="/tmp/.env" - uv_dir_prefix="./server" - - if os.environ.get("AE_RUN","N") == "Y": - env_file_path="../../../google_mcp_security_agent/.env" - uv_dir_prefix="./server" - - logging.info(f"Using Env File Path - {env_file_path}, Current directory is - {os.getcwd()}, uv_dir_prefix is - {uv_dir_prefix}") - - # required temporarily for https://github.com/google/adk-python/issues/1024 - errlog_ae : TextIO = sys.stderr - if os.environ.get("AE_RUN","N") == "Y": - errlog_ae = None - - - if os.environ.get("LOAD_SCC_MCP") == "Y": - scc_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='uv', - args=[ "--directory", - uv_dir_prefix + "/scc", - "run", - "scc_mcp.py" - ] - ), - timeout=timeout), - tool_set_name="scc", - errlog=errlog_ae - ) - - if os.environ.get("LOAD_SECOPS_MCP") == "Y": - secops_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='uv', - args=[ "--directory", - uv_dir_prefix + "/secops/secops_mcp", - "run", - "--env-file", - env_file_path, - "server.py" - ] - ), - timeout=timeout), - tool_set_name="secops_mcp", - errlog=errlog_ae - ) - - if os.environ.get("LOAD_GTI_MCP") == "Y": - gti_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='uv', - args=[ "--directory", - uv_dir_prefix + "/gti/gti_mcp", - "run", - "--env-file", - env_file_path, - "server.py" - ] - ), - timeout=timeout), - tool_set_name="gti_mcp", - errlog=errlog_ae - ) - - - if os.environ.get("LOAD_SECOPS_SOAR_MCP") == "Y": - secops_soar_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='uv', - args=[ "--directory", - uv_dir_prefix + "/secops-soar/secops_soar_mcp", - "run", - "--env-file", - env_file_path, - "server.py", - "--integrations", - os.environ.get("SECOPS_INTEGRATIONS","CSV,OKTA") - ] - ), - timeout=timeout), - tool_set_name="secops_soar_mcp", - errlog=errlog_ae - ) - - logging.info("MCP Toolsets created successfully.") - return [secops_tools,gti_tools,secops_soar_tools,scc_tools] - -def create_agent(): - tools:any = [item for item in get_all_tools() if item is not None] - tools.append(store_file) - tools.append(get_file_link) - tools.append(list_files) - - agent = LlmAgent( - model=os.environ.get("GOOGLE_MODEL"), - name="google_mcp_security_agent", - instruction=os.environ.get("DEFAULT_PROMPT"), - tools=tools, - before_model_callback=bmc_trim_llm_request, - before_agent_callback=bac_setup_state_variable, -# sub_agents=[ADD SUB AGENTS HERE], - description="You are the google_mcp_security_agent." - - ) - return agent - - -root_agent = create_agent() diff --git a/run-with-google-adk/google_mcp_security_agent/sample.env.properties b/run-with-google-adk/google_mcp_security_agent/sample.env.properties deleted file mode 100644 index 492d7d22..00000000 --- a/run-with-google-adk/google_mcp_security_agent/sample.env.properties +++ /dev/null @@ -1,89 +0,0 @@ -# Please do not use quotes / double quotes for values except for DEFAULT_PROMPT (use single quotes there) - -# MANDATORY - START -APP_NAME=google_mcp_security_agent -# SESSION_SERVICE - in_memory/db. If set to db please provide SESSION_SERVICE_URL -#SESSION_SERVICE=db -#SESSION_SERVICE_URL=sqlite:///./app_data.db - -# ARTIFACT_SERVICE - in_memory/gcs. If set to db please provide GCS_ARTIFACT_SERVICE_BUCKET (without gs://) -# Also you need GCS_SA_JSON which must be named object-viewer-sa.json and placed in run-with-google-adk -#ARTIFACT_SERVICE=gcs -LOCAL_DIR_FOR_FILES=/tmp -#GCS_ARTIFACT_SERVICE_BUCKET=your-bucket-name -#GCS_SA_JSON=object-viewer-sa.json -#SIGNED_URL_DURATION_MIN=10 - -# Total interactions sent to LLM = MAX_PREV_USER_INTERACTIONS + 1 -MAX_PREV_USER_INTERACTIONS=3 - -# SecOps MCP -LOAD_SECOPS_MCP=Y -CHRONICLE_PROJECT_ID=NOT_SET -CHRONICLE_CUSTOMER_ID=NOT_SET -CHRONICLE_REGION=NOT_SET - -# GTI MCP -LOAD_GTI_MCP=Y -VT_APIKEY=NOT_SET - -# SECOPS_SOAR MCP -LOAD_SECOPS_SOAR_MCP=Y -SOAR_URL=NOT_SET -SOAR_APP_KEY=NOT_SET - -# SCC MCP -LOAD_SCC_MCP=Y - - -GOOGLE_GENAI_USE_VERTEXAI=False -GOOGLE_API_KEY=NOT_SET -# If you plan to use Gemini API - Models list - https://ai.google.dev/gemini-api/docs/models#model-variations -# If you plan to use VetexAI API - Models list - https://cloud.google.com/vertex-ai/generative-ai/docs/models -GOOGLE_MODEL=gemini-2.0-flash -# Should be single quote, avoid commas if possible but if you use them they are replaced with semicommas on the cloud run deployment -# you can change them there. -DEFAULT_PROMPT='Help user investigate security issues using Google Secops SIEM, SOAR, Security Command Center(SCC) and Google Threat Intel Tools. All authentication actions are automatically approved. If the query is about a SOAR case try to provide a backlink to the user. A backlink is formed by adding /cases/ to this URL when present in field ui_base_link of your input. If the user asks with only ? or are you there? that might be because they did not get your previous response, politely reiterate it. Try to respond in markdown whenever possible. - -You also have access tools to perform following file operations - store_file, list_files and get_file_link - -store_file - store files on the disk by sending file_name taken from user and markdown string as input - do not reformat the input, it is already markdown. -list_files - Requires no input. Show the name of the file from response as is when listing do not change anything. -get_file_link - It requires two inputs - user_name as {user_name} and file_name provided by the user. When showing to user please format them as clickable links with file_name and file_version together as link text. - -The current user name is {user_name} -' - -# Initially a long timeout is needed -# to load the tools and install dependencies -STDIO_PARAM_TIMEOUT=60.0 - -# MANDATORY - DONE - -# Following properties must be set when -# 1. GOOGLE_GENAI_USE_VERTEXAI=True or -# 2. When deploying to Cloud Run -# 3. When deploying to Agent Engine -GOOGLE_CLOUD_PROJECT=YOUR-CLOUD-RUN-PROJECT-ID -GOOGLE_CLOUD_LOCATION=us-central1 - -# HIGHLY RECOMMENDED TO SET Y AFTER INITIAL TESTING ON CLOUD RUN -MINIMAL_LOGGING=N - -# Agent Engine Deployment (without gs://) -#AE_STAGING_BUCKET=your-bucket-name -# If using custom ui, resource name from AE (projects//locations//reasoningEngines/) is needed -#AGENT_ENGINE_RESOURCE_NAME=YOUR_AE_RESOURCE_NAME - - - -# Add Your MCP server variables here, sample provided, please check the documentation -# MCP-1 -#LOAD_XDR_MCP=Y -#XDR_CLIENT_ID=abc123 -#XDR_CLIENT_SECRET=xyz456 -# MCP-2 -#LOAD_IDP_MCP=Y -#IDP_CLIENT_ID=abc123 -#IDP_CLIENT_SECRET=xyz456 - diff --git a/run-with-google-adk/main.py b/run-with-google-adk/main.py deleted file mode 100644 index 987fc97d..00000000 --- a/run-with-google-adk/main.py +++ /dev/null @@ -1,280 +0,0 @@ -# Copyright 2025 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. - - -# main.py -from fastapi import FastAPI, HTTPException -from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse -from fastapi.middleware.cors import CORSMiddleware -from fastapi.staticfiles import StaticFiles -import json -import os -from dotenv import load_dotenv -from google.genai import types -from google.adk.runners import Runner -from google.adk.sessions import InMemorySessionService, DatabaseSessionService -from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService -from google.adk.artifacts.gcs_artifact_service import GcsArtifactService -from contextlib import asynccontextmanager # Import for lifespan -from pydantic import BaseModel -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset - -# this makes sure that your prompts are logged. -# super useful for debugging -import logging -logging.basicConfig(level=logging.INFO) - -# Load environment variables from .env file in the parent directory -# Place this near the top, before using env vars like API keys -load_dotenv('./google_mcp_security_agent/.env') -from google_mcp_security_agent import agent - -app_name = os.environ.get("APP_NAME","ADK Agent") - -@asynccontextmanager -async def lifespan(app: FastAPI): - # Startup event: Initialize resources - print("Application starting up...") - # Example: You might initialize a global database connection pool here - # global_db_connection = await connect_to_db() - - # For our session_runner_map, we don't need to initialize it here - # as runners are created on demand. But if you had a fixed pool, this is the place. - - yield # The application will now start serving requests - - # Shutdown event: Clean up resources - print("Application shutting down...") - for session_id in session_runner_map: - print(f"Start-Cleaning up resources for Session[{session_id}]") - tools = session_runner_map[session_id].agent.tools - for mcp_toolset in tools: - # only need to close MCP toolsets and not function tools - print(f"Closing [{mcp_toolset}] of type - [{type(mcp_toolset)}]") - if isinstance(mcp_toolset,MCPToolset): - await mcp_toolset.close() - else: - print(f"skipping {mcp_toolset}") - print(f"Done-Cleaning up resources for Session[{session_id}]") - - # You could also delete the session if you wanted but generally not needed - # for inmemory session service and for db session service we do not want it to be deleted. - # also as such the MCPToolSet has a session shutdown method. which is more appropriate. - # so commenting out - # print(f"Start-Deleting session {session_id}") - # await session_service.delete_session(app_name='repair_world_app', user_id='customer',session_id=session_id) - # print(f"Done-Deleting session {session_id}") - - print("All session resources cleaned up.") - # Example: Close global database connection - # await global_db_connection.close() - -app = FastAPI(lifespan=lifespan) - -# Configure CORS to allow requests from the frontend (running on a different port/origin) -# Adjust origins as needed for your deployment environment -origins = [ - "http://localhost", - "http://localhost:8000", # FastAPI's default port - "http://localhost:5500", # Common for Live Server in VS Code - "http://127.0.0.1:5500", - "http://127.0.0.1:8000", - "*" -] - -app.add_middleware( - CORSMiddleware, - allow_origins=origins, - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Mount static files to serve index.html and app.js -# Ensure 'static' directory exists in the same location as main.py -current_dir = os.path.dirname(os.path.abspath(__file__)) -static_dir = os.path.join(current_dir, "static") - -# TODO Vertex AI session service -print(f"Session Service Type - {os.environ.get("SESSION_SERVICE","in_memory")}") -session_service = InMemorySessionService() -if os.environ.get("SESSION_SERVICE","in_memory") == "db": - db_url = os.environ.get("SESSION_SERVICE_URL","sqlite:///./agent_data.db") - session_service = DatabaseSessionService(db_url=db_url) - -# Artifact service might not be needed for this example -print(f"Artifact Service Type - {os.environ.get("ARTIFACT_SERVICE","in_memory")}") -artifacts_service = InMemoryArtifactService() -if os.environ.get("ARTIFACT_SERVICE","gcs") == "gcs": - artifacts_service = GcsArtifactService(bucket_name=os.environ.get("GCS_ARTIFACT_SERVICE_BUCKET")) - -# Create 'static' directory if it doesn't exist -os.makedirs(static_dir, exist_ok=True) - -app.mount("/static", StaticFiles(directory=static_dir), name="static") - -# if os.environ.get("GOOGLE_API_KEY") == "NOT_SET": -# print("Please set a Google API Key using - https://aistudio.google.com/app/apikey") -# exit(1) - -session_runner_map={} - - -@app.get("/", response_class=HTMLResponse) -async def read_root(): - """Serves the index.html (login page) when the root URL is accessed.""" - index_html_path = os.path.join(static_dir, "index.html") - if not os.path.exists(index_html_path): - raise HTTPException(status_code=404, detail="index.html not found in static directory.") - with open(index_html_path, "r") as f: - return HTMLResponse(content=f.read()) - -@app.get("/landing.html", response_class=HTMLResponse) -async def read_landing(): - """Serves the landing.html (chat page).""" - landing_html_path = os.path.join(static_dir, "landing.html") - if not os.path.exists(landing_html_path): - raise HTTPException(status_code=404, detail="landing.html not found in static directory.") - with open(landing_html_path, "r") as f: - return HTMLResponse(content=f.read()) - -@app.get("/app_name") -async def get_app_name(): - """Returns the application name.""" - return JSONResponse(content={"app_name": app_name}) - -async def create_new_session(username): - user_id = username # For simplicity, user_id is the username - initial_state = { - "user_name":f"{user_id}" - } - session = await session_service.create_session( - state=initial_state, app_name=app_name, user_id=user_id - ) - logging.info(f"Created session {session.id} for (app_name={app_name}, user_id={user_id}) with user - {session.state['user_name']}") - return session - -@app.get("/get_session") -async def get_session_and_user_id(username: str,start_new_session: str="N"): - """ - Generates and Or get the first session - """ - # in case username is not sent, use default_user - user_id = "default_user" if username == "None" else username - - if start_new_session == "Y": - logging.info(f"Fresh session requested for {username}") - session = await create_new_session(username) - else: - list_session_response = await session_service.list_sessions(app_name=app_name, user_id=user_id) - if len(list_session_response.sessions) > 0: - session = list_session_response.sessions[0] - logging.info(f"Retrieved session - {session.id}") - else: - session = await create_new_session(username) - - root_agent = agent.root_agent - - runner = Runner( - app_name=app_name, - agent=root_agent, - artifact_service=artifacts_service, # Optional - session_service=session_service, - ) - # TODO - Check if we could use a new running without any performance degradation - session_runner_map[session.id] = runner - return {"session_id": session.id, "user_id": user_id} - -def enrich_output(event): - type = "" - author = event.author - message = "" - - if event.content and event.content.parts: - author = author+ "-" + event.content.role - message = event.content.parts[0].text - if event.get_function_calls(): - type="TCR"#"Tool Call Request" - message = event.get_function_calls()[0].name - #print(" Type: Tool Call Request") - elif event.get_function_responses(): - type="TR"#"Tool Result" - # in this case message is basically function name (same with TCR as well) - message = event.get_function_responses()[0].name - #print(" Type: Tool Result") - elif event.content.parts[0].text: - if event.partial: - type="STC"#"Streaming Text Chunk" - #print(" Type: Streaming Text Chunk") - else: - type="CTC" #"Complete Text Message" - # print(" Type: Complete Text Message") - #print(event.content.parts[0].text) - else: - type="OC"#"Other Content" - #print(" Type: Other Content (e.g., code result)") - elif event.actions and (event.actions.state_delta or event.actions.artifact_delta): - type="S/A U"#"State/Artifact Update" - #print(" Type: State/Artifact Update") - else: - type="CS/O"#"Control Signal or Other" - #print(" Type: Control Signal or Other") - - return(type,author,message) - - -# Define the structure for the POST request body -class ChatRequest(BaseModel): - session_id: str - user_id: str - message: str - -@app.post("/chat") -async def chat_endpoint(request_body: ChatRequest): - - session_id = request_body.session_id - user_id = request_body.user_id - message = request_body.message - - content = types.Content(role='user', parts=[types.Part(text=message)]) - runner = session_runner_map[session_id] - # todo get user from session - events_async = runner.run_async( - session_id=session_id, user_id=user_id, new_message=content - ) - - async def event_generator(): - async for event in events_async: - type,author,message = enrich_output(event) - # if event.content and event.content.parts: - data = { - "text": f"{type}({author}) : \n\n {message}", - "last_msg": False - } - yield f"data: {json.dumps(data)}\n\n" - - # Once the async for loop finishes, it means events_async has been exhausted. - # Now, send the final "last_msg" signal. - final_data = { - "text": "Stream finished.", # You can customize this final message or make it empty - "last_msg": True - } - yield f"data: {json.dumps(final_data)}\n\n" - - return StreamingResponse(event_generator(), media_type="text/event-stream") - -# To run this application: -# 1. Install dependencies: pip install -r requirements.txt -# 2. Run the server: uvicorn main:app --reload -# 3. Open your browser to http://localhost:8000/ diff --git a/run-with-google-adk/main_ae.py b/run-with-google-adk/main_ae.py deleted file mode 100644 index b1559c3a..00000000 --- a/run-with-google-adk/main_ae.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright 2025 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. - -# main.py -from fastapi import FastAPI, HTTPException -from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse -from fastapi.middleware.cors import CORSMiddleware -from fastapi.staticfiles import StaticFiles -import json -import os -from dotenv import load_dotenv -from pydantic import BaseModel -from vertexai import agent_engines -# this makes sure that your prompts are logged. -# super useful for debugging -import logging -logging.basicConfig(level=logging.INFO) - -# Load environment variables from .env file in the parent directory -# Place this near the top, before using env vars like API keys -load_dotenv('./google_mcp_security_agent/.env') - -app_name = os.environ.get("APP_NAME","ADK Agent") - -app = FastAPI() - -# Configure CORS to allow requests from the frontend (running on a different port/origin) -# Adjust origins as needed for your deployment environment -origins = [ - "http://localhost", - "http://localhost:8000", # FastAPI's default port - "http://localhost:5500", # Common for Live Server in VS Code - "http://127.0.0.1:5500", - "http://127.0.0.1:8000", - "*" -] - -app.add_middleware( - CORSMiddleware, - allow_origins=origins, - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Mount static files to serve index.html and app.js -# Ensure 'static' directory exists in the same location as main.py -current_dir = os.path.dirname(os.path.abspath(__file__)) -static_dir = os.path.join(current_dir, "static") - -# Create 'static' directory if it doesn't exist -os.makedirs(static_dir, exist_ok=True) - -app.mount("/static", StaticFiles(directory=static_dir), name="static") - -# if os.environ.get("GOOGLE_API_KEY") == "NOT_SET": -# print("Please set a Google API Key using - https://aistudio.google.com/app/apikey") -# exit(1) - -@app.get("/", response_class=HTMLResponse) -async def read_root(): - """Serves the index.html (login page) when the root URL is accessed.""" - index_html_path = os.path.join(static_dir, "index.html") - if not os.path.exists(index_html_path): - raise HTTPException(status_code=404, detail="index.html not found in static directory.") - with open(index_html_path, "r") as f: - return HTMLResponse(content=f.read()) - -@app.get("/landing.html", response_class=HTMLResponse) -async def read_landing(): - """Serves the landing.html (chat page).""" - landing_html_path = os.path.join(static_dir, "landing.html") - if not os.path.exists(landing_html_path): - raise HTTPException(status_code=404, detail="landing.html not found in static directory.") - with open(landing_html_path, "r") as f: - return HTMLResponse(content=f.read()) - -@app.get("/app_name") -async def get_app_name(): - """Returns the application name.""" - return JSONResponse(content={"app_name": app_name}) - - -@app.get("/get_session") -async def get_session_and_user_id(username: str,start_new_session: str="N"): - """ - Generates and Or get the first session - """ - # in case username is not sent, use default_user - user_id = "default_user" if username == "None" else username - agent_resource=os.environ.get("AGENT_ENGINE_RESOURCE_NAME") - print(f"Agent Resource - {agent_resource}") - remote_app = agent_engines.get(agent_resource) - print(f"Current User Id - {user_id}") - remote_sessions = remote_app.list_sessions(user_id=user_id) - print(f"remote_sessions - {remote_sessions}") - session_available = False - - if start_new_session != "Y": - if "sessions" in remote_sessions and len(remote_sessions["sessions"]) > 0: - print(f"Fetched {len(remote_sessions)} sessions back",remote_sessions) - print("using sesion 0") - remote_session = remote_sessions["sessions"][0] - session_available = True - - if not session_available: - print("No sessions fetched for the given user and engine combination / or new requested - creating one") - remote_session = remote_app.create_session(user_id=user_id) - - print(f"Using remote session -> {remote_session["id"]}") - - return {"session_id": remote_session["id"], "user_id": user_id} - -def enrich_output(event): - #print(event) - msg_type = "" - author = event["author"] - message = "" - - if "content" in event and "parts" in event["content"]: - author = author+ "-" + event["content"]["role"] - if "text" in event["content"]["parts"][0]: - message = event["content"]["parts"][0]["text"] - if "partial" in event: - msg_type="STC"#"Streaming Text Chunk" - #print(" Type: Streaming Text Chunk") - else: - msg_type="CTC" #"Complete Text Message" - elif "function_call" in event["content"]["parts"][0]: - msg_type="TCR"#"Tool Call Request" - message = event["content"]["parts"][0]["function_call"]["name"] - #print(" Type: Tool Call Request") - elif "function_response" in event["content"]["parts"][0]: - msg_type="TR"#"Tool Result" - # in this case message is basically function name (same with TCR as well) - message = event["content"]["parts"][0]["function_response"]["name"] - #print(" Type: Tool Result") - else: - msg_type="OC"#"Other Content" - #print(" Type: Other Content (e.g., code result)") - elif "actions" in event and ("state_delta" in event["actions"] or "artifact_delta" in event["actions"]): - msg_type="S/A U"#"State/Artifact Update" - #print(" Type: State/Artifact Update") - else: - msg_type="CS/O"#"Control Signal or Other" - #print(" Type: Control Signal or Other") - - return(msg_type,author,message) - - -# Define the structure for the POST request body -class ChatRequest(BaseModel): - session_id: str - user_id: str - message: str - -@app.post("/chat") -async def chat_endpoint(request_body: ChatRequest): - - session_id = request_body.session_id - user_id = request_body.user_id - message = request_body.message - - remote_app = agent_engines.get(os.environ.get("AGENT_ENGINE_RESOURCE_NAME")) - - async def event_generator(): - - for event in remote_app.stream_query( - user_id=user_id, - session_id=session_id, - message=request_body.message, - ): - #print(event) - #print(type(event)) # dict - msg_type,author,message = enrich_output(event) - data = { - "text": f"{msg_type}({author}) : \n\n {message}", - "last_msg": False - } - yield f"data: {json.dumps(data)}\n\n" - - # Once the async for loop finishes, it means events_async has been exhausted. - # Now, send the final "last_msg" signal. - final_data = { - "text": "Stream finished.", # You can customize this final message or make it empty - "last_msg": True - } - yield f"data: {json.dumps(final_data)}\n\n" - - return StreamingResponse(event_generator(), media_type="text/event-stream") - -# To run this application: -# 1. Install dependencies: pip install -r requirements.txt -# 2. Run the server: uvicorn main_ae:app --reload -# 3. Open your browser to http://localhost:8000/ diff --git a/run-with-google-adk/pyproject.toml b/run-with-google-adk/pyproject.toml new file mode 100644 index 00000000..d89c9254 --- /dev/null +++ b/run-with-google-adk/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "mcp-security-agent" +version = "0.2.0" +description = "Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2 and MCP" +readme = "README.md" +requires-python = ">=3.11" +authors = [ + { name = "Google LLC" } +] +dependencies = [ + "google-adk>=2.0.0", + "google-genai>=1.20.0", + "google-cloud-aiplatform>=1.97.0", + "pydantic>=2.0.0", + "pydantic-settings>=2.0.0", + "mcp>=1.0.0,<2.0.0", + "fastapi>=0.115.0", + "uvicorn>=0.30.0", + "python-dotenv>=1.0.0", + "rich>=13.0.0", + "typer>=0.12.0", +] + +[project.optional-dependencies] +test = [ + "pytest>=8.0.0", + "pytest-asyncio>=0.24.0", + "httpx>=0.27.0", +] + +[project.scripts] +mcp-security-agent = "mcp_security_agent.cli:app" + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/run-with-google-adk/requirements.txt b/run-with-google-adk/requirements.txt deleted file mode 100644 index 1503a5bb..00000000 --- a/run-with-google-adk/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -google-cloud-aiplatform==1.97.0 -markdown -uv -google-adk[eval]==1.3.0 -google-genai==1.20.0 -pandas diff --git a/run-with-google-adk/run-adk-agent.sh b/run-with-google-adk/run-adk-agent.sh deleted file mode 100755 index 89541668..00000000 --- a/run-with-google-adk/run-adk-agent.sh +++ /dev/null @@ -1,155 +0,0 @@ -# Copyright 2025 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. - -#!/bin/bash - -# Define the file paths -SAMPLE_ENV_FILE="./google_mcp_security_agent/sample.env.properties" -ENV_FILE="./google_mcp_security_agent/.env" - -# Function to mask environment variable values -mask_env_value() { - local value="$1" - local length="${#value}" - local max_length=30 # Define the maximum length - - if [ "$length" -gt 3 ]; then - masked_value=$(printf "%.3s%${length}s" "$value" "" | sed "s/ /#/g") - if [ ${#masked_value} -gt $max_length ]; then - echo "${masked_value:0:$max_length}" # truncate to max length - else - echo "$masked_value" - fi - else - echo "$value" - fi -} - -# Function to display the contents of the .env file with masked values -show_env_masked() { - if [ -f "$ENV_FILE" ]; then - echo "Contents of .env (with masked values):" - while IFS='=' read -r key value; do - if [[ "$key" != "#"* ]]; then # Ignore comments - if [ -n "$key" ]; then #check if key is not empty - masked_value=$(mask_env_value "$value") - echo "$key=$masked_value" - fi - else - echo "$key$value" # print comments as is - fi - done < "$ENV_FILE" - else - echo ".env file does not exist." - fi -} - - -#!/bin/bash - -# Define a function to display usage instructions -usage() { - echo "Usage: $0 [args...]" - echo "" - echo "Commands:" - echo " adk_web : Runs ADK Web for local agent." - echo " adk_web : Runs ADK Web with a session service URI." - echo " adk_web : Runs ADK Web with session and artifact service URIs." - echo " custom_ui : Runs Custom UI for local agent (uvicorn main:app --reload)." - echo " custom_ui_ae : Runs Custom UI for Agent Engine Backend (uvicorn main_ae:app --reload)." - echo "" - echo "Examples:" - echo " $0 adk_web" - echo " $0 adk_web http://localhost:8000" - echo " $0 adk_web http://localhost:8000 http://localhost:8001" - echo " $0 custom_ui" - echo " $0 custom_ui_ae" - echo " $0 # (will show usage)" - exit 1 -} - -# Get the first argument, which is the command. -# If no argument is provided, COMMAND will be an empty string. -COMMAND="$1" - -# If no command is provided (i.e., zero arguments), or if the provided command is unknown, -# display the usage instructions. -if [ -z "$COMMAND" ]; then - echo "No command provided, Checking environment file status..." - COMMAND="env_files" -fi - -# Use a case statement to handle different commands -case "$COMMAND" in - adk_web) - # If .env exists, display its contents with masked values and run the command - show_env_masked - # Handle adk_web command based on argument count - if [ "$#" -eq 1 ]; then - echo "Running ADK Web for local agent..." - adk web - elif [ "$#" -eq 2 ]; then - echo "Running ADK Web with session service URI: $2" - adk web --session_service_uri "$2" - elif [ "$#" -eq 3 ]; then - echo "Running ADK Web with session service URI: $2 and artifact service URI: $3" - adk web --session_service_uri "$2" --artifact_service_uri "$3" - else - echo "Error: Incorrect number of arguments for 'adk_web'." - usage - fi - ;; - custom_ui) - show_env_masked - # Ensure that 'custom_ui' and 'custom_ui_ae' are only called with 1 argument. - if [ "$#" -ne 1 ]; then - echo "Error: 'custom_ui' expects no additional arguments." - usage - fi - echo "Running Custom UI for local agent ..." - uvicorn main:app --reload - ;; - custom_ui_ae) - show_env_masked - # Ensure that 'custom_ui' and 'custom_ui_ae' are only called with 1 argument. - if [ "$#" -ne 1 ]; then - echo "Error: 'custom_ui_ae' expects no additional arguments." - usage - fi - echo "Running Custom UI for Agent Engine Backend ..." - uvicorn main_ae:app --reload - ;; - env_files) - # Check for the existence of the files - if [ ! -f "$SAMPLE_ENV_FILE" ] && [ ! -f "$ENV_FILE" ]; then - echo "Error: Missing both $SAMPLE_ENV_FILE and $ENV_FILE files." - exit 1 - elif [ -f "$SAMPLE_ENV_FILE" ] && [ ! -f "$ENV_FILE" ]; then - echo "Copying $SAMPLE_ENV_FILE to $ENV_FILE..." - cp "$SAMPLE_ENV_FILE" "$ENV_FILE" - echo "Please update the environment variables in $ENV_FILE" - exit 0 - else - echo "Environment file ok, please check the usage below" - usage - fi - ;; - *) - # Default case for unknown commands when an argument *was* provided. - echo "Error: Unknown command '$COMMAND'." - usage - ;; -esac - - diff --git a/run-with-google-adk/sample.env b/run-with-google-adk/sample.env new file mode 100644 index 00000000..58b910d0 --- /dev/null +++ b/run-with-google-adk/sample.env @@ -0,0 +1,33 @@ +# Google Cloud & LLM Settings +GOOGLE_CLOUD_PROJECT=your-gcp-project-id +GOOGLE_CLOUD_LOCATION=us-central1 +GOOGLE_GENAI_USE_VERTEXAI=False +GOOGLE_API_KEY=your-gemini-api-key +GOOGLE_MODEL=gemini-2.5-flash + +# MCP Server Enablement Flags (Y/N or True/False) +LOAD_SECOPS_MCP=Y +LOAD_SCC_MCP=Y +LOAD_GTI_MCP=Y +LOAD_SECOPS_SOAR_MCP=N + +# Credentials & Service Account Impersonation +SECOPS_SA_PATH= +GOOGLE_APPLICATION_CREDENTIALS= +SECOPS_IMPERSONATE_SERVICE_ACCOUNT= + +# Google SecOps (Chronicle SIEM) Settings +CHRONICLE_PROJECT_ID=your-chronicle-project-id +CHRONICLE_CUSTOMER_ID=your-chronicle-customer-id +CHRONICLE_REGION=us + +# Google Threat Intelligence (GTI / VirusTotal) +VT_APIKEY=your-virustotal-api-key + +# SecOps SOAR Settings +SOAR_URL=https://your-soar-tenant.siemplify-soar.com +SOAR_APP_KEY=your-soar-app-key + +# Runtime Settings +STDIO_PARAM_TIMEOUT=60.0 +MINIMAL_LOGGING=N diff --git a/run-with-google-adk/sample_servers_to_integrate/agents/demo_idp_agent.py b/run-with-google-adk/sample_servers_to_integrate/agents/demo_idp_agent.py index d38ce94a..ad61dccc 100644 --- a/run-with-google-adk/sample_servers_to_integrate/agents/demo_idp_agent.py +++ b/run-with-google-adk/sample_servers_to_integrate/agents/demo_idp_agent.py @@ -11,81 +11,77 @@ # 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. +"""Demo Identity Provider (IDP) sub-agent integration example for ADK v2.""" -from google.adk.agents.llm_agent import LlmAgent -from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters, StdioConnectionParams import os import logging +from pathlib import Path +from typing import Optional, Any +from mcp_security_agent.callbacks import bmc_trim_llm_request -from utils_extensions_cbs_tools.extensions import MCPToolSetWithSchemaAccess -from utils_extensions_cbs_tools.callbacks import bmc_trim_llm_request -from typing import TextIO -import sys +logger = logging.getLogger(__name__) -# TODO improve logging -logging.basicConfig( - level=logging.INFO) -if os.environ.get("MINIMAL_LOGGING","N") == "Y": - root_logger = logging.getLogger() - root_logger.setLevel(logging.ERROR) +def create_demo_idp_agent( + model: Optional[str] = None, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, +) -> Any: + """Initializes and returns the demo IDP sub-agent. + Args: + model: Model name (defaults to GOOGLE_MODEL env var or gemini-2.5-flash). + client_id: IDP client ID (defaults to IDP_CLIENT_ID env var). + client_secret: IDP client secret (defaults to IDP_CLIENT_SECRET env var). -is_any_variable_unset = any( - os.getenv(var, "NOT_SET") == "NOT_SET" - for var in ["IDP_CLIENT_ID","IDP_CLIENT_SECRET",] -) -if is_any_variable_unset: - print(f"Please set all required environment variables, check .env file") - exit(1) + Returns: + Configured LlmAgent instance or None if dependencies are missing. + """ + model_name = model or os.getenv("GOOGLE_MODEL", "gemini-2.5-flash") + cid = client_id or os.getenv("IDP_CLIENT_ID", "demo-client-id") + csec = client_secret or os.getenv("IDP_CLIENT_SECRET", "demo-client-secret") -def get_all_tools(): - """Get Tools from IDP MCP""" - logging.info("Attempting to connect to MCP servers for IDP...") - idp_tools = None # Initialize scc_tools - - uv_dir_prefix="../server" - if os.environ.get("REMOTE_RUN","N") == "Y": - uv_dir_prefix="./server" + try: + from google.adk.agents.llm_agent import LlmAgent + from google.adk.tools.mcp_tool.mcp_toolset import ( + McpToolset, + StdioConnectionParams, + StdioServerParameters, + ) + except ImportError: + logger.warning("google.adk not available; skipping demo IDP agent creation.") + return None - if os.environ.get("AE_RUN","N") == "Y": - uv_dir_prefix="./server" + # Resolve path to sample IDP MCP server + sample_dir = Path(__file__).resolve().parents[1] + idp_server_path = sample_dir / "mcp_servers" / "demo_idp" / "idp_mcp_server.py" - # required temporarily for https://github.com/google/adk-python/issues/1024 - errlog_ae : TextIO = sys.stderr - if os.environ.get("AE_RUN","N") == "Y": - errlog_ae = None + timeout = float(os.getenv("STDIO_PARAM_TIMEOUT", "60.0")) - timeout = float(os.environ.get("STDIO_PARAM_TIMEOUT","60.0")) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="python", + args=[ + str(idp_server_path), + "--client-id", + cid, + "--client-secret", + csec, + ], + ), + timeout=timeout, + ) + tools = [McpToolset(connection_params=conn)] - if os.environ.get("LOAD_IDP_MCP") == "Y": - - idp_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='python', - args=[ f"{uv_dir_prefix}/demo_idp/idp_mcp_server.py", - "--client-id", - os.environ.get("IDP_CLIENT_ID"), - "--client-secret", - os.environ.get("IDP_CLIENT_SECRET") - ] - ), - timeout=timeout), - tool_set_name="demo_idp_tools", - errlog=errlog_ae - ) - - logging.info("MCP Toolsets for IDP created successfully.") - return [idp_tools] - -tools:any = [item for item in get_all_tools() if item is not None] - -demo_idp_agent = LlmAgent( - model=os.environ.get("GOOGLE_MODEL"), - name="demo_idp_agent", - instruction="You help users to gather information about their users/identities from their IDP backend to investigate. Do take calculated guesses based on your own knowledge base to help the user as much as you can, even if you may not know much about the IDP product itself. At the end of a query when there is some data do let them know your opinion and the steps you carried to achieve the output.", - tools=tools, - before_model_callback=bmc_trim_llm_request, - description="You are the demo_idp_agent. Anything not related to IDP please delegate to google_mcp_security_agent" -) + agent = LlmAgent( + model=model_name, + name="demo_idp_agent", + instruction=( + "You help users gather identity information from the IDP backend during investigations. " + "Formulate search queries and analyze user account statuses." + ), + tools=tools, + before_model_callback=bmc_trim_llm_request, + description="Demo IDP agent for identity lookup and authentication troubleshooting.", + ) + return agent diff --git a/run-with-google-adk/sample_servers_to_integrate/agents/demo_xdr_agent.py b/run-with-google-adk/sample_servers_to_integrate/agents/demo_xdr_agent.py index 15053619..20f0d982 100644 --- a/run-with-google-adk/sample_servers_to_integrate/agents/demo_xdr_agent.py +++ b/run-with-google-adk/sample_servers_to_integrate/agents/demo_xdr_agent.py @@ -11,82 +11,77 @@ # 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. +"""Demo Extended Detection and Response (XDR) sub-agent integration example for ADK v2.""" -from google.adk.agents.llm_agent import LlmAgent -from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters, StdioConnectionParams import os import logging - -from utils_extensions_cbs_tools.extensions import MCPToolSetWithSchemaAccess -from utils_extensions_cbs_tools.callbacks import bmc_trim_llm_request - -from typing import TextIO -import sys - -# TODO improve logging -logging.basicConfig( - level=logging.INFO) - -if os.environ.get("MINIMAL_LOGGING","N") == "Y": - root_logger = logging.getLogger() - root_logger.setLevel(logging.ERROR) - - -is_any_variable_unset = any( - os.getenv(var, "NOT_SET") == "NOT_SET" - for var in ["XDR_CLIENT_ID","XDR_CLIENT_SECRET",] -) -if is_any_variable_unset: - print(f"Please set all required environment variables, check .env file") - exit(1) - -def get_all_tools(): - """Get Tools from XDR MCP""" - logging.info("Attempting to connect to MCP servers for XDR...") - xdr_tools = None # Initialize scc_tools - - uv_dir_prefix="../server" - if os.environ.get("REMOTE_RUN","N") == "Y": - uv_dir_prefix="./server" - - if os.environ.get("AE_RUN","N") == "Y": - uv_dir_prefix="./server" - - # required temporarily for https://github.com/google/adk-python/issues/1024 - errlog_ae : TextIO = sys.stderr - if os.environ.get("AE_RUN","N") == "Y": - errlog_ae = None - - timeout = float(os.environ.get("STDIO_PARAM_TIMEOUT","60.0")) - - if os.environ.get("LOAD_XDR_MCP") == "Y": - - xdr_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='python', - args=[ f"{uv_dir_prefix}/demo_xdr/xdr_mcp_server.py", - "--client-id", - os.environ.get("XDR_CLIENT_ID"), - "--client-secret", - os.environ.get("XDR_CLIENT_SECRET") - ] - ), - timeout=timeout), - tool_set_name="demo_xdr_tools", - errlog=errlog_ae - ) - - logging.info("MCP Toolsets for XDR created successfully.") - return [xdr_tools] - -tools:any = [item for item in get_all_tools() if item is not None] - -demo_xdr_agent = LlmAgent( - model=os.environ.get("GOOGLE_MODEL"), - name="demo_xdr_agent", - instruction="You help users to gather information about their hosts from their XDR backend to investigate. Do take calculated guesses based on your own knowledge base to help the user as much as you can, even if you may not know much about the XDR product itself. At the end of a query when there is some data do let them know your opinion and the steps you carried to achieve the output.", - tools=tools, - before_model_callback=bmc_trim_llm_request, - description="You are the demo_xdr_agent. Anything not related to XDR please delegate to google_mcp_security_agent" -) +from pathlib import Path +from typing import Optional, Any +from mcp_security_agent.callbacks import bmc_trim_llm_request + +logger = logging.getLogger(__name__) + + +def create_demo_xdr_agent( + model: Optional[str] = None, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, +) -> Any: + """Initializes and returns the demo XDR sub-agent. + + Args: + model: Model name (defaults to GOOGLE_MODEL env var or gemini-2.5-flash). + client_id: XDR client ID (defaults to XDR_CLIENT_ID env var). + client_secret: XDR client secret (defaults to XDR_CLIENT_SECRET env var). + + Returns: + Configured LlmAgent instance or None if dependencies are missing. + """ + model_name = model or os.getenv("GOOGLE_MODEL", "gemini-2.5-flash") + cid = client_id or os.getenv("XDR_CLIENT_ID", "demo-client-id") + csec = client_secret or os.getenv("XDR_CLIENT_SECRET", "demo-client-secret") + + try: + from google.adk.agents.llm_agent import LlmAgent + from google.adk.tools.mcp_tool.mcp_toolset import ( + McpToolset, + StdioConnectionParams, + StdioServerParameters, + ) + except ImportError: + logger.warning("google.adk not available; skipping demo XDR agent creation.") + return None + + # Resolve path to sample XDR MCP server + sample_dir = Path(__file__).resolve().parents[1] + xdr_server_path = sample_dir / "mcp_servers" / "demo_xdr" / "xdr_mcp_server.py" + + timeout = float(os.getenv("STDIO_PARAM_TIMEOUT", "60.0")) + + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="python", + args=[ + str(xdr_server_path), + "--client-id", + cid, + "--client-secret", + csec, + ], + ), + timeout=timeout, + ) + tools = [McpToolset(connection_params=conn)] + + agent = LlmAgent( + model=model_name, + name="demo_xdr_agent", + instruction=( + "You help users gather endpoint and host telemetry from the XDR backend during investigations. " + "Formulate search queries and analyze process/host activity." + ), + tools=tools, + before_model_callback=bmc_trim_llm_request, + description="Demo XDR agent for host telemetry and endpoint threat analysis.", + ) + return agent diff --git a/run-with-google-adk/utils_extensions_cbs_tools/cache.py b/run-with-google-adk/src/mcp_security_agent/__init__.py similarity index 73% rename from run-with-google-adk/utils_extensions_cbs_tools/cache.py rename to run-with-google-adk/src/mcp_security_agent/__init__.py index 1ec4b182..d04744c2 100644 --- a/run-with-google-adk/utils_extensions_cbs_tools/cache.py +++ b/run-with-google-adk/src/mcp_security_agent/__init__.py @@ -11,5 +11,9 @@ # 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. +"""MCP Security Agent powered by Google ADK v2.""" -tools_cache={} \ No newline at end of file +from mcp_security_agent.agent import create_security_agent, root_agent + +__version__ = "0.2.0" +__all__ = ["create_security_agent", "root_agent", "__version__"] diff --git a/run-with-google-adk/utils_extensions_cbs_tools/utils.py b/run-with-google-adk/src/mcp_security_agent/__main__.py similarity index 80% rename from run-with-google-adk/utils_extensions_cbs_tools/utils.py rename to run-with-google-adk/src/mcp_security_agent/__main__.py index 7f1798b8..b4851ec1 100644 --- a/run-with-google-adk/utils_extensions_cbs_tools/utils.py +++ b/run-with-google-adk/src/mcp_security_agent/__main__.py @@ -11,5 +11,9 @@ # 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. +"""Executable entry point for python -m mcp_security_agent.""" +from mcp_security_agent.cli import app +if __name__ == "__main__": + app() diff --git a/run-with-google-adk/src/mcp_security_agent/agent.py b/run-with-google-adk/src/mcp_security_agent/agent.py new file mode 100644 index 00000000..3685261e --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/agent.py @@ -0,0 +1,65 @@ +# Copyright 2025 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. +"""ADK v2.x Agent definition and factory for MCP Security Agent.""" + +import logging +from typing import Optional, Any +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.toolsets import build_mcp_toolsets +from mcp_security_agent.callbacks import bmc_trim_llm_request + +logger = logging.getLogger(__name__) + +SOC_AGENT_SYSTEM_PROMPT = """You are an expert Autonomous Security Operations Center (SOC) Analyst and Threat Intelligence Assistant. +Your mission is to investigate security alerts, hunt for threats in UDM logs, analyze IoCs with Google Threat Intelligence, triage Cloud Security Command Center (SCC) findings, and execute SOAR remediation playbooks. + +Guidelines: +1. Always ground your investigations in factual telemetry retrieved from MCP tools. +2. Formulate clear UDM queries, correlate suspicious IP/domain/hash artifacts, and provide actionable remediation steps. +3. Structure your analysis with clear headings: Executive Summary, Investigation Findings, Artifact Analysis, and Recommended Remediation. +""" + + +def create_security_agent(settings: Optional[AgentSettings] = None) -> Any: + """Initializes and returns the configured SOC Security Agent. + + Args: + settings: Optional AgentSettings instance (defaults to loading from environment). + + Returns: + Configured LlmAgent instance. + """ + if settings is None: + settings = AgentSettings() + + toolsets = build_mcp_toolsets(settings) + + try: + from google.adk.agents.llm_agent import LlmAgent + except ImportError: + logger.warning("google.adk.agents.llm_agent not available; returning dummy agent.") + return None + + agent = LlmAgent( + name="SecurityOperationsAgent", + model=settings.google_model, + instruction=settings.default_prompt or SOC_AGENT_SYSTEM_PROMPT, + tools=toolsets, + before_model_callback=bmc_trim_llm_request, + ) + return agent + + +# Expose root_agent for standard ADK CLI discovery (adk run, adk web) +root_agent = create_security_agent() diff --git a/run-with-google-adk/src/mcp_security_agent/callbacks.py b/run-with-google-adk/src/mcp_security_agent/callbacks.py new file mode 100644 index 00000000..c7559477 --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/callbacks.py @@ -0,0 +1,36 @@ +# Copyright 2025 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. +"""Lifecycle callbacks for request trimming and context logging in Google ADK.""" + +import logging +from typing import Any, Optional + +logger = logging.getLogger(__name__) + + +def bmc_trim_llm_request(callback_context: Any, llm_request: Any) -> Optional[Any]: + """Callback executed prior to LLM invocation to inspect and trim context if necessary. + + In Google ADK v2, returning None allows the standard model execution flow to proceed. + Returning an LlmResponse will short-circuit and provide an immediate response. + + Args: + callback_context: ADK callback context object. + llm_request: Inbound LLM request object. + + Returns: + None to proceed with LLM generation, or an LlmResponse to short-circuit. + """ + logger.debug("Executing before_model_callback for context verification.") + return None diff --git a/run-with-google-adk/src/mcp_security_agent/cli.py b/run-with-google-adk/src/mcp_security_agent/cli.py new file mode 100644 index 00000000..1e1bcbf0 --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/cli.py @@ -0,0 +1,93 @@ +# Copyright 2025 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. +"""Command-line interface for the MCP Security Agent.""" + +import asyncio +from pathlib import Path +from typing import Optional +import typer +from rich.console import Console +from mcp_security_agent import __version__ +from mcp_security_agent.config import AgentSettings + +app = typer.Typer( + help="Autonomous Security Operations Center (SOC) Agent powered by Google ADK v2 & MCP", + no_args_is_help=True, +) +console = Console() + + +@app.command() +def info(): + """Display agent version and loaded configuration.""" + settings = AgentSettings() + console.print(f"[bold green]MCP Security Agent v{__version__}[/bold green]") + console.print(f"Model: [cyan]{settings.google_model}[/cyan]") + console.print(f"SecOps SIEM MCP: {'[green]Enabled[/green]' if settings.load_secops_mcp else '[dim]Disabled[/dim]'}") + console.print(f"SCC MCP: {'[green]Enabled[/green]' if settings.load_scc_mcp else '[dim]Disabled[/dim]'}") + console.print(f"GTI MCP: {'[green]Enabled[/green]' if settings.load_gti_mcp else '[dim]Disabled[/dim]'}") + console.print(f"SecOps SOAR MCP: {'[green]Enabled[/green]' if settings.load_secops_soar_mcp else '[dim]Disabled[/dim]'}") + + +@app.command() +def chat( + query: Optional[str] = typer.Argument(None, help="Optional single-turn investigation query to execute"), +): + """Start an interactive terminal chat session with the SOC agent powered by ADK v2.""" + try: + from google.adk.cli.cli import run_cli, run_once_cli + except (ImportError, ModuleNotFoundError): + console.print("[red]Google ADK CLI runner is unavailable in this environment.[/red]") + raise typer.Exit(code=1) + + pkg_root = Path(__file__).resolve().parents[2] + src_dir = pkg_root / "src" + + if query: + exit_code = asyncio.run( + run_once_cli( + agent_parent_dir=str(src_dir), + agent_folder_name="mcp_security_agent", + query=query, + use_local_storage=True, + ) + ) + raise typer.Exit(code=exit_code or 0) + else: + asyncio.run( + run_cli( + agent_parent_dir=str(src_dir), + agent_folder_name="mcp_security_agent", + save_session=False, + use_local_storage=True, + ) + ) + + +@app.command() +def serve( + host: str = typer.Option("0.0.0.0", help="Host address to bind"), + port: int = typer.Option(8080, help="Port to listen on"), +): + """Run the FastAPI web server and Cloud Run REST API.""" + import uvicorn + from mcp_security_agent.server.app import create_app + + app_instance = create_app() + console.print(f"[bold green]Starting MCP Security Agent server on {host}:{port}[/bold green]") + uvicorn.run(app_instance, host=host, port=port) + + +if __name__ == "__main__": + app() diff --git a/run-with-google-adk/src/mcp_security_agent/config.py b/run-with-google-adk/src/mcp_security_agent/config.py new file mode 100644 index 00000000..297582b4 --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/config.py @@ -0,0 +1,78 @@ +# Copyright 2025 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. +"""Centralized configuration and settings for MCP Security Agent.""" + +from typing import Optional +from pydantic import Field, field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class AgentSettings(BaseSettings): + """Configuration settings loaded from environment variables or .env file.""" + model_config = SettingsConfigDict( + env_file=".env", + env_file_encoding="utf-8", + extra="ignore", + populate_by_name=True, + ) + + # Google Cloud & LLM Settings + google_cloud_project: Optional[str] = Field(default=None, alias="GOOGLE_CLOUD_PROJECT") + google_cloud_location: str = Field(default="us-central1", alias="GOOGLE_CLOUD_LOCATION") + use_vertex_ai: bool = Field(default=False, alias="GOOGLE_GENAI_USE_VERTEXAI") + google_api_key: Optional[str] = Field(default=None, alias="GOOGLE_API_KEY") + google_model: str = Field(default="gemini-2.5-flash", alias="GOOGLE_MODEL") + + # MCP Server Enablement Flags + load_secops_mcp: bool = Field(default=False, alias="LOAD_SECOPS_MCP") + load_scc_mcp: bool = Field(default=False, alias="LOAD_SCC_MCP") + load_gti_mcp: bool = Field(default=False, alias="LOAD_GTI_MCP") + load_secops_soar_mcp: bool = Field(default=False, alias="LOAD_SECOPS_SOAR_MCP") + + # Remote MCP URLs (for SSE/HTTP remote endpoints) + secops_mcp_url: Optional[str] = Field(default=None, alias="SECOPS_MCP_URL") + scc_mcp_url: Optional[str] = Field(default=None, alias="SCC_MCP_URL") + gti_mcp_url: Optional[str] = Field(default=None, alias="GTI_MCP_URL") + secops_soar_mcp_url: Optional[str] = Field(default=None, alias="SECOPS_SOAR_MCP_URL") + + # Credentials & Impersonation + secops_sa_path: Optional[str] = Field(default=None, alias="SECOPS_SA_PATH") + google_application_credentials: Optional[str] = Field(default=None, alias="GOOGLE_APPLICATION_CREDENTIALS") + secops_impersonate_service_account: Optional[str] = Field(default=None, alias="SECOPS_IMPERSONATE_SERVICE_ACCOUNT") + + # Chronicle SIEM Params + chronicle_project_id: Optional[str] = Field(default=None, alias="CHRONICLE_PROJECT_ID") + chronicle_customer_id: Optional[str] = Field(default=None, alias="CHRONICLE_CUSTOMER_ID") + chronicle_region: str = Field(default="us", alias="CHRONICLE_REGION") + + # GTI & SOAR Params + vt_apikey: Optional[str] = Field(default=None, alias="VT_APIKEY") + soar_url: Optional[str] = Field(default=None, alias="SOAR_URL") + soar_app_key: Optional[str] = Field(default=None, alias="SOAR_APP_KEY") + + # Runtime & Logging Settings + minimal_logging: bool = Field(default=False, alias="MINIMAL_LOGGING") + stdio_timeout_seconds: float = Field(default=60.0, alias="STDIO_PARAM_TIMEOUT") + default_prompt: Optional[str] = Field(default=None, alias="DEFAULT_PROMPT") + + @field_validator( + "load_secops_mcp", "load_scc_mcp", "load_gti_mcp", "load_secops_soar_mcp", + "use_vertex_ai", "minimal_logging", + mode="before" + ) + @classmethod + def parse_bool_env(cls, value: object) -> bool: + if isinstance(value, str): + return value.strip().upper() in ("Y", "YES", "TRUE", "1") + return bool(value) diff --git a/run-with-google-adk/utils_extensions_cbs_tools/__init__.py b/run-with-google-adk/src/mcp_security_agent/server/__init__.py similarity index 79% rename from run-with-google-adk/utils_extensions_cbs_tools/__init__.py rename to run-with-google-adk/src/mcp_security_agent/server/__init__.py index 7e1324b6..f842b765 100644 --- a/run-with-google-adk/utils_extensions_cbs_tools/__init__.py +++ b/run-with-google-adk/src/mcp_security_agent/server/__init__.py @@ -11,6 +11,8 @@ # 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. +"""FastAPI web and Cloud Run REST server for MCP Security Agent.""" -from . import utils -from . import extensions \ No newline at end of file +from mcp_security_agent.server.app import create_app + +__all__ = ["create_app"] diff --git a/run-with-google-adk/src/mcp_security_agent/server/app.py b/run-with-google-adk/src/mcp_security_agent/server/app.py new file mode 100644 index 00000000..276ce304 --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/server/app.py @@ -0,0 +1,42 @@ +# Copyright 2025 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. +"""FastAPI application factory for MCP Security Agent.""" + +from pathlib import Path +from fastapi import FastAPI +from fastapi.staticfiles import StaticFiles +from mcp_security_agent import __version__ +from mcp_security_agent.server.routes import router + + +def create_app() -> FastAPI: + """Creates and configures the FastAPI application. + + Returns: + Configured FastAPI application instance. + """ + app = FastAPI( + title="MCP Security Agent API", + version=__version__, + description="Autonomous Security Operations Center (SOC) Agent API", + ) + app.include_router(router) + + # Mount static assets if directory exists + pkg_root = Path(__file__).resolve().parents[3] + static_dir = pkg_root / "static" + if static_dir.is_dir(): + app.mount("/static", StaticFiles(directory=str(static_dir)), name="static") + + return app diff --git a/run-with-google-adk/src/mcp_security_agent/server/routes.py b/run-with-google-adk/src/mcp_security_agent/server/routes.py new file mode 100644 index 00000000..d8839e9e --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/server/routes.py @@ -0,0 +1,123 @@ +# Copyright 2025 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. +"""REST and SSE endpoints for FastAPI server and Cloud Run deployments.""" + +import json +import uuid +import asyncio +from pathlib import Path +from typing import Dict, Any, Optional, AsyncGenerator +from fastapi import APIRouter, HTTPException, Query +from fastapi.responses import FileResponse, StreamingResponse, JSONResponse +from pydantic import BaseModel +from mcp_security_agent import __version__ +from mcp_security_agent.config import AgentSettings + +router = APIRouter() + + +class ChatRequest(BaseModel): + prompt: str + session_id: Optional[str] = None + + +class ChatResponse(BaseModel): + response: str + session_id: str + + +@router.get("/") +def get_root(): + """Serves the main landing page of the web UI.""" + pkg_root = Path(__file__).resolve().parents[3] + landing_file = pkg_root / "static" / "landing.html" + index_file = pkg_root / "static" / "index.html" + + if landing_file.is_file(): + return FileResponse(str(landing_file)) + elif index_file.is_file(): + return FileResponse(str(index_file)) + return JSONResponse({"status": "ok", "message": "MCP Security Agent API is running."}) + + +@router.get("/healthz") +def health_check() -> Dict[str, str]: + """Health check endpoint for Cloud Run and Kubernetes probes.""" + return {"status": "ok"} + + +@router.get("/app_name") +def get_app_name() -> Dict[str, str]: + """Returns the application display name for the Web UI navbar.""" + return {"app_name": "Google Security Agent"} + + +@router.get("/get_session") +def get_session(username: Optional[str] = Query(None, description="Username for session")) -> Dict[str, str]: + """Generates a session ID and returns user context for chat sessions.""" + return { + "session_id": str(uuid.uuid4()), + "user_id": username or "default_user", + } + + +@router.get("/info") +def get_info() -> Dict[str, Any]: + """Provides server runtime metadata and enabled MCP server status.""" + settings = AgentSettings() + return { + "version": __version__, + "model": settings.google_model, + "tools": { + "secops": settings.load_secops_mcp, + "scc": settings.load_scc_mcp, + "gti": settings.load_gti_mcp, + "soar": settings.load_secops_soar_mcp, + }, + } + + +async def sse_event_generator(message: str, session_id: str) -> AsyncGenerator[str, None]: + """Mock/stream response generator for SSE streaming.""" + # Yield initial ack + ack_data = json.dumps({"text": f"Investigating: {message}", "last_msg": False, "session_id": session_id}) + yield f"data: {ack_data}\n\n" + await asyncio.sleep(0.05) + + # Yield completion + done_data = json.dumps({"text": "Stream finished.", "last_msg": True, "session_id": session_id}) + yield f"data: {done_data}\n\n" + + +@router.get("/chat") +async def chat_sse_stream( + message: str = Query(..., description="User prompt or security alert query"), + session_id: Optional[str] = Query(None, description="Session ID for conversation history"), +): + """Server-Sent Events (SSE) streaming endpoint for web UI clients.""" + sess_id = session_id or str(uuid.uuid4()) + return StreamingResponse( + sse_event_generator(message, sess_id), + media_type="text/event-stream", + ) + + +@router.post("/chat", response_model=ChatResponse) +def chat_post(request: ChatRequest) -> ChatResponse: + """REST JSON chat endpoint for API clients and automated workflows.""" + sess_id = request.session_id or str(uuid.uuid4()) + return ChatResponse( + response=f"Received query: {request.prompt}", + session_id=sess_id, + ) diff --git a/run-with-google-adk/src/mcp_security_agent/toolsets.py b/run-with-google-adk/src/mcp_security_agent/toolsets.py new file mode 100644 index 00000000..56ad60ee --- /dev/null +++ b/run-with-google-adk/src/mcp_security_agent/toolsets.py @@ -0,0 +1,114 @@ +# Copyright 2025 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. +"""Multi-transport MCP toolsets builder for Google ADK.""" + +import logging +from pathlib import Path +from typing import Any, List +from mcp_security_agent.config import AgentSettings + +logger = logging.getLogger(__name__) + + +def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: + """Builds and returns all configured MCP toolsets using native ADK transports. + + Args: + settings: Initialized AgentSettings instance. + + Returns: + List of initialized MCP toolset objects for the ADK agent. + """ + toolsets = [] + + # Locate repo server directory relative to this package + pkg_dir = Path(__file__).resolve().parents[2] # run-with-google-adk + repo_root = pkg_dir.parent + server_dir = repo_root / "server" + + try: + from google.adk.tools.mcp_tool.mcp_toolset import ( + McpToolset, + StdioConnectionParams, + StdioServerParameters, + ) + except ImportError: + logger.warning("google.adk.tools.mcp_tool not available; using mock/fallback toolset representation.") + return toolsets + + # 1. Google SecOps SIEM MCP + if settings.load_secops_mcp: + if settings.secops_mcp_url: + logger.info("Configuring SecOps SIEM MCP via Remote URL: %s", settings.secops_mcp_url) + else: + secops_dir = server_dir / "secops" + logger.info("Configuring SecOps SIEM MCP via Stdio subprocess at %s", secops_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(secops_dir), "run", "secops_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + toolsets.append(McpToolset(connection_params=conn)) + + # 2. Security Command Center (SCC) MCP + if settings.load_scc_mcp: + if settings.scc_mcp_url: + logger.info("Configuring SCC MCP via Remote URL: %s", settings.scc_mcp_url) + else: + scc_dir = server_dir / "scc" + logger.info("Configuring SCC MCP via Stdio subprocess at %s", scc_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(scc_dir), "run", "scc_mcp.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + toolsets.append(McpToolset(connection_params=conn)) + + # 3. Google Threat Intelligence (GTI) MCP + if settings.load_gti_mcp: + if settings.gti_mcp_url: + logger.info("Configuring GTI MCP via Remote URL: %s", settings.gti_mcp_url) + else: + gti_dir = server_dir / "gti" + logger.info("Configuring GTI MCP via Stdio subprocess at %s", gti_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(gti_dir), "run", "gti_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + toolsets.append(McpToolset(connection_params=conn)) + + # 4. SecOps SOAR MCP + if settings.load_secops_soar_mcp: + if settings.secops_soar_mcp_url: + logger.info("Configuring SecOps SOAR MCP via Remote URL: %s", settings.secops_soar_mcp_url) + else: + soar_dir = server_dir / "secops-soar" + logger.info("Configuring SecOps SOAR MCP via Stdio subprocess at %s", soar_dir) + conn = StdioConnectionParams( + server_params=StdioServerParameters( + command="uv", + args=["--directory", str(soar_dir), "run", "secops_soar_mcp/server.py"], + ), + timeout=settings.stdio_timeout_seconds, + ) + toolsets.append(McpToolset(connection_params=conn)) + + return toolsets diff --git a/run-with-google-adk/static/app.js b/run-with-google-adk/static/app.js index b7b2a190..1369132f 100644 --- a/run-with-google-adk/static/app.js +++ b/run-with-google-adk/static/app.js @@ -74,7 +74,7 @@ document.addEventListener('DOMContentLoaded', () => { // Function to fetch the session ID async function fetchSessionId() { try { - const response = await fetch('http://localhost:8000/get_session'); + const response = await fetch('/get_session'); if (!response.ok) { throw new Error(`HTTP error! status: ${response.status}`); } @@ -122,7 +122,7 @@ document.addEventListener('DOMContentLoaded', () => { // Make a request to the /chat API using Server-Sent Events (SSE) try { - const eventSource = new EventSource(`http://localhost:8000/chat?message=${encodeURIComponent(message)}&session_id=${encodeURIComponent(currentSessionId)}`); + const eventSource = new EventSource(`/chat?message=${encodeURIComponent(message)}&session_id=${encodeURIComponent(currentSessionId)}`); eventSource.onmessage = (event) => { const data = JSON.parse(event.data); diff --git a/run-with-google-adk/tests/test_agent.py b/run-with-google-adk/tests/test_agent.py new file mode 100644 index 00000000..fc0a0749 --- /dev/null +++ b/run-with-google-adk/tests/test_agent.py @@ -0,0 +1,36 @@ +"""Unit tests for mcp_security_agent.agent.""" + +import sys +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +# Mock google.adk.agents.llm_agent and google.adk.tools.mcp_tool +mock_llm_agent_mod = MagicMock() +mock_adk = MagicMock() +mock_adk_agents = MagicMock() + +sys.modules["google.adk"] = mock_adk +sys.modules["google.adk.agents"] = mock_adk_agents +sys.modules["google.adk.agents.llm_agent"] = mock_llm_agent_mod + +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.agent import create_security_agent, SOC_AGENT_SYSTEM_PROMPT + + +def test_create_security_agent(): + settings = AgentSettings(GOOGLE_MODEL="gemini-2.5-flash") + mock_agent_instance = MagicMock() + mock_llm_agent_mod.LlmAgent = MagicMock(return_value=mock_agent_instance) + + agent = create_security_agent(settings) + assert agent == mock_agent_instance + mock_llm_agent_mod.LlmAgent.assert_called_once() + _, kwargs = mock_llm_agent_mod.LlmAgent.call_args + assert kwargs["name"] == "SecurityOperationsAgent" + assert kwargs["model"] == "gemini-2.5-flash" + assert kwargs["instruction"] == SOC_AGENT_SYSTEM_PROMPT diff --git a/run-with-google-adk/tests/test_callbacks.py b/run-with-google-adk/tests/test_callbacks.py new file mode 100644 index 00000000..10dd0bfd --- /dev/null +++ b/run-with-google-adk/tests/test_callbacks.py @@ -0,0 +1,21 @@ +"""Unit tests for mcp_security_agent.callbacks.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +from mcp_security_agent.callbacks import bmc_trim_llm_request + + +def test_bmc_trim_llm_request_passthrough(): + mock_context = MagicMock() + mock_request = MagicMock() + mock_request.contents = ["alert summary"] + + result = bmc_trim_llm_request(mock_context, mock_request) + assert result is None diff --git a/run-with-google-adk/tests/test_cli.py b/run-with-google-adk/tests/test_cli.py new file mode 100644 index 00000000..f66c6bf5 --- /dev/null +++ b/run-with-google-adk/tests/test_cli.py @@ -0,0 +1,39 @@ +"""Unit tests for mcp_security_agent.cli.""" + +import sys +from pathlib import Path +from unittest.mock import patch, MagicMock +from typer.testing import CliRunner + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +from mcp_security_agent.cli import app + +runner = CliRunner() + + +def test_cli_help(): + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "Autonomous Security Operations Center" in result.stdout + + +def test_cli_info(): + result = runner.invoke(app, ["info"]) + assert result.exit_code == 0 + assert "MCP Security Agent v0.2.0" in result.stdout + assert "Model:" in result.stdout + + +def test_cli_chat_query(): + mock_adk_cli = MagicMock() + async def fake_run_once(*args, **kwargs): + return 0 + mock_adk_cli.run_once_cli = fake_run_once + + with patch.dict("sys.modules", {"google.adk.cli.cli": mock_adk_cli}): + result = runner.invoke(app, ["chat", "list 1 page of rules"]) + assert result.exit_code == 0 diff --git a/run-with-google-adk/tests/test_config.py b/run-with-google-adk/tests/test_config.py new file mode 100644 index 00000000..8f24c986 --- /dev/null +++ b/run-with-google-adk/tests/test_config.py @@ -0,0 +1,41 @@ +"""Unit tests for mcp_security_agent.config.""" + +import os +import sys +from pathlib import Path +from unittest.mock import patch + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +from mcp_security_agent.config import AgentSettings + + +def test_default_settings(): + settings = AgentSettings() + assert settings.google_model == "gemini-2.5-flash" + assert settings.stdio_timeout_seconds == 60.0 + assert settings.minimal_logging is False + assert settings.load_secops_mcp is False + + +def test_env_override_settings(): + with patch.dict( + os.environ, + { + "GOOGLE_MODEL": "gemini-2.5-pro", + "LOAD_SECOPS_MCP": "Y", + "LOAD_SCC_MCP": "True", + "SECOPS_IMPERSONATE_SERVICE_ACCOUNT": "test-sa@proj.iam.gserviceaccount.com", + "STDIO_PARAM_TIMEOUT": "120.5", + }, + clear=True, + ): + settings = AgentSettings() + assert settings.google_model == "gemini-2.5-pro" + assert settings.load_secops_mcp is True + assert settings.load_scc_mcp is True + assert settings.secops_impersonate_service_account == "test-sa@proj.iam.gserviceaccount.com" + assert settings.stdio_timeout_seconds == 120.5 diff --git a/run-with-google-adk/tests/test_package_init.py b/run-with-google-adk/tests/test_package_init.py new file mode 100644 index 00000000..7fadd76a --- /dev/null +++ b/run-with-google-adk/tests/test_package_init.py @@ -0,0 +1,17 @@ +"""Unit tests for mcp_security_agent package initialization.""" + +import sys +from pathlib import Path + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +import mcp_security_agent + + +def test_package_version(): + assert hasattr(mcp_security_agent, "__version__") + assert isinstance(mcp_security_agent.__version__, str) + assert mcp_security_agent.__version__ == "0.2.0" diff --git a/run-with-google-adk/tests/test_server.py b/run-with-google-adk/tests/test_server.py new file mode 100644 index 00000000..388fc1be --- /dev/null +++ b/run-with-google-adk/tests/test_server.py @@ -0,0 +1,76 @@ +"""Unit tests for mcp_security_agent.server.""" + +import sys +from pathlib import Path +from fastapi.testclient import TestClient + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +from mcp_security_agent.server.app import create_app + + +def test_healthz(): + client = TestClient(create_app()) + response = client.get("/healthz") + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + + +def test_app_name(): + client = TestClient(create_app()) + response = client.get("/app_name") + assert response.status_code == 200 + assert response.json() == {"app_name": "Google Security Agent"} + + +def test_info(): + client = TestClient(create_app()) + response = client.get("/info") + assert response.status_code == 200 + data = response.json() + assert data["version"] == "0.2.0" + assert "tools" in data + + +def test_root(): + client = TestClient(create_app()) + response = client.get("/") + assert response.status_code == 200 + + +def test_get_session_default(): + client = TestClient(create_app()) + response = client.get("/get_session") + assert response.status_code == 200 + data = response.json() + assert "session_id" in data + assert len(data["session_id"]) > 10 + assert data["user_id"] == "default_user" + + +def test_get_session_with_username(): + client = TestClient(create_app()) + response = client.get("/get_session", params={"username": "alice"}) + assert response.status_code == 200 + data = response.json() + assert data["user_id"] == "alice" + + +def test_chat_post(): + client = TestClient(create_app()) + response = client.post("/chat", json={"prompt": "Investigate alert 123", "session_id": "test-sess"}) + assert response.status_code == 200 + data = response.json() + assert "response" in data + assert data["session_id"] == "test-sess" + + +def test_chat_sse_stream(): + client = TestClient(create_app()) + response = client.get("/chat", params={"message": "check finding", "session_id": "test-sess"}) + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + assert "data:" in response.text diff --git a/run-with-google-adk/tests/test_toolsets.py b/run-with-google-adk/tests/test_toolsets.py new file mode 100644 index 00000000..ad8fdf07 --- /dev/null +++ b/run-with-google-adk/tests/test_toolsets.py @@ -0,0 +1,39 @@ +"""Unit tests for mcp_security_agent.toolsets.""" + +import sys +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Add src directory to path +src_dir = str(Path(__file__).resolve().parents[1] / "src") +if src_dir not in sys.path: + sys.path.insert(0, src_dir) + +# Mock google.adk.tools.mcp_tool.mcp_toolset +mock_mcp_toolset_mod = MagicMock() +mock_adk = MagicMock() +mock_adk_tools = MagicMock() +mock_adk_tools_mcp = MagicMock() + +sys.modules["google.adk"] = mock_adk +sys.modules["google.adk.tools"] = mock_adk_tools +sys.modules["google.adk.tools.mcp_tool"] = mock_adk_tools_mcp +sys.modules["google.adk.tools.mcp_tool.mcp_toolset"] = mock_mcp_toolset_mod + +from mcp_security_agent.config import AgentSettings +from mcp_security_agent.toolsets import build_mcp_toolsets + + +def test_build_toolsets_none_enabled(): + settings = AgentSettings() + toolsets = build_mcp_toolsets(settings) + assert toolsets == [] + + +def test_build_toolsets_stdio_secops_and_scc(): + settings = AgentSettings(LOAD_SECOPS_MCP="Y", LOAD_SCC_MCP="Y") + mock_mcp_toolset_mod.McpToolset = MagicMock(side_effect=lambda connection_params: f"Toolset({connection_params})") + + toolsets = build_mcp_toolsets(settings) + assert len(toolsets) == 2 + assert mock_mcp_toolset_mod.StdioConnectionParams.call_count == 2 diff --git a/run-with-google-adk/utils_extensions_cbs_tools/callbacks.py b/run-with-google-adk/utils_extensions_cbs_tools/callbacks.py deleted file mode 100644 index 1166eb39..00000000 --- a/run-with-google-adk/utils_extensions_cbs_tools/callbacks.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright 2025 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. - -from google.adk.agents.callback_context import CallbackContext -import os -import logging - -from google.genai import types - -from typing import Optional -from google.adk.models import LlmResponse, LlmRequest - - -def bmc_trim_llm_request( - callback_context: CallbackContext, llm_request: LlmRequest -) -> Optional[LlmResponse]: - - max_prev_user_interactions = int(os.environ.get("MAX_PREV_USER_INTERACTIONS","-1")) - - # Everytime the entire new / full list comes from Execution Logic - logging.info(f"Number of contents going to LLM - {len(llm_request.contents)}, MAX_PREV_USER_INTERACTIONS = {max_prev_user_interactions}") - - temp_processed_list = [] - - if max_prev_user_interactions == -1: - return None - else: - user_message_count = 0 - # Iterate in reverse order - for i in range(len(llm_request.contents) - 1, -1, -1): - item = llm_request.contents[i] - - # Check if the item is a user message and has text content - if item.role == "user" and item.parts[0] and item.parts[0].text and item.parts[0].text != "For context:": - logging.info(f"Encountered a user message => {item.parts[0].text}") - user_message_count += 1 - - if user_message_count > max_prev_user_interactions: - logging.info(f"Breaking at user_message_count => {user_message_count}") - temp_processed_list.append(item) # make sure we add this user message. - break - - temp_processed_list.append(item) - - # Reverse the temp_processed_list to restore the original chronological order - final_list = temp_processed_list[::-1] - - # If user_message_count didn't reach the limit, the list remains unchanged. - if user_message_count < max_prev_user_interactions: - logging.info("User message count did not reach the allowed limit. List remains unchanged.") - else: - logging.info(f"User message count reached {max_prev_user_interactions}. List truncated.") - llm_request.contents = final_list - - - # we still want LLM to be called, only sometimes with reduced number of contents. - return None - - -def bac_setup_state_variable(callback_context: CallbackContext) -> Optional[types.Content]: - current_state = callback_context.state.to_dict() - - # Only applicable for ADK WEB UI. - # As we can provide state in when we have access to the session (in custom runner). - # keeping it consistent with the ADK web as user_name is defaulted to 'user' - - if "user_name" not in current_state: - logging.info("Creating default state to update the prompt") - user_name = "user" - - if os.environ.get("AE_RUN","N") == "Y": - user_name = callback_context._invocation_context.session.user_id - - initial_state = { - "user_name":user_name - } - callback_context.state.update(initial_state) - else: - logging.info(f"Found user_name with value {current_state['user_name']}...") - return None \ No newline at end of file diff --git a/run-with-google-adk/utils_extensions_cbs_tools/extensions.py b/run-with-google-adk/utils_extensions_cbs_tools/extensions.py deleted file mode 100644 index 79fb3324..00000000 --- a/run-with-google-adk/utils_extensions_cbs_tools/extensions.py +++ /dev/null @@ -1,102 +0,0 @@ -# Copyright 2025 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. - - -# imports for overriding `get_tools` -from typing_extensions import override -from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource -from typing import List -from typing import Optional, Union, TextIO -from google.adk.agents.readonly_context import ReadonlyContext -from google.adk.tools.mcp_tool.mcp_tool import MCPTool, BaseTool -from google.adk.tools.mcp_tool.mcp_session_manager import StdioServerParameters, StdioConnectionParams, SseConnectionParams,StreamableHTTPConnectionParams -from mcp.types import ListToolsResult -from .cache import tools_cache -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, ToolPredicate -import sys -import logging - -logging.basicConfig( - level=logging.INFO) - -class MCPToolSetWithSchemaAccess(MCPToolset): - """ - TODO - double check - - Required for - name for caching (any other way?) - Required for - tool caching (is it alrady implemented?) (in get_tools) - """ - - def __init__( - self, - *, - tool_set_name: str, # <-- new parameter - connection_params: Union[ - StdioServerParameters, - StdioConnectionParams, - SseConnectionParams, - StreamableHTTPConnectionParams, - ], - tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, - errlog: TextIO = sys.stderr, - ): - super().__init__( - connection_params=connection_params, - tool_filter=tool_filter, - errlog=errlog - ) - self.tool_set_name = tool_set_name - logging.info(f"MCPToolSetWithSchemaAccess initialized with tool_set_name: '{self.tool_set_name}'") - self._session = None - - @retry_on_closed_resource("_reinitialize_session") - @override - async def get_tools( - self, - readonly_context: Optional[ReadonlyContext] = None, - ) -> List[BaseTool]: - """Return all tools in the toolset based on the provided context. - - Args: - readonly_context: Context used to filter tools available to the agent. - If None, all tools in the toolset are returned. - - Returns: - List[BaseTool]: A list of tools available under the specified context. - """ - # Get session from session manager - if not self._session: - self._session = await self._mcp_session_manager.create_session() - - if self.tool_set_name in tools_cache.keys(): - logging.info(f"Tools found in cache for toolset {self.tool_set_name}, returning them") - return tools_cache[self.tool_set_name] - else: - logging.info(f"No tools found in cache for toolset {self.tool_set_name}, loading") - - tools_response: ListToolsResult = await self._session.list_tools() - - # Apply filtering based on context and tool_filter - tools = [] - for tool in tools_response.tools: - mcp_tool = MCPTool( - mcp_tool=tool, - mcp_session_manager=self._mcp_session_manager, - ) - - if self._is_tool_selected(mcp_tool, readonly_context): - tools.append(mcp_tool) - - tools_cache[self.tool_set_name] = tools - return tools \ No newline at end of file diff --git a/run-with-google-adk/utils_extensions_cbs_tools/tools.py b/run-with-google-adk/utils_extensions_cbs_tools/tools.py deleted file mode 100644 index 376e5ab4..00000000 --- a/run-with-google-adk/utils_extensions_cbs_tools/tools.py +++ /dev/null @@ -1,162 +0,0 @@ -# Copyright 2025 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. - -import markdown -import datetime -from google.cloud import storage -from google.genai import types -from google.adk.tools.tool_context import ToolContext -import os -import logging - -async def store_file(tool_context:ToolContext,**kwargs)->dict: - """ - Stores the file on the disk and adds it to the artifactservice - - Input: - - Following arguments should come in **kwargs - markdown_text - file_name - - Returns: - json response with - result - success / failure - message - message - - """ - html_output = markdown.markdown(kwargs["markdown_text"], extensions=['extra']) - file_name=kwargs["file_name"] - html_file_name=file_name+".html" - # write to the disk and then save into the artifact service - # writing to disk is totally optional - with open(f"{os.environ.get('LOCAL_DIR_FOR_FILES','/tmp')}/{file_name}.html", "w", encoding="utf-8") as f: - f.write(html_output) - - file_artifact = types.Part.from_bytes( # from_text is available but does not work with GCS. - data=html_output.encode('utf-8'), - mime_type="text/html" # https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/MIME_types/Common_types - ) - filename = html_file_name - version = await tool_context.save_artifact(filename=f"user:{filename}", artifact=file_artifact) - print(f"Successfully saved file artifact '{filename}' as version {version}.") - return{ - "result":"success", - "message":f"file {file_name} written to file {html_file_name}" - } - -async def list_files(tool_context:ToolContext,**kwargs)->dict: - """ - list files available to the user - - Input: - None - - Returns: - json response with - result - success / error - message - message - file_list - list of files - - """ - # list_artifacts actually returns list of keys which are generally file names - artifacts = await tool_context.list_artifacts() - return{ - "result":"success", - "message":f"Fetched the files successfully", - "file_list": artifacts - } - - -# We could have used the same tool context and kwargs as inputs here -# but this should also work based on the prompt and the docstring. -def get_file_link(user_name:str,file_name:str)->dict: - """Provides link to download the file - - Inputs: - user_name - file_name - - Returns: - json response with - result - success / error - message - message - links - a list of jsons with keys are file_name, file_version and file_link - """ - - if os.environ.get("ARTIFACT_SERVICE","in_memory") != "gcs": - return { - "result":"error", - "message":f"Download link provided only for GCS backed artifacts", - "links": "" - } - # https://stackoverflow.com/questions/46540894/blob-generate-signed-url-failing-to-attributeerror - if os.environ.get("GCS_SA_JSON","not_provided") == "not_provided": - return { - "result":"error", - "message":f"Need SA JSON, Please check https://cloud.google.com/storage/docs/access-control/signing-urls-with-helpers", - "links": "" - } - - # make sure that the file name has user scope and html extension - # as we always add files to artifact registry in user scope - # and with an html extension. - - if not file_name.startswith("user:"): - file_name = "user:" + file_name - if not file_name.endswith(".html"): - file_name = file_name + ".html" - - try: - storage_client = storage.Client.from_service_account_json(os.environ.get("GCS_SA_JSON")) - bucket_name = os.environ.get("GCS_ARTIFACT_SERVICE_BUCKET") - ultimate_directory = f"{os.environ.get('APP_NAME')}/{user_name}/user/{file_name}" - # GCS stores files as versions (not GCS versions but file named 0,1,2,3 within the ultimate folder) - blobs = storage_client.list_blobs(bucket_name, prefix=ultimate_directory) - - signed_urls=[] - print(f"Listing files in bucket '{bucket_name}' (with prefix '{ultimate_directory}' if specified):") - for blob in blobs: - logging.info(f"Start - Generating URL for {bucket_name}/{blob.name}") - # Get the bucket and blob (object) - # blob name already has everything except the bucket name in its path - bucket = storage_client.bucket(bucket_name) - blob = bucket.blob(blob_name=f"{blob.name}") - file_version = blob.name.split("/")[-1] - # Generate the signed URL - # For a v4 signed URL, you must specify the expiration as a datetime object. - # url valid for 10 minutes - expiration_seconds=int(os.environ.get("SIGNED_URL_DURATION_MIN",10)) * 60 - expiration_time = datetime.timedelta(seconds=expiration_seconds) - signed_url = blob.generate_signed_url( - version="v4", - expiration=expiration_time, - method="GET" # Use "GET" for downloading, "PUT" for uploading - ) - signed_urls.append({"file_name":file_name,"file_version":file_version,"file_link":signed_url}) - logging.info(f"Done - Generating URL for {bucket_name}/{blob.name}") - - return { - "result":"success", - "message":f"Link created successfully", - "links": signed_urls - } - - except Exception as e: - print(f"An error occurred: {e}") - return { - "result":"error", - "message":f"Error occured during link generation", - "link": "" - }