diff --git a/soc_agent/__init__.py b/soc_agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/soc_agent/agent.py b/soc_agent/agent.py new file mode 100644 index 00000000..45631a0c --- /dev/null +++ b/soc_agent/agent.py @@ -0,0 +1,49 @@ +import logging +import google.auth +import os +from google.auth.transport.requests import Request +from google.adk.agents import Agent +from google.adk.tools.mcp_tool import McpToolset, StreamableHTTPConnectionParams +from dotenv import load_dotenv + +# Configure logging +logging.basicConfig(level=logging.INFO) + +# Load environment variables +load_dotenv(os.path.join(os.path.dirname(__file__), '.env')) + +# 1. Setup scopes +SCOPES = ["https://www.googleapis.com/auth/chronicle"] + +def get_access_token(): + creds, _ = google.auth.default(scopes=SCOPES) + auth_req = Request() + creds.refresh(auth_req) + return creds.token + +# 2. Configure Toolset +toolset = McpToolset( + connection_params=StreamableHTTPConnectionParams( + url=os.getenv("REMOTE_MCP_URL", "https://chronicle.us.rep.googleapis.com/mcp"), + headers={ + "Authorization": f"Bearer {get_access_token()}", + "Accept": "application/json", + "x-goog-user-project": os.getenv("PROJECT_ID") + } + ) +) + +# 3. Create Agent +root_agent = Agent( + name="oc_agent", + model=os.getenv("GOOGLE_MODEL", "gemini-2.5-flash"), + description="ADK Agent to test the Remote SecOps MCP Server", + instruction=f"""You are an Agent that tests the remote MCP server's tools. + + When using the SecOps MCP, use these parameters for EVERY request: + Customer ID: {os.getenv("CUSTOMER_ID")} + Region: {os.getenv("REGION", "us")} + Project ID: {os.getenv("PROJECT_ID")} + """, + tools=[toolset], +) diff --git a/soc_agent/requirements.in b/soc_agent/requirements.in new file mode 100644 index 00000000..dcfc9fef --- /dev/null +++ b/soc_agent/requirements.in @@ -0,0 +1,5 @@ +google-adk[eval]~=1.22.1 +google-auth +requests +python-dotenv +google-genai diff --git a/soc_agent/sample.env b/soc_agent/sample.env new file mode 100644 index 00000000..2f819255 --- /dev/null +++ b/soc_agent/sample.env @@ -0,0 +1,17 @@ +# Google Cloud Project Config +PROJECT_ID=YOUR_PROJECT_ID +# For SecOps/Chronicle, this is usually a UUID +CUSTOMER_ID=YOUR_CUSTOMER_ID +REGION=us + +# Configuration for ADK/Gemini +GOOGLE_GENAI_USE_VERTEXAI=True +GOOGLE_CLOUD_PROJECT=YOUR_PROJECT_ID +GOOGLE_CLOUD_LOCATION=us-central1 + +# Model Config +GOOGLE_MODEL=gemini-2.5-flash +DEFAULT_PROMPT='You are a helpful security assistant. You can use the available tools to investigate security incidents.' + +# Remote MCP Server URL (Defaults to production if not set) +# REMOTE_MCP_URL="https://chronicle.us.rep.googleapis.com/mcp"