diff --git a/.gitignore b/.gitignore index 7596086b..4d137a1c 100644 --- a/.gitignore +++ b/.gitignore @@ -132,8 +132,10 @@ celerybeat.pid *.sage.py # Environments -.env +.env* +.envrc .venv +.gcloud/ env/ venv/ ENV/ @@ -212,4 +214,7 @@ app_data.db .gemini/ # devcontainer -.devcontainer/ \ No newline at end of file +.devcontainer/ + +# IDE / Project tools +.cm_project \ No newline at end of file diff --git a/README.md b/README.md index 77326dc1..e8a97331 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,7 @@ You can also use environment files with uvx: ## Client Configurations The MCP servers from this repo can be used with the following clients 1. Cline, Claude Desktop, and other MCP supported clients -2. [Google ADK(Agent Development Kit)](https://google.github.io/adk-docs/) Agents (a prebuilt agent is provided, details [below](#using-the-prebuilt-google-adk-agent-as-client)) +2. [Google ADK(Agent Development Kit)](https://google.github.io/adk-docs/) Agents (a prebuilt agent is provided, details [below](#using-the-google-adk-autonomous-soc-agent)) 3. [Google SecOps Extension](https://google.github.io/mcp-security/google_secops_extension.html) - Install our example extension for Gemini CLI to get specialized security skills (Triage, Investigate, Hunt). 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. @@ -172,7 +172,7 @@ It can be run locally via an interactive CLI REPL or launched as a FastAPI servi ```bash cd run-with-google-adk -cp sample.env .env +cp .env.example .env # Interactive terminal investigation REPL uv run mcp-security-agent chat diff --git a/docs/usage_guide.md b/docs/usage_guide.md index 40f9f610..757d6295 100644 --- a/docs/usage_guide.md +++ b/docs/usage_guide.md @@ -50,7 +50,7 @@ The repository provides a prebuilt Autonomous Security Operations Center (SOC) A ```bash cd run-with-google-adk -cp sample.env .env +cp .env.example .env # Interactive terminal investigation REPL uv run mcp-security-agent chat diff --git a/run-with-google-adk/.env.example b/run-with-google-adk/.env.example new file mode 100644 index 00000000..ddfea8ea --- /dev/null +++ b/run-with-google-adk/.env.example @@ -0,0 +1,35 @@ +# Google Cloud & LLM Settings +GOOGLE_CLOUD_PROJECT= +GOOGLE_CLOUD_LOCATION=us-central1 +GOOGLE_GENAI_USE_VERTEXAI=False +GOOGLE_API_KEY= +GOOGLE_MODEL=gemini-2.5-flash + +# MCP Server Enablement Flags +# (SCC uses Google Cloud ADC. SecOps, GTI, and SOAR auto-enable when credentials below are populated. +# You can explicitly set LOAD_*_MCP to Y or N to override auto-detection.) +LOAD_SCC_MCP=Y +# LOAD_SECOPS_MCP= +# LOAD_GTI_MCP= +# LOAD_SECOPS_SOAR_MCP= + +# Credentials & Service Account Impersonation +SECOPS_SA_PATH= +GOOGLE_APPLICATION_CREDENTIALS= +SECOPS_IMPERSONATE_SERVICE_ACCOUNT= + +# Google SecOps (Chronicle SIEM) Settings - Populating enables SecOps MCP +CHRONICLE_PROJECT_ID= +CHRONICLE_CUSTOMER_ID= +CHRONICLE_REGION=us + +# Google Threat Intelligence (GTI / VirusTotal) - Populating enables GTI MCP +VT_APIKEY= + +# SecOps SOAR Settings - Populating enables SOAR MCP +SOAR_URL= +SOAR_APP_KEY= + +# Runtime Settings +STDIO_PARAM_TIMEOUT=60.0 +MINIMAL_LOGGING=N diff --git a/run-with-google-adk/README.md b/run-with-google-adk/README.md index 6c895121..68de6bed 100644 --- a/run-with-google-adk/README.md +++ b/run-with-google-adk/README.md @@ -4,326 +4,258 @@ This guide provides instructions on how to run the Autonomous Security Operation ## Table of Contents -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. [Prerequisites](#prerequisites) +2. [Quickstart: Running Agent Locally](#quickstart-running-agent-locally) +3. [CLI Commands & Options](#cli-commands--options) +4. [Configuration & Environment Variables](#configuration--environment-variables) +5. [Running the Web UI & API Server](#running-the-web-ui--api-server) +6. [Deploying to Google Cloud Run](#deploying-to-google-cloud-run) +7. [Integrating Custom MCP Servers](#integrating-custom-mcp-servers) +8. [Testing & Verification](#testing--verification) --- -## 1. Quickstart: Running Agent Locally +## Prerequisites + +1. **Python 3.11+** +2. [**uv**](https://docs.astral.sh/uv/) (recommended) or `pip` +3. **Google Cloud Account / Project** with access to one or more of: + * Google SecOps (Chronicle SIEM) + * Google Cloud Security Command Center (SCC) + * Google Threat Intelligence (GTI / VirusTotal) + * Google SecOps SOAR (Siemplify) +4. **Authentication**: + * For Google Cloud services (Chronicle, SCC, Vertex AI): Authenticate using Application Default Credentials: + ```bash + gcloud auth application-default login + ``` + * Alternatively, to use the Gemini Developer API without Vertex AI, obtain an API key from [Google AI Studio](https://ai.google.dev/gemini-api/docs/api-key) and configure `GOOGLE_API_KEY`. -### 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 - -### Installation & Execution - -```bash -# Clone the repository -git clone https://github.com/google/mcp-security.git -cd mcp-security/run-with-google-adk +--- -# Copy the sample environment file and configure your API keys / project IDs -cp sample.env .env +## Quickstart: Running Agent Locally -# Start interactive chat session -uv run mcp-security-agent chat -``` +### 1. Setup Environment -Alternatively, install in editable mode: ```bash -python3 -m venv .venv -source .venv/bin/activate -pip install -e . +cd run-with-google-adk -mcp-security-agent info -mcp-security-agent chat +# Copy the sample environment template +cp .env.example .env ``` -## 2. CLI Commands & Subcommands - -The package exposes the `mcp-security-agent` CLI with the following commands: +Configure your credentials and project IDs in `.env` (see [Configuration](#configuration--environment-variables)). -* `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. +### 2. Verify Diagnostics -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. -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). +Run `info` to verify package installation, active model, resolved GCP project/ADC credentials, and configured MCP servers: ```bash -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 -#GCS_ARTIFACT_SERVICE_BUCKET=your-bucket-name -#GCS_SA_JSON=object-viewer-sa.json - -# 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 +uv run mcp-security-agent info +``` -# MANDATORY -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.5-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='Helps 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.' +### 3. Launch Interactive Chat -# Initially a long timeout is needed -# to load the tools and install dependencies -STDIO_PARAM_TIMEOUT=60.0 +Start an interactive threat investigation session with your desired MCP toolsets: +```bash +# Enable Chronicle SIEM MCP tools (alerts, UDM search, rules) +uv run mcp-security-agent chat --secops -# 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 +# Enable both SecOps SIEM and Security Command Center (SCC) +uv run mcp-security-agent chat --secops --scc +``` -# HIGHLY RECOMMENDED TO SET Y AFTER INITIAL TESTING ON CLOUD RUN -MINIMAL_LOGGING=N +You can also provide an initial prompt directly on the command line: -# 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 +```bash +uv run mcp-security-agent chat --secops "List recent critical security alerts from the past 24 hours" +``` +## CLI Commands & Options + +The package provides the `mcp-security-agent` CLI entry point. + +### `mcp-security-agent info` +Displays current runtime diagnostics: +* Package version +* Active Gemini model & Vertex AI status +* Resolved Google Cloud Project ID and Application Default Credentials (ADC) path +* Toolset status for each supported MCP server + +### `mcp-security-agent chat [PROMPT] [OPTIONS]` +Starts an interactive terminal REPL for security investigations. + +| Option | Type | Description | +| :--- | :--- | :--- | +| `PROMPT` | Argument | Optional initial prompt to execute immediately upon startup. | +| `--secops / --no-secops` | Flag | Enable or disable Chronicle SIEM MCP (`server/secops`). | +| `--scc / --no-scc` | Flag | Enable or disable Security Command Center MCP (`server/scc`). | +| `--gti / --no-gti` | Flag | Enable or disable Google Threat Intelligence MCP (`server/gti`). | +| `--soar / --no-soar` | Flag | Enable or disable SecOps SOAR MCP (`server/secops-soar`). | +| `--vertex / --no-vertex` | Flag | Toggle Vertex AI (`--vertex`) vs Gemini Developer API (`--no-vertex`). | +| `--model ` | String | Override Gemini model name (default: `gemini-2.5-flash`). | +| `--project ` | String | Override Google Cloud Project ID. | +| `--customer-id ` | String | Override Chronicle Customer ID. | + +### `mcp-security-agent serve [OPTIONS]` +Launches the FastAPI application for web UI access or Cloud Run hosting. + +| Option | Default | Description | +| :--- | :--- | :--- | +| `--host` | `0.0.0.0` | Network interface to bind to. | +| `--port` | `8080` | Port to listen on (reads `$PORT` environment variable if set). | +| `--reload` | `False` | Enable auto-reload for local development. | +| Tool flags | | Supports the same tool toggles as `chat` (`--secops`, `--scc`, `--gti`, `--soar`, etc.). | +--- -# Add Your MCP server variables here, sample provided -# 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 +## Configuration & Environment Variables +The agent reads configuration from environment variables and an optional `.env` file located in the working directory. +### `.env.example` Template +```properties +# Google Cloud & LLM Settings +GOOGLE_CLOUD_PROJECT= +GOOGLE_CLOUD_LOCATION=us-central1 +GOOGLE_GENAI_USE_VERTEXAI=False +GOOGLE_API_KEY= +GOOGLE_MODEL=gemini-2.5-flash +# MCP Server Enablement Flags +# (SCC uses Google Cloud ADC. SecOps, GTI, and SOAR auto-enable when credentials below are populated. +# You can explicitly set LOAD_*_MCP to Y or N to override auto-detection.) +LOAD_SCC_MCP=Y +# LOAD_SECOPS_MCP= +# LOAD_GTI_MCP= +# LOAD_SECOPS_SOAR_MCP= -``` +# Credentials & Service Account Impersonation +SECOPS_SA_PATH= +GOOGLE_APPLICATION_CREDENTIALS= +SECOPS_IMPERSONATE_SERVICE_ACCOUNT= -Once the variables are updated in `.env`, run the agent (make sure you are in the `mcp-security/run-with-google-adk` directory). +# Google SecOps (Chronicle SIEM) Settings - Populating enables SecOps MCP +CHRONICLE_PROJECT_ID= +CHRONICLE_CUSTOMER_ID= +CHRONICLE_REGION=us -```bash -# Authenticate to use Google Cloud / SecOps APIs -# Skip if running in Google Cloud Shell -gcloud auth application-default login +# Google Threat Intelligence (GTI / VirusTotal) - Populating enables GTI MCP +VT_APIKEY= -# Start interactive terminal chat -uv run mcp-security-agent chat +# SecOps SOAR Settings - Populating enables SOAR MCP +SOAR_URL= +SOAR_APP_KEY= -# Or start the ADK Web interface -adk web src/mcp_security_agent +# Runtime Settings +STDIO_PARAM_TIMEOUT=60.0 +MINIMAL_LOGGING=N ``` -Access the agent interface by navigating to `http://localhost:8000`. +### Environment Variable Reference + +| Variable | Default | Description | +| :--- | :--- | :--- | +| `GOOGLE_CLOUD_PROJECT` | None | Google Cloud Project ID for Vertex AI and Cloud Run deployment. | +| `GOOGLE_CLOUD_LOCATION` | `us-central1` | Google Cloud region for Vertex AI endpoints. | +| `GOOGLE_GENAI_USE_VERTEXAI`| `False` | Set `True` to route LLM requests through Vertex AI (uses ADC). Set `False` to use Gemini API (requires `GOOGLE_API_KEY`). | +| `GOOGLE_API_KEY` | None | Gemini API Key (required when `GOOGLE_GENAI_USE_VERTEXAI=False`). | +| `GOOGLE_MODEL` | `gemini-2.5-flash` | Gemini model name (e.g. `gemini-2.5-flash`, `gemini-2.5-pro`). | +| `LOAD_SCC_MCP` | `False` | Enables Security Command Center MCP (`server/scc`). Uses ADC. | +| `LOAD_SECOPS_MCP` | Auto | Enables Chronicle SIEM MCP. Auto-enables if `CHRONICLE_PROJECT_ID` and `CUSTOMER_ID` are set. | +| `LOAD_GTI_MCP` | Auto | Enables Google Threat Intelligence MCP. Auto-enables if `VT_APIKEY` is set. | +| `LOAD_SECOPS_SOAR_MCP` | Auto | Enables SecOps SOAR MCP. Auto-enables if `SOAR_URL` and `SOAR_APP_KEY` are set. | +| `CHRONICLE_PROJECT_ID` | None | GCP project ID hosting Chronicle SIEM. | +| `CHRONICLE_CUSTOMER_ID` | None | Chronicle Customer ID (UUID). | +| `CHRONICLE_REGION` | `us` | Chronicle regional gateway (`us`, `europe`, `asia`). | +| `VT_APIKEY` | None | VirusTotal / GTI API Key (presence auto-activates GTI tools). | +| `SOAR_URL` | None | Instance URL for SecOps SOAR (presence auto-activates SOAR tools). | +| `SOAR_APP_KEY` | None | API Key for SecOps SOAR. | +| `GOOGLE_APPLICATION_CREDENTIALS` | Auto-detected | Path to service account key file or local `.gcloud/application_default_credentials.json`. | +| `STDIO_PARAM_TIMEOUT` | `60.0` | Timeout in seconds for MCP subprocess initialization and tool execution. | +| `MINIMAL_LOGGING` | `False` | Reduces logging verbosity to suppress sensitive query content in production. | -> **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 Custom Session and Artifact Services +--- -Google ADK provides persistent [sessions](https://google.github.io/adk-docs/sessions/) and [artifacts](https://google.github.io/adk-docs/artifacts/). +## Running the Web UI & API Server -You can run the agent with the session and artifact service of your choice: +The built-in FastAPI server provides an interactive web dashboard and REST / Server-Sent Events (SSE) API endpoints. ```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 +uv run mcp-security-agent serve --port 8080 ``` -When the artifact service is backed by GCS, signed URLs allow easy file sharing. Grant the runtime service account the `roles/storage.objectViewer` role. - +Open `http://localhost:8080` in your browser to access the SOC Agent UI. -## 2. Running Agent as a Cloud Run Service +### REST & Streaming Endpoints -The agent with MCP servers can be deployed as a Cloud Run Service, right from within the source code directory. +* **`GET /`**: Serves the bundled web landing page and interactive investigation console. +* **`GET /healthz`**: Liveness & readiness probe returning `{"status": "ok"}` for Cloud Run. +* **`GET /info`**: Returns JSON metadata including package version, active model, and toolset configurations. +* **`POST /chat`**: Synchronous chat endpoint accepting `{"prompt": "string", "session_id": "optional"}`. +* **`GET /chat?message=...`**: Server-Sent Events (SSE) token streaming endpoint. -Before you do this, please consider following - -1. Do you really need it? Deployment is recommended in scenarios where you need to share agent with your team members who may not have access to all of the backend services (SCC, SecOps - SIEM, SecOps - SOAR, Google Threat Intelligence) -2. Make sure that after initial testing - 1. Require authentication for your agent (steps provided [below](#restrict-service-to-known-developers--testers)) - 2. Implement restrictive logging (steps provided [below](#adjust-logging-verbosity)) - -### Prerequisites - -1. Must have locally run the ADK based agent successfully at least once. Environment variables `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` should have valid values. -2. Must have required APIs enabled and proper IAM access ([details](https://cloud.google.com/run/docs/deploying-source-code#before_you_begin)) +--- -### Costs -In addition to Gemini/ Vertex API costs, running agent will incur cloud costs. Please check [Cloud Run Pricing](https://cloud.google.com/run/pricing). +## Deploying to Google Cloud Run -> ⚠️ **WARNING:** -> 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. +The package includes a production-ready [`Dockerfile`](./Dockerfile) configured to run `mcp-security-agent serve` on Cloud Run. -### Deployment Steps +### 1. Build and Deploy -> **NOTE:** -> It is recommended to switch to Vertex AI (with `GOOGLE_GENAI_USE_VERTEXAI=True`) when deploying to Cloud Run. +Deploy the service directly from the repository root: ```bash -# 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 URL. - -### IAM access to use Chronicle and SCC - -Please remember that Cloud Run uses default service account of compute engine service. Go to IAM and provide the service account access to "Chronicle API Viewer" (in the project associated with your SecOps instance) and appropriate role for SCC (roles starting with Security Center in IAM) - - -### Restrict Service To Known Developers / Testers - -Summarizing the steps from [IAM authentication](https://cloud.google.com/run/docs/authenticating/developers) - -1. Goto Cloud Run - Services - click `mcp-security-agent-service` -2. Click `Security` -3. In `Authentication`, `Use Cloud IAM to authenticate incoming requests` should be already selected. -4. Select the radio button `Require authentication` -5. Click `Save` -6. Cloud Run - Services - select `mcp-security-agent-service` -7. At the top click `permissions`, a pane `Permissions for mcp-security-agent-service` should open on the right hand side. -8. Click `Add principal` -9. Add the users you want to provide access to and provide them `Cloud Run Invoker` role. -10. Wait for some time. - -### Accessing the restricted service - -1. Ask your users to run the following command (replace project id and region with the project id & region in which you have deployed the service) - -```bash -gcloud run services proxy mcp-security-agent-service --project PROJECT-ID --region YOUR-REGION - + --set-env-vars="LOAD_SECOPS_MCP=Y,LOAD_SCC_MCP=Y,GOOGLE_GENAI_USE_VERTEXAI=True,GOOGLE_CLOUD_PROJECT=YOUR_PROJECT_ID,CHRONICLE_PROJECT_ID=YOUR_PROJECT_ID,CHRONICLE_CUSTOMER_ID=YOUR_CUSTOMER_ID" ``` -2. Now they can access the Cloud Run Service locally on `http://localhost:8080` - - -### Vertically scaling your container(s) -In case the Cloud Run logs show errors like below, you can consider increasing the resources for the individual containers - -`Memory limit of 512 MiB exceeded with 543 MiB used. Consider increasing the memory limit, see https://cloud.google.com/run/docs/configuring/memory-limits` - -##### Steps - -1. Goto Cloud Run - Services - click `mcp-security-agent-service` -2. Click `Edit & deploy new revision` -3. In `Container(s)` - `Edit Container(s)` - `Settings` -4. Add resources by updating either Memory/ CPU or both. - -### Adjust Logging Verbosity -Since the entire context and response from the LLM is printed as logs. You might end up logging some sensitive information. Setting the environment variable `MINIMAL_LOGGING` to `Y` should fix this issue. This should also reduce cloud logging costs. Please do this once you have verified the service initially. Changes to be made directly on Cloud Run service and it will result in restarting the service. Verify service logs after the change is made. -## 3. Deploying and Running Agent on Agent Engine +### 2. IAM Roles -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. +Ensure the runtime Service Account used by Cloud Run has the appropriate IAM permissions: +* **Chronicle API Viewer** (`roles/chronicle.viewer`) on the project hosting Chronicle. +* **Security Center Finding Viewer** (`roles/securitycenter.findingsViewer`) for SCC findings. +* **Vertex AI User** (`roles/aiplatform.user`) for Vertex AI model execution. -Here are the deployment steps: +### 3. Restricting Access in Production -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). - -### How to Test - -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). - ---- - -## 4. Improving Performance and Optimizing Costs - -By default, the agent sends the active conversation context to the LLM. - -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 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. - ---- +To protect the service in production: +1. In Cloud Run, select `mcp-security-agent-service` and navigate to **Security**. +2. Select **Require authentication**. +3. Under **Permissions**, grant `Cloud Run Invoker` (`roles/run.invoker`) to authorized users. +4. Authorized users can securely access the service locally via proxy: + ```bash + gcloud run services proxy mcp-security-agent-service --project YOUR_PROJECT_ID --region us-central1 + ``` + The service will then be reachable at `http://localhost:8080`. -## 5. Integrating Custom MCP Servers +### 4. Memory Limits & Logging Verbosity -If your organization uses additional security products (such as identity providers or third-party EDRs), integrating them with Google Security MCP servers provides: +* If MCP toolsets process large result sets, consider increasing container memory in Cloud Run Settings to 1 GiB or 2 GiB. +* In production, set `MINIMAL_LOGGING=Y` to suppress full LLM prompt/response logging and reduce Cloud Logging costs. -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. +## Integrating Custom MCP Servers -### Reference Integration Templates +You can integrate custom security products (such as internal IdPs, custom EDRs, or ticketing systems) using modular sub-agents. -Reference templates are provided in `run-with-google-adk/sample_servers_to_integrate/`: +Reference implementations are provided in [`sample_servers_to_integrate/`](./sample_servers_to_integrate/): +* **Sample MCP Servers**: [`sample_servers_to_integrate/mcp_servers/`](./sample_servers_to_integrate/mcp_servers/) (`demo_idp` and `demo_xdr`). +* **Sample Sub-Agents**: [`sample_servers_to_integrate/agents/`](./sample_servers_to_integrate/agents/) (`demo_idp_agent.py` and `demo_xdr_agent.py`). -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`: +To attach them to the primary agent in [`src/mcp_security_agent/agent.py`](./src/mcp_security_agent/agent.py): ```python -# 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, @@ -334,8 +266,7 @@ agent = LlmAgent( ) ``` -Configure corresponding environment variables in `.env`: - +Configure credentials in `.env`: ```properties LOAD_XDR_MCP=Y XDR_CLIENT_ID=demo_client_id @@ -346,119 +277,25 @@ IDP_CLIENT_ID=demo_client_id IDP_CLIENT_SECRET=demo_client_secret ``` -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:** -> Once tested, you can attach production MCP servers following this modular pattern. - -Reference architecture screenshots: +Sample reference interfaces: -Sample XDR: +**Sample XDR Integration:** ![](./static/demo-xdr.png) -Sample IDP: +**Sample IDP Integration:** ![](./static/demo-idp.png) +--- -## 6. Additional Features - -The prebuilt agent also allows creating files and signed URLs to these files. A possible scenario is when you want to create a report. You can say "add the summary as markdown to summary_146.md". This creates a file and saves it using the artifact service. You can later ask for a shareable link to this file - "create a link to file summary_146.md" - -## 7. Registering Agent Engine Agent to AgentSpace - -1. When an agent is deployed on Agent Engine ([guide](#3-deploying-and-running-agent-on-agent-engine)) you get a resource name. Make sure you have it to carry out next steps -2. Go to the Agentspace [page](https://console.cloud.google.com/gen-app-builder/engines) in Google Cloud Console. -3. Create an App (Type - AgentSpace) -4. Note down the app details including the app name (e.g. google-security-agent-app_1750057151234) -5. Make sure that you have the Agent Space Admin role while performing the following actions -6. Enable Discovery Engine API for your project -7. Provide the following roles to the Discovery Engine Service Account - Vertex AI viewer - Vertex AI user -8. Please note that these roles need to be provided into the project housing your Agent Engine Agent. Also you need to enable the show Google provided role grants to access the Discovery Engine Service Account. -9. Now to register the agent and make it available to your application use the following shell script. Please replace the variables `AGENT_SPACE_PROJECT_ID ,AGENT_SPACE_APP_NAME ,AGENT_ENGINE_PROJECT_NUMBER , AGENT_LOCATION` and `REASONING_ENGINE_NUMBER` before running the script. - -```bash -#!/bin/bash - -TARGET_URL="https://discoveryengine.googleapis.com/v1alpha/projects/AGENT_SPACE_PROJECT_ID/locations/global/collections/default_collection/engines/AGENT_SPACE_APP_NAME/assistants/default_assistant/agents" # - -JSON_DATA=$(cat < Any: return agent -# Expose root_agent for standard ADK CLI discovery (adk run, adk web) -root_agent = create_security_agent() +# Lazy root_agent instantiation for standard ADK CLI discovery (adk run, adk web) +_root_agent: Optional[Any] = None +_root_agent_lock = threading.Lock() + + +def __getattr__(name: str) -> Any: + global _root_agent + if name == "root_agent": + if _root_agent is None: + with _root_agent_lock: + if _root_agent is None: + _root_agent = create_security_agent() + return _root_agent + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return ["create_security_agent", "root_agent", "SOC_AGENT_SYSTEM_PROMPT"] + diff --git a/run-with-google-adk/src/mcp_security_agent/cli.py b/run-with-google-adk/src/mcp_security_agent/cli.py index 1e1bcbf0..5f63e180 100644 --- a/run-with-google-adk/src/mcp_security_agent/cli.py +++ b/run-with-google-adk/src/mcp_security_agent/cli.py @@ -14,6 +14,8 @@ """Command-line interface for the MCP Security Agent.""" import asyncio +import os +import sys from pathlib import Path from typing import Optional import typer @@ -34,17 +36,102 @@ def info(): 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"Project: [cyan]{settings.google_cloud_project or os.environ.get('GOOGLE_CLOUD_PROJECT', 'Not set')}[/cyan]") + console.print(f"ADC: [cyan]{settings.google_application_credentials or 'Default (~/.config/gcloud/...)'}[/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]'}") + +def _apply_cli_overrides( + secops: Optional[bool] = None, + scc: Optional[bool] = None, + gti: Optional[bool] = None, + soar: Optional[bool] = None, + vertex: Optional[bool] = None, + model: Optional[str] = None, + project: Optional[str] = None, + customer_id: Optional[str] = None, +) -> None: + """Applies CLI flag overrides to environment variables before settings initialization.""" + if secops is not None: + os.environ["LOAD_SECOPS_MCP"] = "Y" if secops else "N" + if scc is not None: + os.environ["LOAD_SCC_MCP"] = "Y" if scc else "N" + if gti is not None: + os.environ["LOAD_GTI_MCP"] = "Y" if gti else "N" + if soar is not None: + os.environ["LOAD_SECOPS_SOAR_MCP"] = "Y" if soar else "N" + if vertex is not None: + os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "TRUE" if vertex else "FALSE" + if model: + os.environ["GOOGLE_MODEL"] = model + if project: + os.environ["GOOGLE_CLOUD_PROJECT"] = project + if "CHRONICLE_PROJECT_ID" not in os.environ: + os.environ["CHRONICLE_PROJECT_ID"] = project + if customer_id: + os.environ["CHRONICLE_CUSTOMER_ID"] = customer_id + + +def _display_active_toolsets(settings: AgentSettings) -> None: + enabled_tools = [] + if settings.load_secops_mcp: + enabled_tools.append("SecOps SIEM") + if settings.load_scc_mcp: + enabled_tools.append("SCC") + if settings.load_gti_mcp: + enabled_tools.append("GTI") + if settings.load_secops_soar_mcp: + enabled_tools.append("SecOps SOAR") + + if not enabled_tools: + console.print( + "[bold yellow]Warning:[/bold yellow] No MCP tools are currently enabled.\n" + "Enable tools using environment variables (e.g. [cyan]LOAD_SECOPS_MCP=Y[/cyan]), " + "CLI flags ([cyan]--secops[/cyan], [cyan]--scc[/cyan]), or [cyan].env[/cyan].\n" + ) + else: + console.print(f"[bold green]Active MCP Toolsets:[/bold green] {', '.join(enabled_tools)}\n") + + @app.command() def chat( query: Optional[str] = typer.Argument(None, help="Optional single-turn investigation query to execute"), + secops: Optional[bool] = typer.Option(None, "--secops/--no-secops", help="Enable or disable SecOps SIEM MCP"), + scc: Optional[bool] = typer.Option(None, "--scc/--no-scc", help="Enable or disable SCC MCP"), + gti: Optional[bool] = typer.Option(None, "--gti/--no-gti", help="Enable or disable GTI MCP"), + soar: Optional[bool] = typer.Option(None, "--soar/--no-soar", help="Enable or disable SecOps SOAR MCP"), + vertex: Optional[bool] = typer.Option(None, "--vertex/--no-vertex", help="Use Vertex AI for LLM requests"), + model: Optional[str] = typer.Option(None, "--model", help="Gemini model to use (e.g. gemini-2.5-flash)"), + project: Optional[str] = typer.Option(None, "--project", help="Google Cloud project ID"), + customer_id: Optional[str] = typer.Option(None, "--customer-id", help="Chronicle Customer ID (UUID)"), ): """Start an interactive terminal chat session with the SOC agent powered by ADK v2.""" + _apply_cli_overrides( + secops=secops, + scc=scc, + gti=gti, + soar=soar, + vertex=vertex, + model=model, + project=project, + customer_id=customer_id, + ) + + settings = AgentSettings() + _display_active_toolsets(settings) + + import mcp_security_agent.agent as agent_mod + agent = agent_mod.create_security_agent(settings) + agent_mod._root_agent = agent + agent_mod.root_agent = agent + if "mcp_security_agent" in sys.modules: + sys.modules["mcp_security_agent"]._root_agent = agent + sys.modules["mcp_security_agent"].root_agent = agent + try: from google.adk.cli.cli import run_cli, run_once_cli except (ImportError, ModuleNotFoundError): @@ -78,15 +165,50 @@ def chat( @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"), + port: Optional[int] = typer.Option(None, help="Port to listen on (defaults to $PORT or 8080)"), + reload: bool = typer.Option(False, "--reload", help="Enable auto-reload for development"), + secops: Optional[bool] = typer.Option(None, "--secops/--no-secops", help="Enable or disable SecOps SIEM MCP"), + scc: Optional[bool] = typer.Option(None, "--scc/--no-scc", help="Enable or disable SCC MCP"), + gti: Optional[bool] = typer.Option(None, "--gti/--no-gti", help="Enable or disable GTI MCP"), + soar: Optional[bool] = typer.Option(None, "--soar/--no-soar", help="Enable or disable SecOps SOAR MCP"), + vertex: Optional[bool] = typer.Option(None, "--vertex/--no-vertex", help="Use Vertex AI for LLM requests"), + model: Optional[str] = typer.Option(None, "--model", help="Gemini model to use (e.g. gemini-2.5-flash)"), + project: Optional[str] = typer.Option(None, "--project", help="Google Cloud project ID"), + customer_id: Optional[str] = typer.Option(None, "--customer-id", help="Chronicle Customer ID (UUID)"), ): """Run the FastAPI web server and Cloud Run REST API.""" - import uvicorn - from mcp_security_agent.server.app import create_app + _apply_cli_overrides( + secops=secops, + scc=scc, + gti=gti, + soar=soar, + vertex=vertex, + model=model, + project=project, + customer_id=customer_id, + ) + + settings = AgentSettings() + _display_active_toolsets(settings) - 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) + import mcp_security_agent.agent as agent_mod + agent = agent_mod.create_security_agent(settings) + agent_mod._root_agent = agent + agent_mod.root_agent = agent + if "mcp_security_agent" in sys.modules: + sys.modules["mcp_security_agent"]._root_agent = agent + sys.modules["mcp_security_agent"].root_agent = agent + + bind_port = port if port is not None else int(os.environ.get("PORT", 8080)) + console.print(f"[bold green]Starting MCP Security Agent server on {host}:{bind_port}[/bold green]") + + import uvicorn + if reload: + uvicorn.run("mcp_security_agent.server.app:create_app", factory=True, host=host, port=bind_port, reload=True) + else: + from mcp_security_agent.server.app import create_app + app_instance = create_app() + uvicorn.run(app_instance, host=host, port=bind_port) if __name__ == "__main__": diff --git a/run-with-google-adk/src/mcp_security_agent/config.py b/run-with-google-adk/src/mcp_security_agent/config.py index 297582b4..2a6edc22 100644 --- a/run-with-google-adk/src/mcp_security_agent/config.py +++ b/run-with-google-adk/src/mcp_security_agent/config.py @@ -13,15 +13,132 @@ # limitations under the License. """Centralized configuration and settings for MCP Security Agent.""" +import functools +from pathlib import Path from typing import Optional -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict +_pkg_dir = Path(__file__).resolve().parents[2] +_env_files = ( + str(_pkg_dir.parent / ".env"), + str(_pkg_dir / ".env"), + ".env", +) + + +def _is_configured(val: Optional[str]) -> bool: + """Checks if a configuration string is set and non-empty.""" + return bool(val and val.strip()) + + +def _discover_local_adc() -> Optional[str]: + """Finds local .gcloud/application_default_credentials.json if present.""" + candidates = [ + Path.cwd() / ".gcloud" / "application_default_credentials.json", + _pkg_dir / ".gcloud" / "application_default_credentials.json", + _pkg_dir.parent / ".gcloud" / "application_default_credentials.json", + ] + for c in candidates: + if c.is_file(): + return str(c) + return None + + +@functools.lru_cache(maxsize=1) +def discover_user_identity() -> str: + """Discovers the active user identity from ADC, gcloud config, or system environment. + + Returns: + The detected username or service account email, falling back to 'secops_user'. + """ + import os + import json + import getpass + import shutil + import subprocess + + # 1. Explicit user override from environment + explicit = os.getenv("SECOPS_USER") or os.getenv("AGENT_USER") + if explicit and explicit.strip(): + return explicit.strip() + + # 2. Impersonated service account + impersonate_sa = os.getenv("SECOPS_IMPERSONATE_SERVICE_ACCOUNT") + if impersonate_sa and impersonate_sa.strip(): + return impersonate_sa.strip() + + # 3. Discovered ADC file (check for account or client_email) + adc_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") or _discover_local_adc() + if not adc_path: + cloudsdk_config = os.getenv("CLOUDSDK_CONFIG") + if cloudsdk_config: + p = Path(cloudsdk_config) / "application_default_credentials.json" + if p.is_file(): + adc_path = str(p) + if not adc_path: + p = Path.home() / ".config" / "gcloud" / "application_default_credentials.json" + if p.is_file(): + adc_path = str(p) + + if adc_path and Path(adc_path).is_file(): + try: + with open(adc_path, "r", encoding="utf-8") as f: + data = json.load(f) + acct = data.get("account") or data.get("client_email") + if acct and isinstance(acct, str) and acct.strip(): + return acct.strip() + except Exception: + pass + + # 4. google.auth.default() credentials inspection + try: + import google.auth + creds, _ = google.auth.default() + sa_email = getattr(creds, "service_account_email", None) + if sa_email and isinstance(sa_email, str) and sa_email.strip() and sa_email != "default": + return sa_email.strip() + acct = getattr(creds, "account", None) + if acct and isinstance(acct, str) and acct.strip(): + return acct.strip() + except Exception: + pass + + # 5. gcloud active account + if shutil.which("gcloud"): + try: + res = subprocess.run( + ["gcloud", "config", "get-value", "account"], + capture_output=True, + text=True, + timeout=2, + ) + if res.returncode == 0: + acct = res.stdout.strip() + if acct and acct != "(unset)": + return acct + except Exception: + pass + + # 6. System username (LDAP / OS user) + system_user = os.getenv("USER") or os.getenv("USERNAME") + if not system_user: + try: + system_user = getpass.getuser() + except Exception: + system_user = None + if system_user and system_user.strip(): + return system_user.strip() + + return "secops_user" + + + class AgentSettings(BaseSettings): """Configuration settings loaded from environment variables or .env file.""" model_config = SettingsConfigDict( - env_file=".env", + env_file=_env_files, env_file_encoding="utf-8", extra="ignore", populate_by_name=True, @@ -34,11 +151,11 @@ class AgentSettings(BaseSettings): 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") + # MCP Server Enablement Flags (Optional; auto-detected from credentials if None) + load_secops_mcp: Optional[bool] = Field(default=None, 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") + load_gti_mcp: Optional[bool] = Field(default=None, alias="LOAD_GTI_MCP") + load_secops_soar_mcp: Optional[bool] = Field(default=None, 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") @@ -48,8 +165,54 @@ class AgentSettings(BaseSettings): # 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") + google_application_credentials: Optional[str] = Field( + default_factory=lambda: _discover_local_adc(), alias="GOOGLE_APPLICATION_CREDENTIALS" + ) secops_impersonate_service_account: Optional[str] = Field(default=None, alias="SECOPS_IMPERSONATE_SERVICE_ACCOUNT") + + def __init__(self, **values): + super().__init__(**values) + self.bootstrap_environment() + + @model_validator(mode="after") + def resolve_tool_enablement(self) -> "AgentSettings": + """Auto-detects MCP tool enablement based on presence of API keys and credentials.""" + if self.load_gti_mcp is None: + self.load_gti_mcp = _is_configured(self.vt_apikey) + + if self.load_secops_soar_mcp is None: + self.load_secops_soar_mcp = _is_configured(self.soar_url) and _is_configured(self.soar_app_key) + + if self.load_secops_mcp is None: + self.load_secops_mcp = _is_configured(self.chronicle_project_id) and _is_configured(self.chronicle_customer_id) + + return self + + def bootstrap_environment(self) -> None: + """Configures environment variables for Google Cloud authentication and Cloudtop compatibility.""" + import os + if self.google_application_credentials and not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"): + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.google_application_credentials + if not os.environ.get("CLOUDSDK_CONFIG"): + os.environ["CLOUDSDK_CONFIG"] = str(Path(self.google_application_credentials).parent) + + if "GOOGLE_API_USE_CLIENT_CERTIFICATE" not in os.environ: + os.environ["GOOGLE_API_USE_CLIENT_CERTIFICATE"] = "false" + if "GOOGLE_API_USE_MTLS_ENDPOINT" not in os.environ: + os.environ["GOOGLE_API_USE_MTLS_ENDPOINT"] = "never" + if "CLOUDSDK_CONTEXT_AWARE_USE_CLIENT_CERTIFICATE" not in os.environ: + os.environ["CLOUDSDK_CONTEXT_AWARE_USE_CLIENT_CERTIFICATE"] = "false" + + if self.google_api_key and not os.environ.get("GOOGLE_API_KEY"): + os.environ["GOOGLE_API_KEY"] = self.google_api_key + + target_project = self.google_cloud_project or os.environ.get("GCP_PROJECT_ID") + if target_project and not os.environ.get("GOOGLE_CLOUD_PROJECT"): + os.environ["GOOGLE_CLOUD_PROJECT"] = target_project + + if self.use_vertex_ai or (target_project and not self.google_api_key and not os.environ.get("GOOGLE_API_KEY")): + os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "TRUE" + # Chronicle SIEM Params chronicle_project_id: Optional[str] = Field(default=None, alias="CHRONICLE_PROJECT_ID") @@ -67,12 +230,31 @@ class AgentSettings(BaseSettings): 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", + "load_secops_mcp", "load_gti_mcp", "load_secops_soar_mcp", + mode="before" + ) + @classmethod + def parse_optional_bool_env(cls, value: object) -> Optional[bool]: + if value is None: + return None + if isinstance(value, str): + val_clean = value.strip().upper() + if not val_clean: + return None + return val_clean in ("Y", "YES", "TRUE", "1") + return bool(value) + + @field_validator( + "load_scc_mcp", "use_vertex_ai", "minimal_logging", mode="before" ) @classmethod def parse_bool_env(cls, value: object) -> bool: + if value is None: + return False if isinstance(value, str): - return value.strip().upper() in ("Y", "YES", "TRUE", "1") + val_clean = value.strip().upper() + if not val_clean: + return False + return val_clean in ("Y", "YES", "TRUE", "1") return bool(value) 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 index 276ce304..3df56f58 100644 --- a/run-with-google-adk/src/mcp_security_agent/server/app.py +++ b/run-with-google-adk/src/mcp_security_agent/server/app.py @@ -39,4 +39,15 @@ def create_app() -> FastAPI: if static_dir.is_dir(): app.mount("/static", StaticFiles(directory=str(static_dir)), name="static") + @app.middleware("http") + async def add_no_cache_headers(request, call_next): + response = await call_next(request) + path = request.url.path + if path.startswith("/static") or path in ["/", "/login", "/landing.html", "/index.html", "/chat.html"]: + response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate" + response.headers["Pragma"] = "no-cache" + response.headers["Expires"] = "0" + return response + 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 index d8839e9e..85f09adf 100644 --- a/run-with-google-adk/src/mcp_security_agent/server/routes.py +++ b/run-with-google-adk/src/mcp_security_agent/server/routes.py @@ -16,20 +16,31 @@ import json import uuid import asyncio +import logging +import threading +from collections import OrderedDict from pathlib import Path from typing import Dict, Any, Optional, AsyncGenerator -from fastapi import APIRouter, HTTPException, Query +from fastapi import APIRouter, HTTPException, Query, Request, Response from fastapi.responses import FileResponse, StreamingResponse, JSONResponse from pydantic import BaseModel +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.genai import types +import mcp_security_agent.agent as agent_mod from mcp_security_agent import __version__ -from mcp_security_agent.config import AgentSettings +from mcp_security_agent.config import AgentSettings, discover_user_identity + +logger = logging.getLogger(__name__) router = APIRouter() class ChatRequest(BaseModel): - prompt: str + prompt: Optional[str] = None + message: Optional[str] = None session_id: Optional[str] = None + user_id: Optional[str] = None class ChatResponse(BaseModel): @@ -38,25 +49,32 @@ class ChatResponse(BaseModel): @router.get("/") +@router.get("/index.html") +@router.get("/landing.html") +@router.get("/chat.html") +@router.get("/login") def get_root(): - """Serves the main landing page of the web UI.""" + """Serves the unified investigation console 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(): + if 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("/favicon.ico") +def get_favicon(): + """Returns 204 No Content for browser favicon requests.""" + return Response(status_code=204) + + @router.get("/app_name") def get_app_name() -> Dict[str, str]: """Returns the application display name for the Web UI navbar.""" @@ -66,19 +84,36 @@ def get_app_name() -> Dict[str, str]: @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.""" + detected_user = discover_user_identity() + user = username if (username and username.strip() and username != "secops_user") else detected_user return { "session_id": str(uuid.uuid4()), - "user_id": username or "default_user", + "user_id": user, } +_settings: Optional[AgentSettings] = None +_settings_lock = threading.Lock() + + +def get_settings() -> AgentSettings: + """Returns cached AgentSettings singleton to avoid redundant env parsing.""" + global _settings + if _settings is None: + with _settings_lock: + if _settings is None: + _settings = AgentSettings() + return _settings + + @router.get("/info") def get_info() -> Dict[str, Any]: - """Provides server runtime metadata and enabled MCP server status.""" - settings = AgentSettings() + """Provides server runtime metadata, active user identity, and enabled MCP server status.""" + settings = get_settings() return { "version": __version__, "model": settings.google_model, + "user": discover_user_identity(), "tools": { "secops": settings.load_secops_mcp, "scc": settings.load_scc_mcp, @@ -88,15 +123,229 @@ def get_info() -> Dict[str, Any]: } -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}) +class BoundedSessionService(InMemorySessionService): + """InMemorySessionService with FIFO session eviction to prevent unbounded memory growth.""" + + def __init__(self, max_sessions: int = 1000): + super().__init__() + self.max_sessions = max_sessions + self._session_order: OrderedDict = OrderedDict() + self._lock = threading.RLock() + + def _create_session_impl( + self, + *, + app_name: str, + user_id: str, + state: Optional[Dict[str, Any]] = None, + session_id: Optional[str] = None, + ) -> Any: + with self._lock: + while len(self._session_order) >= self.max_sessions: + (old_app, old_user, old_sess), _ = self._session_order.popitem(last=False) + if old_app in self.sessions and old_user in self.sessions[old_app]: + self.sessions[old_app][old_user].pop(old_sess, None) + if not self.sessions[old_app][old_user]: + self.sessions[old_app].pop(old_user, None) + if old_app in self.sessions and not self.sessions[old_app]: + self.sessions.pop(old_app, None) + sess = super()._create_session_impl( + app_name=app_name, + user_id=user_id, + state=state, + session_id=session_id, + ) + self._session_order[(app_name, user_id, sess.id)] = True + return sess + + def _get_session_impl( + self, + *, + app_name: str, + user_id: str, + session_id: str, + config: Optional[Any] = None, + ) -> Optional[Any]: + with self._lock: + return super()._get_session_impl( + app_name=app_name, + user_id=user_id, + session_id=session_id, + config=config, + ) + + def _list_sessions_impl( + self, + *, + app_name: str, + user_id: Optional[str] = None, + ) -> Any: + with self._lock: + return super()._list_sessions_impl( + app_name=app_name, + user_id=user_id, + ) + + async def append_event(self, session: Any, event: Any) -> Any: + with self._lock: + app_name = getattr(session, "app_name", None) + user_id = getattr(session, "user_id", None) + session_id = getattr(session, "id", None) + if ( + not app_name + or not user_id + or not session_id + or app_name not in self.sessions + or user_id not in self.sessions[app_name] + or session_id not in self.sessions[app_name][user_id] + ): + return event + return await super().append_event(session=session, event=event) + + def _delete_session_impl( + self, + *, + app_name: str, + user_id: str, + session_id: str, + ) -> None: + with self._lock: + self._session_order.pop((app_name, user_id, session_id), None) + super()._delete_session_impl( + app_name=app_name, + user_id=user_id, + session_id=session_id, + ) + if app_name in self.sessions and user_id in self.sessions[app_name]: + if not self.sessions[app_name][user_id]: + self.sessions[app_name].pop(user_id, None) + if app_name in self.sessions and not self.sessions[app_name]: + self.sessions.pop(app_name, None) + + +_runner: Optional[Runner] = None +_runner_lock = threading.Lock() +_session_service: Optional[BoundedSessionService] = None + + +def get_session_service() -> BoundedSessionService: + """Returns the singleton BoundedSessionService for active chat sessions.""" + global _session_service + if _session_service is None: + with _runner_lock: + if _session_service is None: + _session_service = BoundedSessionService(max_sessions=1000) + return _session_service + + +def get_runner() -> Optional[Runner]: + """Returns or initializes the thread-safe ADK Runner instance for the security agent.""" + global _runner + if _runner is None: + with _runner_lock: + if _runner is None: + settings = get_settings() + agent = getattr(agent_mod, "_root_agent", None) + if agent is None: + agent = agent_mod.create_security_agent(settings) + agent_mod._root_agent = agent + if agent is None: + return None + session_svc = get_session_service() + _runner = Runner( + agent=agent, + app_name="mcp_security_agent", + session_service=session_svc, + ) + return _runner + + +async def sse_event_generator( + message: str, + session_id: str, + user_id: Optional[str] = None, +) -> AsyncGenerator[str, None]: + """Streams real ADK agent execution events over Server-Sent Events (SSE).""" + uid = user_id or discover_user_identity() + session_svc = get_session_service() + runner = get_runner() + + if runner is None: + err_msg = json.dumps({ + "text": "[Warning] Agent runner is unavailable in this environment.", + "last_msg": False, + "session_id": session_id, + }) + yield f"data: {err_msg}\n\n" + done_msg = json.dumps({"text": "Stream finished.", "last_msg": True, "session_id": session_id}) + yield f"data: {done_msg}\n\n" + return + + session = await session_svc.get_session( + app_name="mcp_security_agent", + user_id=uid, + session_id=session_id, + ) + if not session: + session = await session_svc.create_session( + app_name="mcp_security_agent", + user_id=uid, + session_id=session_id, + ) + + content = types.Content(role="user", parts=[types.Part(text=message)]) + + try: + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=content, + ): + if event.content and event.content.parts: + for part in event.content.parts: + if part.text: + data = json.dumps({ + "text": part.text, + "last_msg": False, + "session_id": session_id, + "author": event.author or "SecurityOperationsAgent", + "event_type": "content", + }) + yield f"data: {data}\n\n" + elif part.function_call: + call_info = json.dumps({ + "text": f"[Tool] **Calling tool `{part.function_call.name}`**\n```json\n{json.dumps(part.function_call.args, indent=2)}\n```", + "last_msg": False, + "session_id": session_id, + "event_type": "tool_call", + }) + yield f"data: {call_info}\n\n" + elif part.function_response: + resp_info = json.dumps({ + "text": f"[Tool] **Received tool response from `{part.function_response.name}`**", + "last_msg": False, + "session_id": session_id, + "event_type": "tool_response", + }) + yield f"data: {resp_info}\n\n" + except asyncio.CancelledError: + logger.info(f"SSE client disconnected for session {session_id}") + raise + except Exception as e: + logger.error(f"Error during agent execution: {e}", exc_info=True) + err_data = json.dumps({ + "text": f"[Error] **Error during investigation:** {str(e)}", + "last_msg": False, + "session_id": session_id, + "event_type": "error", + }) + yield f"data: {err_data}\n\n" + + done_data = json.dumps({ + "text": "Stream finished.", + "last_msg": True, + "session_id": session_id, + }) yield f"data: {done_data}\n\n" @@ -104,20 +353,46 @@ async def sse_event_generator(message: str, session_id: str) -> AsyncGenerator[s 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"), + user_id: Optional[str] = Query(None, description="Active user ID"), ): """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), + sse_event_generator(message, sess_id, user_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.""" +async def chat_post(request: ChatRequest, http_request: Request): + """REST and SSE chat endpoint for API clients, automated workflows, and web UI.""" sess_id = request.session_id or str(uuid.uuid4()) - return ChatResponse( - response=f"Received query: {request.prompt}", - session_id=sess_id, + query_text = request.message or request.prompt or "" + uid = request.user_id or discover_user_identity() + + accept_header = http_request.headers.get("accept", "") + if "text/event-stream" in accept_header: + return StreamingResponse( + sse_event_generator(query_text, sess_id, uid), + media_type="text/event-stream", + ) + + # For non-streaming REST queries, collect text parts from runner + accumulated = [] + async for chunk_str in sse_event_generator(query_text, sess_id, uid): + if chunk_str.startswith("data: "): + try: + data = json.loads(chunk_str[6:].strip()) + if not data.get("last_msg") and data.get("text"): + accumulated.append(data["text"]) + except Exception: + pass + + response_text = "\n\n".join(accumulated) if accumulated else f"No response generated for: {query_text}" + return JSONResponse( + content={ + "response": response_text, + "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 index 56ad60ee..855e5be4 100644 --- a/run-with-google-adk/src/mcp_security_agent/toolsets.py +++ b/run-with-google-adk/src/mcp_security_agent/toolsets.py @@ -47,6 +47,82 @@ def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: logger.warning("google.adk.tools.mcp_tool not available; using mock/fallback toolset representation.") return toolsets + def _build_stdio_env() -> dict[str, str]: + """Constructs environment dictionary for Stdio subprocesses with credential and project isolation.""" + import os + env = dict(os.environ) + + # Propagate local ADC and CloudSDK config + if settings.google_application_credentials: + env["GOOGLE_APPLICATION_CREDENTIALS"] = settings.google_application_credentials + if "CLOUDSDK_CONFIG" not in env: + env["CLOUDSDK_CONFIG"] = str(Path(settings.google_application_credentials).parent) + + # Propagate credentials & impersonation + if settings.secops_sa_path: + env["SECOPS_SA_PATH"] = settings.secops_sa_path + if settings.secops_impersonate_service_account: + env["SECOPS_IMPERSONATE_SERVICE_ACCOUNT"] = settings.secops_impersonate_service_account + + # Propagate GCP Project (for SCC, Vertex AI, and general Cloud SDK) + gcp_project = ( + settings.google_cloud_project + or os.environ.get("GOOGLE_CLOUD_PROJECT") + or os.environ.get("GCP_PROJECT_ID") + or settings.chronicle_project_id + ) + if gcp_project: + env["GOOGLE_CLOUD_PROJECT"] = gcp_project + + # Propagate Chronicle SIEM Project (independent from SCC/GCP project) + chronicle_project = ( + settings.chronicle_project_id + or gcp_project + ) + if chronicle_project: + env["CHRONICLE_PROJECT_ID"] = chronicle_project + + if settings.chronicle_customer_id: + env["CHRONICLE_CUSTOMER_ID"] = settings.chronicle_customer_id + if settings.chronicle_region: + env["CHRONICLE_REGION"] = settings.chronicle_region + + # Propagate GTI / SOAR params if set + if settings.vt_apikey: + env["VT_APIKEY"] = settings.vt_apikey + if settings.soar_url: + env["SOAR_URL"] = settings.soar_url + if settings.soar_app_key: + env["SOAR_APP_KEY"] = settings.soar_app_key + + # Cloudtop mTLS bypass + env.setdefault("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") + env.setdefault("GOOGLE_API_USE_MTLS_ENDPOINT", "never") + env.setdefault("CLOUDSDK_CONTEXT_AWARE_USE_CLIENT_CERTIFICATE", "false") + + return env + + stdio_env = _build_stdio_env() + + def _get_stdio_cmd_args(server_path: Path, script_relpath: str) -> tuple[str, list[str], dict[str, str]]: + import shutil + import sys + import os + + env = dict(stdio_env) + if shutil.which("uv"): + return "uv", ["--directory", str(server_path), "run", script_relpath], env + + logger.warning( + "uv executable not found in PATH; falling back to sys.executable (%s) for %s", + sys.executable, + server_path.name, + ) + existing_pp = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{server_path}{os.pathsep}{existing_pp}" if existing_pp else str(server_path) + script_full = str(server_path / script_relpath) + return sys.executable, [script_full], env + # 1. Google SecOps SIEM MCP if settings.load_secops_mcp: if settings.secops_mcp_url: @@ -54,10 +130,12 @@ def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: else: secops_dir = server_dir / "secops" logger.info("Configuring SecOps SIEM MCP via Stdio subprocess at %s", secops_dir) + cmd, args, env = _get_stdio_cmd_args(secops_dir, "secops_mcp/server.py") conn = StdioConnectionParams( server_params=StdioServerParameters( - command="uv", - args=["--directory", str(secops_dir), "run", "secops_mcp/server.py"], + command=cmd, + args=args, + env=env, ), timeout=settings.stdio_timeout_seconds, ) @@ -70,10 +148,12 @@ def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: else: scc_dir = server_dir / "scc" logger.info("Configuring SCC MCP via Stdio subprocess at %s", scc_dir) + cmd, args, env = _get_stdio_cmd_args(scc_dir, "scc_mcp.py") conn = StdioConnectionParams( server_params=StdioServerParameters( - command="uv", - args=["--directory", str(scc_dir), "run", "scc_mcp.py"], + command=cmd, + args=args, + env=env, ), timeout=settings.stdio_timeout_seconds, ) @@ -86,10 +166,12 @@ def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: else: gti_dir = server_dir / "gti" logger.info("Configuring GTI MCP via Stdio subprocess at %s", gti_dir) + cmd, args, env = _get_stdio_cmd_args(gti_dir, "gti_mcp/server.py") conn = StdioConnectionParams( server_params=StdioServerParameters( - command="uv", - args=["--directory", str(gti_dir), "run", "gti_mcp/server.py"], + command=cmd, + args=args, + env=env, ), timeout=settings.stdio_timeout_seconds, ) @@ -102,13 +184,16 @@ def build_mcp_toolsets(settings: AgentSettings) -> List[Any]: else: soar_dir = server_dir / "secops-soar" logger.info("Configuring SecOps SOAR MCP via Stdio subprocess at %s", soar_dir) + cmd, args, env = _get_stdio_cmd_args(soar_dir, "secops_soar_mcp/server.py") conn = StdioConnectionParams( server_params=StdioServerParameters( - command="uv", - args=["--directory", str(soar_dir), "run", "secops_soar_mcp/server.py"], + command=cmd, + args=args, + env=env, ), timeout=settings.stdio_timeout_seconds, ) toolsets.append(McpToolset(connection_params=conn)) return toolsets + diff --git a/run-with-google-adk/static/app.css b/run-with-google-adk/static/app.css new file mode 100644 index 00000000..480585ac --- /dev/null +++ b/run-with-google-adk/static/app.css @@ -0,0 +1,827 @@ +/* Google SecOps AI Assistant - Modernized Styling */ +:root { + --bg-primary: #121316; + --bg-secondary: #1a1b1f; + --bg-surface: #22242a; + --bg-surface-hover: #2c2e35; + --border-color: #33353d; + --border-subtle: #272830; + + --text-primary: #e8eaed; + --text-secondary: #9aa0a6; + --text-tertiary: #5f6368; + + --accent-primary: #8ab4f8; + --accent-primary-hover: #aecbfa; + --accent-surface: rgba(138, 180, 248, 0.12); + + --color-success: #81c995; + --color-success-bg: rgba(129, 201, 149, 0.15); + --color-warning: #fdd663; + --color-warning-bg: rgba(253, 214, 99, 0.15); + --color-danger: #f28b82; + --color-danger-bg: rgba(242, 139, 130, 0.15); + + --code-bg: #18191c; + --font-sans: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif; + --font-mono: 'Roboto Mono', 'SFMono-Regular', Menlo, Monaco, Consolas, monospace; + + --radius-sm: 6px; + --radius-md: 10px; + --radius-lg: 16px; + --shadow-sm: 0 1px 3px rgba(0, 0, 0, 0.2); + --shadow-md: 0 4px 12px rgba(0, 0, 0, 0.3); + --shadow-lg: 0 8px 24px rgba(0, 0, 0, 0.4); +} + +[data-theme="light"] { + --bg-primary: #f8f9fa; + --bg-secondary: #ffffff; + --bg-surface: #f1f3f4; + --bg-surface-hover: #e8eaed; + --border-color: #dadce0; + --border-subtle: #e8eaed; + + --text-primary: #202124; + --text-secondary: #5f6368; + --text-tertiary: #80868b; + + --accent-primary: #1a73e8; + --accent-primary-hover: #1557b0; + --accent-surface: rgba(26, 115, 232, 0.08); + + --code-bg: #f1f3f4; + --shadow-sm: 0 1px 2px rgba(60, 64, 67, 0.15); + --shadow-md: 0 4px 8px rgba(60, 64, 67, 0.15); + --shadow-lg: 0 8px 16px rgba(60, 64, 67, 0.15); +} + +* { + box-sizing: border-box; + margin: 0; + padding: 0; + scrollbar-width: thin; + scrollbar-color: var(--border-color) transparent; +} + +::-webkit-scrollbar { + width: 8px; + height: 8px; +} + +::-webkit-scrollbar-track { + background: transparent; +} + +::-webkit-scrollbar-thumb { + background-color: var(--border-color); + border-radius: 4px; +} + +::-webkit-scrollbar-thumb:hover { + background-color: var(--text-tertiary); +} + +html { + height: 100%; + height: 100vh; + height: 100dvh; +} + +body { + font-family: var(--font-sans); + background-color: var(--bg-primary); + color: var(--text-primary); + height: 100%; + height: 100vh; + height: 100dvh; + max-height: 100vh; + max-height: 100dvh; + display: flex; + flex-direction: column; + overflow: hidden; + transition: background-color 0.2s ease, color 0.2s ease; +} + +/* Header */ +.app-header { + height: 56px; + background-color: var(--bg-secondary); + border-bottom: 1px solid var(--border-color); + display: flex; + align-items: center; + justify-content: space-between; + padding: 0 20px; + z-index: 10; + flex-shrink: 0; +} + +.brand-section { + display: flex; + align-items: center; + gap: 12px; +} + +.brand-icon { + width: 28px; + height: 28px; + border-radius: 6px; + background: linear-gradient(135deg, #1a73e8, #4285f4); + display: flex; + align-items: center; + justify-content: center; + color: #fff; + font-weight: 700; + font-size: 16px; +} + +.brand-title { + font-size: 16px; + font-weight: 600; + letter-spacing: -0.2px; + color: var(--text-primary); +} + +.brand-badge { + font-size: 11px; + padding: 2px 6px; + border-radius: 4px; + background-color: var(--accent-surface); + color: var(--accent-primary); + font-weight: 500; +} + +.header-status-pills { + display: flex; + align-items: center; + gap: 8px; +} + +.status-pill { + font-size: 11px; + padding: 3px 8px; + border-radius: 12px; + border: 1px solid var(--border-color); + color: var(--text-secondary); + display: flex; + align-items: center; + gap: 5px; +} + +.status-dot { + width: 6px; + height: 6px; + border-radius: 50%; + background-color: var(--text-tertiary); +} + +.status-dot.active { + background-color: var(--color-success); + box-shadow: 0 0 6px var(--color-success); +} + +.header-actions { + display: flex; + align-items: center; + gap: 12px; +} + +.user-chip { + display: flex; + align-items: center; + gap: 8px; + padding: 4px 10px; + background-color: var(--bg-surface); + border: 1px solid var(--border-color); + border-radius: 20px; + font-size: 12px; + color: var(--text-primary); + cursor: pointer; + transition: background-color 0.2s ease; +} + +.user-chip:hover { + background-color: var(--bg-surface-hover); +} + +.user-avatar { + width: 20px; + height: 20px; + border-radius: 50%; + background-color: var(--accent-primary); + color: var(--bg-primary); + display: flex; + align-items: center; + justify-content: center; + font-weight: 700; + font-size: 10px; +} + +.icon-btn { + background: transparent; + border: 1px solid var(--border-color); + border-radius: var(--radius-sm); + color: var(--text-secondary); + width: 32px; + height: 32px; + display: flex; + align-items: center; + justify-content: center; + cursor: pointer; + transition: all 0.2s ease; +} + +.icon-btn:hover { + color: var(--text-primary); + background-color: var(--bg-surface); +} + +/* App Container */ +.app-layout { + display: flex; + flex: 1 1 0%; + min-height: 0; + height: calc(100vh - 56px); + height: calc(100dvh - 56px); + overflow: hidden; +} + +/* Sidebar */ +.app-sidebar { + width: 280px; + background-color: var(--bg-secondary); + border-right: 1px solid var(--border-color); + display: flex; + flex-direction: column; + padding: 16px; + gap: 16px; + flex-shrink: 0; + height: 100%; + min-height: 0; + overflow-y: auto; +} + +.sidebar-action-btn { + width: 100%; + padding: 10px 14px; + background-color: var(--accent-primary); + color: #121316; + font-weight: 600; + font-size: 13px; + border: none; + border-radius: var(--radius-sm); + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + gap: 8px; + transition: background-color 0.2s ease; +} + +.sidebar-action-btn:hover { + background-color: var(--accent-primary-hover); +} + +.sidebar-section-title { + font-size: 11px; + font-weight: 600; + text-transform: uppercase; + letter-spacing: 0.5px; + color: var(--text-tertiary); + margin-bottom: 6px; +} + +.quick-prompts-list { + display: flex; + flex-direction: column; + gap: 6px; +} + +.quick-prompt-card { + padding: 8px 10px; + background-color: var(--bg-surface); + border: 1px solid var(--border-color); + border-radius: var(--radius-sm); + cursor: pointer; + font-size: 12px; + color: var(--text-secondary); + text-align: left; + transition: all 0.2s ease; +} + +.quick-prompt-card:hover { + background-color: var(--bg-surface-hover); + color: var(--text-primary); + border-color: var(--accent-primary); +} + +.session-metadata { + margin-top: auto; + padding: 12px; + background-color: var(--bg-surface); + border: 1px solid var(--border-color); + border-radius: var(--radius-sm); + font-size: 11px; + color: var(--text-secondary); + display: flex; + flex-direction: column; + gap: 4px; +} + +.session-metadata code { + font-family: var(--font-mono); + color: var(--accent-primary); + word-break: break-all; +} + +/* Chat Main Area */ +.chat-workspace { + flex: 1 1 0%; + min-width: 0; + min-height: 0; + height: 100%; + display: flex; + flex-direction: column; + background-color: var(--bg-primary); + position: relative; + overflow: hidden; +} + +.messages-container { + flex: 1 1 0%; + min-height: 0; + overflow-y: auto; + overflow-x: hidden; + padding: 24px; + display: flex; + flex-direction: column; + gap: 20px; +} + +/* Empty State */ +.empty-state { + margin: auto; + max-width: 540px; + text-align: center; + display: flex; + flex-direction: column; + align-items: center; + gap: 16px; + padding: 40px 20px; +} + +.empty-icon { + width: 56px; + height: 56px; + border-radius: 16px; + background-color: var(--accent-surface); + color: var(--accent-primary); + display: flex; + align-items: center; + justify-content: center; + font-size: 28px; +} + +.empty-title { + font-size: 20px; + font-weight: 600; + color: var(--text-primary); +} + +.empty-desc { + font-size: 14px; + color: var(--text-secondary); + line-height: 1.5; +} + +/* Messages */ +.message-row { + display: flex; + gap: 12px; + max-width: 860px; + width: 100%; +} + +.message-row.user { + align-self: flex-end; + flex-direction: row-reverse; +} + +.message-row.agent { + align-self: flex-start; +} + +.msg-avatar { + width: 32px; + height: 32px; + border-radius: 50%; + flex-shrink: 0; + display: flex; + align-items: center; + justify-content: center; + font-size: 13px; + font-weight: 600; +} + +.message-row.user .msg-avatar { + background-color: var(--accent-primary); + color: var(--bg-primary); +} + +.message-row.agent .msg-avatar { + background-color: #34a853; + color: #fff; +} + +.message-body-wrap { + display: flex; + flex-direction: column; + gap: 4px; + max-width: calc(100% - 44px); +} + +.message-header-info { + display: flex; + align-items: center; + gap: 8px; + font-size: 11px; + color: var(--text-tertiary); +} + +.message-row.user .message-header-info { + justify-content: flex-end; +} + +.msg-bubble { + padding: 12px 16px; + border-radius: var(--radius-md); + font-size: 14px; + line-height: 1.6; + word-break: break-word; +} + +.message-row.user .msg-bubble { + background-color: var(--accent-surface); + color: var(--text-primary); + border: 1px solid rgba(138, 180, 248, 0.3); + border-top-right-radius: 2px; +} + +.message-row.agent .msg-bubble { + background-color: var(--bg-secondary); + color: var(--text-primary); + border: 1px solid var(--border-color); + border-top-left-radius: 2px; +} + +/* Markdown styling inside agent bubble */ +.msg-bubble p:not(:last-child) { + margin-bottom: 10px; +} + +.msg-bubble ul, .msg-bubble ol { + margin: 8px 0 8px 20px; +} + +.msg-bubble li { + margin-bottom: 4px; +} + +.msg-bubble code:not(pre code) { + font-family: var(--font-mono); + font-size: 12px; + padding: 2px 6px; + background-color: var(--code-bg); + border: 1px solid var(--border-color); + border-radius: 4px; + color: var(--accent-primary); +} + +.msg-bubble pre { + margin: 10px 0; + position: relative; + border-radius: var(--radius-sm); + background-color: var(--code-bg); + border: 1px solid var(--border-color); + overflow-x: auto; +} + +.msg-bubble pre code { + font-family: var(--font-mono); + font-size: 13px; + display: block; + padding: 12px 14px; + color: var(--text-primary); + line-height: 1.45; +} + +.code-copy-btn { + position: absolute; + top: 6px; + right: 6px; + background: var(--bg-surface); + border: 1px solid var(--border-color); + border-radius: 4px; + color: var(--text-secondary); + font-size: 11px; + padding: 3px 8px; + cursor: pointer; + transition: all 0.15s ease; +} + +.code-copy-btn:hover { + background: var(--bg-surface-hover); + color: var(--text-primary); +} + +/* Tool Execution Accordion */ +.tool-accordion { + margin: 8px 0; + border: 1px solid var(--border-color); + border-radius: var(--radius-sm); + background-color: var(--bg-surface); + overflow: hidden; +} + +.tool-header { + padding: 8px 12px; + font-size: 12px; + font-weight: 500; + color: var(--text-secondary); + cursor: pointer; + display: flex; + align-items: center; + justify-content: space-between; + background-color: var(--bg-surface); + transition: background-color 0.15s ease; +} + +.tool-header:hover { + background-color: var(--bg-surface-hover); +} + +.tool-badge { + font-family: var(--font-mono); + font-size: 11px; + padding: 2px 6px; + border-radius: 4px; + background-color: var(--accent-surface); + color: var(--accent-primary); +} + +.tool-content { + padding: 10px 12px; + border-top: 1px solid var(--border-color); + background-color: var(--code-bg); + font-family: var(--font-mono); + font-size: 12px; + color: var(--text-secondary); + overflow-x: auto; +} + +/* Waiting / Streaming Indicator */ +.streaming-pulse { + display: inline-block; + width: 8px; + height: 8px; + border-radius: 50%; + background-color: var(--accent-primary); + animation: pulse 1s infinite alternate; + margin-right: 6px; +} + +@keyframes pulse { + from { opacity: 0.3; transform: scale(0.8); } + to { opacity: 1; transform: scale(1.1); } +} + +/* Input Area */ +.chat-input-area { + padding: 16px 24px 20px; + background-color: var(--bg-primary); + border-top: 1px solid var(--border-color); + display: flex; + flex-direction: column; + gap: 8px; + flex-shrink: 0; +} + +.input-box-wrapper { + background-color: var(--bg-secondary); + border: 1px solid var(--border-color); + border-radius: var(--radius-md); + padding: 8px 12px; + display: flex; + align-items: flex-end; + gap: 10px; + transition: border-color 0.2s ease, box-shadow 0.2s ease; +} + +.input-box-wrapper:focus-within { + border-color: var(--accent-primary); + box-shadow: 0 0 0 2px var(--accent-surface); +} + +.chat-textarea { + flex: 1; + background: transparent; + border: none; + outline: none; + resize: none; + font-family: var(--font-sans); + font-size: 14px; + line-height: 1.5; + color: var(--text-primary); + max-height: 160px; + min-height: 24px; +} + +.chat-textarea::placeholder { + color: var(--text-tertiary); +} + +.send-btn { + padding: 8px 16px; + border: none; + border-radius: var(--radius-sm); + background-color: var(--accent-primary); + color: #121316; + font-weight: 600; + font-size: 13px; + cursor: pointer; + display: flex; + align-items: center; + gap: 6px; + transition: background-color 0.15s ease, opacity 0.15s ease; + flex-shrink: 0; +} + +.send-btn:hover:not(:disabled) { + background-color: var(--accent-primary-hover); +} + +.send-btn:disabled { + opacity: 0.4; + cursor: not-allowed; +} + +.cancel-btn { + background-color: var(--color-danger); + color: #fff; +} + +.input-footer-hints { + font-size: 11px; + color: var(--text-tertiary); + display: flex; + justify-content: space-between; + padding: 0 4px; +} + +/* Toast System */ +.toast-container { + position: fixed; + top: 68px; + right: 20px; + display: flex; + flex-direction: column; + gap: 10px; + z-index: 100; + pointer-events: none; +} + +.toast { + pointer-events: auto; + min-width: 280px; + max-width: 400px; + padding: 10px 14px; + border-radius: var(--radius-sm); + font-size: 13px; + display: flex; + align-items: center; + justify-content: space-between; + box-shadow: var(--shadow-md); + animation: slideIn 0.25s ease-out forwards; +} + +@keyframes slideIn { + from { opacity: 0; transform: translateX(30px); } + to { opacity: 1; transform: translateX(0); } +} + +.toast.info { + background-color: var(--bg-surface); + border: 1px solid var(--accent-primary); + color: var(--text-primary); +} + +.toast.success { + background-color: var(--color-success-bg); + border: 1px solid var(--color-success); + color: var(--text-primary); +} + +.toast.warning { + background-color: var(--color-warning-bg); + border: 1px solid var(--color-warning); + color: var(--text-primary); +} + +.toast.error { + background-color: var(--color-danger-bg); + border: 1px solid var(--color-danger); + color: var(--text-primary); +} + +.toast-close { + background: transparent; + border: none; + color: var(--text-secondary); + cursor: pointer; + font-size: 16px; + padding: 0 4px; +} + +/* User Profile Dialog */ +.modal-overlay { + position: fixed; + inset: 0; + background-color: rgba(0, 0, 0, 0.6); + display: flex; + align-items: center; + justify-content: center; + z-index: 200; + animation: fadeIn 0.2s ease; +} + +@keyframes fadeIn { + from { opacity: 0; } + to { opacity: 1; } +} + +.modal-card { + background-color: var(--bg-secondary); + border: 1px solid var(--border-color); + border-radius: var(--radius-md); + padding: 24px; + width: 100%; + max-width: 380px; + box-shadow: var(--shadow-lg); + display: flex; + flex-direction: column; + gap: 16px; +} + +.modal-title { + font-size: 16px; + font-weight: 600; +} + +.modal-input { + width: 100%; + padding: 8px 12px; + border-radius: var(--radius-sm); + border: 1px solid var(--border-color); + background-color: var(--bg-surface); + color: var(--text-primary); + font-size: 14px; + outline: none; +} + +.modal-input:focus { + border-color: var(--accent-primary); +} + +.modal-actions { + display: flex; + justify-content: flex-end; + gap: 10px; +} + +.btn-secondary { + background: transparent; + border: 1px solid var(--border-color); + color: var(--text-secondary); + padding: 6px 12px; + border-radius: var(--radius-sm); + font-size: 13px; + cursor: pointer; +} + +.btn-primary { + background-color: var(--accent-primary); + border: none; + color: #121316; + font-weight: 600; + padding: 6px 14px; + border-radius: var(--radius-sm); + font-size: 13px; + cursor: pointer; +} + +/* Responsive adjustments */ +@media (max-width: 768px) { + .app-sidebar { + display: none; + } +} diff --git a/run-with-google-adk/static/app.js b/run-with-google-adk/static/app.js index 1369132f..61599340 100644 --- a/run-with-google-adk/static/app.js +++ b/run-with-google-adk/static/app.js @@ -1,176 +1,526 @@ -// app.js +// Google SecOps AI Assistant - Modernized Zero-Build Web Client +// Single-Page Architecture with SSE Streaming, Toast Notifications, and Dynamic MCP Status document.addEventListener('DOMContentLoaded', () => { - const chatWindow = document.getElementById('chat-window'); - const userInput = document.getElementById('user-input'); - const submitBtn = document.getElementById('submit-btn'); - const clearBtn = document.getElementById('clear-btn'); - const sessionInfoDiv = document.getElementById('session-info'); - const darkModeToggle = document.getElementById('darkModeToggle'); - - let currentSessionId = null; // Variable to store the session ID - let requestSentTime = 0; // Timestamp when the user message was sent - let lastAgentMessageTime = 0; // Timestamp of the last received agent message - - // --- Dark Mode Logic --- - function applyTheme(isDarkMode) { - if (isDarkMode) { - document.body.classList.add('dark-mode'); - } else { - document.body.classList.remove('dark-mode'); - } - } + const API_BASE_URL = window.location.origin; + + // DOM Elements + const messagesContainer = document.getElementById('messagesContainer'); + const chatTextarea = document.getElementById('chatTextarea'); + const sendBtn = document.getElementById('sendBtn'); + const cancelBtn = document.getElementById('cancelBtn'); + const newInvestigationBtn = document.getElementById('newInvestigationBtn'); + const appTitle = document.getElementById('appTitle'); + const currentSessionDisplay = document.getElementById('currentSessionDisplay'); + const currentUserDisplay = document.getElementById('currentUserDisplay'); + const userAvatar = document.getElementById('userAvatar'); + const userChip = document.getElementById('userChip'); + const themeToggleBtn = document.getElementById('themeToggleBtn'); + const toastContainer = document.getElementById('toastContainer'); + const emptyState = document.getElementById('emptyState'); + + // Modal Elements + const userModal = document.getElementById('userModal'); + const usernameInput = document.getElementById('usernameInput'); + const saveUserBtn = document.getElementById('saveUserBtn'); + const cancelUserBtn = document.getElementById('cancelUserBtn'); + + // MCP Pills + const pillSecops = document.getElementById('pillSecops'); + const pillScc = document.getElementById('pillScc'); + const pillGti = document.getElementById('pillGti'); + const pillSoar = document.getElementById('pillSoar'); + + // Application State + let currentSessionId = null; + const savedUser = localStorage.getItem('username'); + let currentUserId = (savedUser && savedUser !== 'secops_user') ? savedUser : null; + let isStreaming = false; + let activeAbortController = null; + let requestStartTime = 0; + let lastChunkTime = 0; + // Initialize Theme + function initTheme() { const savedTheme = localStorage.getItem('theme'); - if (savedTheme === 'dark') { - darkModeToggle.checked = true; - applyTheme(true); + if (savedTheme) { + document.documentElement.setAttribute('data-theme', savedTheme); + } else if (window.matchMedia && window.matchMedia('(prefers-color-scheme: light)').matches) { + document.documentElement.setAttribute('data-theme', 'light'); } else { - darkModeToggle.checked = false; - applyTheme(false); + document.documentElement.setAttribute('data-theme', 'dark'); } + } + + function toggleTheme() { + const currentTheme = document.documentElement.getAttribute('data-theme') || 'dark'; + const newTheme = currentTheme === 'dark' ? 'light' : 'dark'; + document.documentElement.setAttribute('data-theme', newTheme); + localStorage.setItem('theme', newTheme); + showToast(`Theme switched to ${newTheme} mode`, 'info', 2000); + } + + // Toast Notification System + function showToast(message, type = 'info', duration = 3500) { + const toast = document.createElement('div'); + toast.className = `toast ${type}`; + + const textSpan = document.createElement('span'); + textSpan.textContent = message; + toast.appendChild(textSpan); + + const closeBtn = document.createElement('button'); + closeBtn.className = 'toast-close'; + closeBtn.innerHTML = '×'; + closeBtn.onclick = () => removeToast(toast); + toast.appendChild(closeBtn); + + toastContainer.appendChild(toast); + + const timer = setTimeout(() => { + removeToast(toast); + }, duration); - darkModeToggle.addEventListener('change', () => { - if (darkModeToggle.checked) { - applyTheme(true); - localStorage.setItem('theme', 'dark'); - } else { - applyTheme(false); - localStorage.setItem('theme', 'light'); + function removeToast(el) { + clearTimeout(timer); + el.style.opacity = '0'; + el.style.transform = 'translateX(20px)'; + el.style.transition = 'all 0.2s ease-out'; + setTimeout(() => { + if (el.parentNode) { + el.parentNode.removeChild(el); } - }); - // --- End Dark Mode Logic --- - - - // Function to append a message to the chat window - // Now accepts timeElapsed and timeDiff parameters - function appendMessage(text, sender, timeElapsed = null, timeDiff = null) { - const messageDiv = document.createElement('div'); - messageDiv.classList.add('message', sender); - - // Create and append the time display element if timeElapsed is provided - if (timeElapsed !== null) { - const timeDisplay = document.createElement('div'); - timeDisplay.classList.add('message-time'); - let timeText = `${timeElapsed}ms`; - if (timeDiff !== null && sender === 'agent') { // Only show diff for agent messages - timeText += ` (${timeDiff}ms)`; - } - timeDisplay.textContent = timeText; - messageDiv.appendChild(timeDisplay); + }, 200); + } + } + + // Update User UI + function updateUserUI() { + if (!currentUserId) { + currentUserDisplay.textContent = 'secops_user'; + userAvatar.textContent = 'S'; + return; + } + currentUserDisplay.textContent = currentUserId; + userAvatar.textContent = currentUserId.charAt(0).toUpperCase(); + localStorage.setItem('username', currentUserId); + } + + // Fetch App Metadata + async function fetchAppMetadata() { + try { + const res = await fetch(`${API_BASE_URL}/app_name`); + if (res.ok) { + const data = await res.json(); + if (data.app_name) { + appTitle.textContent = data.app_name; + document.title = data.app_name; } + } + } catch (e) { + console.warn('Could not fetch app name:', e); + } + } - const messageContent = document.createElement('div'); // Container for the actual message text/markdown - if (sender === 'agent') { - messageContent.innerHTML = marked.parse(text); - } else { - messageContent.textContent = text; + // Fetch MCP Server Status Pills + async function fetchMcpInfo() { + try { + const res = await fetch(`${API_BASE_URL}/info`); + if (res.ok) { + const data = await res.json(); + if (data.user && !currentUserId) { + currentUserId = data.user; + updateUserUI(); } - messageDiv.appendChild(messageContent); // Append content after time - chatWindow.appendChild(messageDiv); - chatWindow.scrollTop = chatWindow.scrollHeight; + const mcp = data.tools || data.mcp_servers || {}; + updatePill(pillSecops, mcp.secops); + updatePill(pillScc, mcp.scc); + updatePill(pillGti, mcp.gti); + updatePill(pillSoar, mcp.soar); + } + } catch (e) { + console.warn('Could not fetch MCP server info:', e); } + } - // Function to fetch the session ID - async function fetchSessionId() { + function updatePill(pillEl, isEnabled) { + if (!pillEl) return; + const dot = pillEl.querySelector('.status-dot'); + if (dot) { + if (isEnabled) { + dot.classList.add('active'); + pillEl.title = 'Enabled and connected'; + } else { + dot.classList.remove('active'); + pillEl.title = 'Disabled or not configured'; + } + } + } + + // Fetch / Initialize Session + async function initSession(startNew = false) { + try { + let url = `${API_BASE_URL}/get_session`; + if (currentUserId) { + url += `?username=${encodeURIComponent(currentUserId)}`; + } + if (startNew) { + url += (currentUserId ? '&' : '?') + 'start_new_session=Y'; + } + const res = await fetch(url); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + currentSessionId = data.session_id; + if (data.user_id) { + currentUserId = data.user_id; + updateUserUI(); + } + currentSessionDisplay.textContent = currentSessionId.slice(0, 8) + '...'; + currentSessionDisplay.title = currentSessionId; + console.log('Active session initialized:', currentSessionId, 'User:', currentUserId); + } catch (err) { + console.error('Session initialization failed:', err); + showToast('Could not initialize session with backend', 'error'); + } + } + + // Safe Markdown Rendering with DOMPurify Sanitization + function renderMarkdown(rawText) { + if (window.marked && typeof window.marked.parse === 'function') { + const rawHtml = window.marked.parse(rawText); + if (window.DOMPurify && typeof window.DOMPurify.sanitize === 'function') { + return window.DOMPurify.sanitize(rawHtml); + } + console.warn('DOMPurify not available, failing closed to escaped text for security.'); + const div = document.createElement('div'); + div.textContent = rawText; + return `
${div.innerHTML}
`; + } + const div = document.createElement('div'); + div.textContent = rawText; + return div.innerHTML; + } + + // Append Message + function appendMessage(text, sender, timeElapsed = null, timeDiff = null) { + if (emptyState && emptyState.style.display !== 'none') { + emptyState.style.display = 'none'; + } + + const row = document.createElement('div'); + row.className = `message-row ${sender}`; + + const avatar = document.createElement('div'); + avatar.className = 'msg-avatar'; + avatar.textContent = sender === 'user' ? currentUserId.charAt(0).toUpperCase() : 'G'; + row.appendChild(avatar); + + const bodyWrap = document.createElement('div'); + bodyWrap.className = 'message-body-wrap'; + + const headerInfo = document.createElement('div'); + headerInfo.className = 'message-header-info'; + const authorSpan = document.createElement('span'); + authorSpan.textContent = sender === 'user' ? currentUserId : 'Google SecOps Agent'; + headerInfo.appendChild(authorSpan); + + if (timeElapsed !== null) { + const timeSpan = document.createElement('span'); + let timingStr = `${timeElapsed}ms`; + if (timeDiff !== null && sender === 'agent') { + timingStr += ` (${timeDiff}ms)`; + } + timeSpan.textContent = timingStr; + headerInfo.appendChild(timeSpan); + } + + bodyWrap.appendChild(headerInfo); + + const bubble = document.createElement('div'); + bubble.className = 'msg-bubble'; + + if (sender === 'agent') { + bubble.innerHTML = renderMarkdown(text); + addCopyButtonsToPre(bubble); + } else { + bubble.textContent = text; + } + + bodyWrap.appendChild(bubble); + row.appendChild(bodyWrap); + + messagesContainer.appendChild(row); + scrollToBottom(); + return bubble; + } + + // Code Copy Buttons + function addCopyButtonsToPre(container) { + const preBlocks = container.querySelectorAll('pre'); + preBlocks.forEach((pre) => { + if (pre.querySelector('.code-copy-btn')) return; + const code = pre.querySelector('code'); + const copyBtn = document.createElement('button'); + copyBtn.className = 'code-copy-btn'; + copyBtn.textContent = 'Copy'; + copyBtn.onclick = async () => { try { - const response = await fetch('/get_session'); - if (!response.ok) { - throw new Error(`HTTP error! status: ${response.status}`); - } - const data = await response.json(); - currentSessionId = data.session_id; - sessionInfoDiv.textContent = `Session ID: ${currentSessionId}`; - sessionInfoDiv.classList.remove('alert-info'); - sessionInfoDiv.classList.add('alert-success'); - console.log('Session ID fetched:', currentSessionId); - submitBtn.disabled = false; // Enable submit button once session is loaded - userInput.disabled = false; // Enable input once session is loaded - } catch (error) { - console.error('Error fetching session ID:', error); - sessionInfoDiv.textContent = 'Error loading session ID.'; - sessionInfoDiv.classList.remove('alert-info'); - sessionInfoDiv.classList.add('alert-danger'); - // Disable submit button if session ID cannot be fetched - submitBtn.disabled = true; - userInput.disabled = true; + const textToCopy = code ? code.innerText : pre.innerText; + await navigator.clipboard.writeText(textToCopy); + copyBtn.textContent = 'Copied!'; + setTimeout(() => { + copyBtn.textContent = 'Copy'; + }, 1800); + } catch (e) { + showToast('Failed to copy code to clipboard', 'warning'); } + }; + pre.appendChild(copyBtn); + }); + } + + function scrollToBottom() { + messagesContainer.scrollTop = messagesContainer.scrollHeight; + } + + // Stream Response Handler + async function handleUserSubmit() { + const prompt = chatTextarea.value.trim(); + if (!prompt || isStreaming) return; + + if (!currentSessionId) { + await initSession(); + if (!currentSessionId) { + showToast('Waiting for session initialization. Try again in a moment.', 'warning'); + return; + } } - // Fetch session ID on page load - fetchSessionId(); - // Initially disable submit button until session ID is loaded - submitBtn.disabled = true; - userInput.disabled = true; + appendMessage(prompt, 'user', 0); + chatTextarea.value = ''; + autoResizeTextarea(); + // Prepare Agent Streaming Container + isStreaming = true; + sendBtn.style.display = 'none'; + cancelBtn.style.display = 'flex'; + requestStartTime = Date.now(); + lastChunkTime = requestStartTime; - // Event listener for the Submit button - submitBtn.addEventListener('click', async () => { - const message = userInput.value.trim(); - if (!currentSessionId) { - appendMessage('Error: Session ID not available. Please refresh the page.', 'agent'); - return; - } - if (message) { - // Display user message immediately with 0ms elapsed time - appendMessage(message, 'user', 0); - userInput.value = ''; // Clear input field + const streamRow = document.createElement('div'); + streamRow.className = 'message-row agent'; + + const avatar = document.createElement('div'); + avatar.className = 'msg-avatar'; + avatar.textContent = 'G'; + streamRow.appendChild(avatar); + + const bodyWrap = document.createElement('div'); + bodyWrap.className = 'message-body-wrap'; + + const headerInfo = document.createElement('div'); + headerInfo.className = 'message-header-info'; + headerInfo.innerHTML = 'Investigating...'; + bodyWrap.appendChild(headerInfo); + + const bubble = document.createElement('div'); + bubble.className = 'msg-bubble'; + bubble.innerHTML = 'Analyzing security telemetry...'; + bodyWrap.appendChild(bubble); + + streamRow.appendChild(bodyWrap); + messagesContainer.appendChild(streamRow); + scrollToBottom(); - // Capture the time just before sending the request - requestSentTime = Date.now(); - lastAgentMessageTime = requestSentTime; // Reset for new request + activeAbortController = new AbortController(); + let accumulatedText = ''; - // Make a request to the /chat API using Server-Sent Events (SSE) + try { + const response = await fetch(`${API_BASE_URL}/chat`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'Accept': 'text/event-stream, application/json', + }, + body: JSON.stringify({ + session_id: currentSessionId, + user_id: currentUserId, + message: prompt, + }), + signal: activeAbortController.signal, + }); + + if (!response.ok) { + throw new Error(`Server returned HTTP ${response.status}`); + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder('utf-8'); + let buffer = ''; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + let boundary = buffer.indexOf('\n\n'); + + while (boundary !== -1) { + const chunk = buffer.substring(0, boundary); + buffer = buffer.substring(boundary + 2); + + if (chunk.startsWith('data: ')) { try { - const eventSource = new EventSource(`/chat?message=${encodeURIComponent(message)}&session_id=${encodeURIComponent(currentSessionId)}`); - - eventSource.onmessage = (event) => { - const data = JSON.parse(event.data); - const receivedTime = Date.now(); - const timeElapsed = receivedTime - requestSentTime; - const timeDiff = receivedTime - lastAgentMessageTime; // Calculate difference from previous agent message - lastAgentMessageTime = receivedTime; // Update last agent message time - - // Do not show "Stream finished." message - if (data.last_msg && data.text === 'Stream finished.') { - eventSource.close(); // Close the connection when last_msg is true - console.log('SSE connection closed.'); - return; // Do not append this message - } - - // Display agent message with calculated elapsed time and difference - appendMessage(data.text, 'agent', timeElapsed, timeDiff); - - if (data.last_msg) { - eventSource.close(); // Close the connection when last_msg is true - console.log('SSE connection closed.'); - } - }; - - eventSource.onerror = (error) => { - console.error('EventSource failed:', error); - eventSource.close(); - appendMessage('Error receiving response from agent or stream ended unexpectedly.', 'agent'); - }; - - } catch (error) { - console.error('Failed to connect to SSE:', error); - appendMessage('Failed to initiate chat session.', 'agent'); + const data = JSON.parse(chunk.substring(6)); + const now = Date.now(); + const timeElapsed = now - requestStartTime; + const timeDiff = now - lastChunkTime; + lastChunkTime = now; + + if (data.last_msg && data.text === 'Stream finished.') { + headerInfo.innerHTML = `Google SecOps Agent${timeElapsed}ms`; + reader.cancel(); + break; + } else if (data.text) { + const isBlock = data.event_type === 'tool_call' || + data.event_type === 'tool_response' || + data.event_type === 'error' || + (!data.event_type && (data.text.startsWith('[Tool]') || data.text.startsWith('[Error]') || data.text.startsWith('[Warning]'))); + + if (!accumulatedText) { + accumulatedText = data.text; + } else if (isBlock) { + accumulatedText += '\n\n' + data.text; + } else { + accumulatedText += data.text; + } + bubble.innerHTML = renderMarkdown(accumulatedText); + addCopyButtonsToPre(bubble); + headerInfo.innerHTML = `${timeElapsed}ms (${timeDiff}ms)`; + scrollToBottom(); + } + } catch (jsonErr) { + console.warn('Could not parse SSE chunk:', chunk, jsonErr); } + } + boundary = buffer.indexOf('\n\n'); } - }); + } - // Event listener for the Clear button - clearBtn.addEventListener('click', () => { - chatWindow.innerHTML = ''; // Clear all messages from the chat window - // Optionally, re-fetch session ID or clear it if desired - }); + const totalTime = Date.now() - requestStartTime; + headerInfo.innerHTML = `Google SecOps Agent${totalTime}ms`; + } catch (err) { + if (err.name === 'AbortError') { + headerInfo.innerHTML = 'Google SecOps Agent(Cancelled)'; + showToast('Investigation query cancelled', 'info'); + } else { + console.error('Chat error:', err); + headerInfo.innerHTML = 'Google SecOps Agent(Error)'; + bubble.innerHTML = `Failed to get response from server: ${err.message}`; + showToast(`Request failed: ${err.message}`, 'error'); + } + } finally { + isStreaming = false; + activeAbortController = null; + sendBtn.style.display = 'flex'; + cancelBtn.style.display = 'none'; + scrollToBottom(); + chatTextarea.focus(); + } + } - // Allow sending message with Enter key - userInput.addEventListener('keypress', (event) => { - if (event.key === 'Enter' && !event.shiftKey) { // Shift+Enter for new line - event.preventDefault(); // Prevent default Enter behavior (new line) - submitBtn.click(); // Trigger submit button click - } + // Cancel Streaming + function cancelStreaming() { + if (activeAbortController) { + activeAbortController.abort(); + } + } + + // Auto-resize Textarea + function autoResizeTextarea() { + chatTextarea.style.height = 'auto'; + chatTextarea.style.height = Math.min(chatTextarea.scrollHeight, 160) + 'px'; + } + + // Modal Handling + function openUserModal() { + usernameInput.value = currentUserId; + userModal.style.display = 'flex'; + usernameInput.focus(); + } + + function closeUserModal() { + userModal.style.display = 'none'; + } + + function saveUser() { + const newName = usernameInput.value.trim(); + if (!newName) { + showToast('Username cannot be empty', 'warning'); + return; + } + if (isStreaming) { + cancelStreaming(); + } + currentUserId = newName; + updateUserUI(); + closeUserModal(); + initSession(true); + showToast(`Active user set to: ${currentUserId}`, 'success'); + } + + // Event Listeners + sendBtn.addEventListener('click', handleUserSubmit); + cancelBtn.addEventListener('click', cancelStreaming); + + chatTextarea.addEventListener('input', autoResizeTextarea); + chatTextarea.addEventListener('keydown', (e) => { + if (e.key === 'Enter' && !e.shiftKey) { + e.preventDefault(); + handleUserSubmit(); + } + }); + + themeToggleBtn.addEventListener('click', toggleTheme); + + userChip.addEventListener('click', openUserModal); + saveUserBtn.addEventListener('click', saveUser); + cancelUserBtn.addEventListener('click', closeUserModal); + userModal.addEventListener('click', (e) => { + if (e.target === userModal) closeUserModal(); + }); + + newInvestigationBtn.addEventListener('click', () => { + if (isStreaming) { + cancelStreaming(); + } + messagesContainer.innerHTML = ''; + if (emptyState) { + emptyState.style.display = 'flex'; + messagesContainer.appendChild(emptyState); + } + initSession(true); + showToast('Started a new investigation session', 'info'); + chatTextarea.focus(); + }); + + // Quick Prompt Cards + document.querySelectorAll('.quick-prompt-card').forEach((card) => { + card.addEventListener('click', () => { + const prompt = card.getAttribute('data-prompt'); + if (prompt) { + chatTextarea.value = prompt; + autoResizeTextarea(); + chatTextarea.focus(); + } }); + }); + + // Initialization + initTheme(); + updateUserUI(); + fetchAppMetadata(); + fetchMcpInfo(); + initSession(false); + autoResizeTextarea(); }); diff --git a/run-with-google-adk/static/index.html b/run-with-google-adk/static/index.html index e9814fea..846a6fe8 100644 --- a/run-with-google-adk/static/index.html +++ b/run-with-google-adk/static/index.html @@ -1,133 +1,143 @@ - - - Login - ADK Agent - - + + + Google SecOps AI Assistant + + - + -