diff --git a/.claude/rules/subagent-development.md b/.claude/rules/subagent-development.md deleted file mode 100644 index 3804b645..00000000 --- a/.claude/rules/subagent-development.md +++ /dev/null @@ -1,111 +0,0 @@ ---- -paths: - - "pkg/agents/**" ---- - -# Sub-Agent Development Rules - -## Package Structure - -Each sub-agent lives under `pkg/agents//` with the following files: - -``` -pkg/agents// -├── factory.go # AgentFactory implementation, CLI flags, Configure() -├── agent.go # Agent core, SubAgent creation, middleware -├── tool.go # Internal tools (API calls) -├── auth.go # Authentication logic (if needed) -├── prompt.go # System prompt builder -├── extract.go # Record extraction from session history -├── export_test.go # Test-only exports -├── prompt/ -│ ├── system.md # System prompt template -│ └── extract.md # Extraction prompt template -└── README.md # Configuration guide -``` - -## Factory Pattern - -- Implement `agents.AgentFactory` interface -- CLI flags: `--agent--`, env vars: `WARREN_AGENT__` -- `Configure()` returns `(nil, nil)` when required credentials are not set (skip registration) -- Log configuration at INFO level with credentials summary (never log secrets, only their length) - -## Middleware Pattern (agent.go) - -All sub-agents follow this middleware flow: -1. **Pre-execution**: Memory search via `memoryService.SearchAndSelectMemories` -2. **Execution**: Internal agent processes the query -3. **Post-execution**: Memory save, record extraction, internal field cleanup - -## msg.Trace Usage - -Use `msg.Trace` extensively to show real-time progress to the user in Slack threads. This is critical for observability. - -### Required Trace Points in tool.go - -- **Before each API call**: What operation is starting (with key parameters like filter/query) - - `msg.Trace(ctx, "🔍 Searching alerts (filter: \`%s\`)", filter)` -- **After successful API call**: Confirmation with result summary - - `msg.Trace(ctx, "✅ Alert search completed")` -- **On API failure**: Error details - - `msg.Trace(ctx, "❌ Alert search failed: %v", err)` -- **On retry**: What's being retried and why - - `msg.Trace(ctx, "🔄 Received 401, refreshing token and retrying...")` -- **For async operations**: Progress updates (e.g., polling status) - - `msg.Trace(ctx, "⏳ Event search job created (job_id: \`%s\`), polling...")` - - `msg.Trace(ctx, "⚠️ Event search reached poll limit, returning %d partial results")` - -### Required Trace Points in agent.go - -- **On request received**: Show the incoming query - - `msg.Trace(ctx, "🦅 *[Falcon Agent]* Request: \`%s\`", request)` - -### Emoji Convention - -| Emoji | Usage | -|-------|-------| -| 🔍 | Search/query operations | -| 📋 | Retrieving details by ID | -| 📊 | Metrics/scores retrieval | -| ✅ | Successful completion | -| ❌ | Failure/error | -| 🔄 | Retry/refresh | -| ⏳ | Async waiting/polling | -| ⚠️ | Warning/partial results | - -## Error Handling - -### Authentication Errors -- Log token refresh failures at WARN level with status code and response body -- Accept both HTTP 200 and 201 for OAuth token responses (CrowdStrike returns 201) -- Log API call failures at WARN level with status, path, and response body - -### Sub-Agent Error Detection -- When record extraction yields 0 records AND the response text contains error indicators (auth failures, connectivity issues), return an error instead of `"status": "success"` with empty records -- This prevents silent failures from being treated as "no data found" - -## HTTP Client Best Practices - -- URL-encode query parameter values (`url.QueryEscape`) -- Use `safe.Close(ctx, resp.Body)` for response body cleanup -- Implement 401 retry: clear cached token, retry once -- Log all API request paths at DEBUG level for traceability - -## Testing - -- Use `export_test.go` for exposing test-only wrappers -- Use `httptest.NewServer` for mocking external APIs -- Include a spec count test to catch accidental tool additions/removals -- Test error cases (auth failure, API errors, 401 retry) - -## Registration - -Add the factory to `pkg/agents/agents.go`: -```go -var All = []AgentFactory{ - &slack.Factory{}, - &bigquery.Factory{}, - &falcon.Factory{}, // Add new agent here -} -``` diff --git a/README.md b/README.md index b3ba32be..1fe9d5fc 100644 --- a/README.md +++ b/README.md @@ -33,20 +33,21 @@ This is not a generic AI agent with security tools bolted on. Warren is purpose- ## How It Works -### Slack-Based Multi-Agent Investigation +### Slack-Based AI Investigation -Warren operates as a **Slack-native multi-agent system**. When an alert arrives, it is posted to a Slack channel with AI-generated analysis. Team members interact with Warren directly in Slack threads — `@warren` triggers an investigation agent that can delegate work to specialized sub-agents in parallel: +Warren operates as a **Slack-native investigation system**. When an alert arrives, it is posted to a Slack channel with AI-generated analysis. Team members interact with Warren directly in Slack threads — `@warren` triggers an investigation agent that can call specialized tools as it works: ``` User asks @warren in Slack thread - └─ Orchestrator Agent - ├─ BigQuery Agent → query audit logs, access patterns - ├─ Falcon Agent → pull EDR endpoint data from CrowdStrike - ├─ Slack Agent → search related conversations - └─ Direct tools → VirusTotal, OTX, Shodan, AbuseIPDB, URLScan + └─ Investigation Agent + ├─ BigQuery → query audit logs, access patterns + ├─ Falcon → pull EDR endpoint data from CrowdStrike + ├─ Slack → search related conversations + ├─ Jira → correlate with tracked tickets + └─ Threat intel → VirusTotal, OTX, Shodan, AbuseIPDB, URLScan ``` -Each sub-agent autonomously decides what queries to run and how to interpret results. Real-time progress traces in the Slack thread show what the agent is doing as it works. +The agent autonomously decides which tools to call and how to interpret results. Real-time progress traces in the Slack thread show what the agent is doing as it works.

Slack integration with interactive investigation @@ -126,15 +127,12 @@ Visit http://127.0.0.1:8080 to access the dashboard. - [**GitHub App**](./pkg/tool/github/README.md) — code search, issue search, file content retrieval, commit history, file blame - [**Microsoft Intune**](./pkg/tool/intune/README.md) — device compliance status, sign-in history +- [**CrowdStrike Falcon**](./pkg/tool/falcon/README.md) — EDR incidents, alerts, behaviors, devices, CrowdScores, and telemetry events - [**Slack Message Search**](./pkg/tool/slack/README.md) — search workspace messages for context +- [**Jira**](./pkg/tool/jira/README.md) — list projects, search issues via JQL, fetch issue content +- [**BigQuery**](./doc/reference/configuration.md#bigquery-tool) — query security log data with SQL and runbooks - [**WebFetch**](./pkg/tool/webfetch/README.md) — fetch a web page and return clean Markdown with indirect-prompt-injection screening -### Sub-Agents - -- [**BigQuery Agent**](./pkg/agents/bigquery/README.md) — query security log data via natural language -- [**CrowdStrike Falcon Agent**](./pkg/agents/falcon/README.md) — query EDR incidents, alerts, and endpoint events -- [**Slack Search Agent**](./pkg/agents/slack/README.md) — search and summarize Slack conversations - ### Collaboration & UI - **Slack** — native bot with interactive buttons, thread-based investigation, real-time progress traces diff --git a/doc/concepts.md b/doc/concepts.md index 22bf1daa..258ccba3 100644 --- a/doc/concepts.md +++ b/doc/concepts.md @@ -225,14 +225,3 @@ Slack command: Limitations: - Unbound alert processing is capped at 100 alerts per run - Each consolidation group contains 2-10 alerts - -### Agent Memory - -Warren's AI agents can learn from previous investigations through agent memory: - -- **Quality Scoring**: Memories rated from -10 (harmful) to +10 (helpful) -- **Adaptive Search**: Re-ranks memories by similarity (50%) + quality (30%) + recency (20%) -- **LLM Feedback**: Automatically evaluates memory usefulness after each agent execution -- **Conservative Pruning**: Strict deletion criteria to preserve valuable memories - -Agent memory is enabled automatically for sub-agents (BigQuery, Falcon, Slack). Memories are scoped per agent and accumulate over time. diff --git a/doc/migration/v0.18.0.md b/doc/migration/v0.18.0.md new file mode 100644 index 00000000..ead5472f --- /dev/null +++ b/doc/migration/v0.18.0.md @@ -0,0 +1,73 @@ +# Migration Guide: v0.18.0 + +This guide covers the tool-related changes in Warren v0.18.0. The release unifies all external integrations under the `pkg/tool/*` path (backed by [`github.com/gollem-dev/tools`](https://github.com/gollem-dev/tools)) and removes the legacy `pkg/agents` "sub-agent" framework. + +## Overview + +Before v0.18.0, Warren exposed two parallel mechanisms for external integrations: + +- **Tools** (`pkg/tool/*`) — single-call functions invoked directly by the investigation agent. +- **Sub-agents** (`pkg/agents/*`) — BigQuery, CrowdStrike Falcon, and Slack, each wired through a separate `WARREN_AGENT_*` flag namespace. + +v0.18.0 retires the sub-agent mechanism entirely. CrowdStrike Falcon and Slack are now plain tools, BigQuery already had a fully featured tool equivalent, and a new **Jira** tool is added. The change is **configuration-only** for operators — no data migration is required. + +## Breaking Changes Summary + +| Area | Removed / Changed | Replacement | +|------|-------------------|-------------| +| CrowdStrike Falcon | `WARREN_AGENT_FALCON_CLIENT_ID` / `_CLIENT_SECRET` / `_BASE_URL` (and `--agent-falcon-*` flags) | `WARREN_FALCON_CLIENT_ID` / `_CLIENT_SECRET` / `_BASE_URL` (`--falcon-*`) | +| Slack search | `WARREN_AGENT_SLACK_USER_TOKEN` (`--agent-slack-user-token`) | `WARREN_SLACK_TOOL_USER_TOKEN` (`--slack-tool-user-token`) — the existing Slack tool | +| BigQuery | `WARREN_AGENT_BIGQUERY_*` (`--agent-bigquery-*`) | `WARREN_BIGQUERY_*` (`--bigquery-*`) — the existing BigQuery tool | +| Jira | _(new)_ | `WARREN_JIRA_BASE_URL` / `_USER_EMAIL` / `_API_TOKEN` (`--jira-*`) | +| Internals | `pkg/agents` package; the chat "sub-agent" budget concept (`subAgentToolNames`, shared-tracker context plumbing) | none — tools are budgeted via the standard per-task tracker | + +There are **no Firestore, Cloud Storage, or policy changes** in this release. Pre-existing alert / triage / enrichment policies are unaffected. + +## Environment Variable Migration + +Rename the variables you currently set. The semantics are unchanged; only the namespace differs. + +| Old (≤ v0.17.x) | New (v0.18.0) | +|-----------------|---------------| +| `WARREN_AGENT_FALCON_CLIENT_ID` | `WARREN_FALCON_CLIENT_ID` | +| `WARREN_AGENT_FALCON_CLIENT_SECRET` | `WARREN_FALCON_CLIENT_SECRET` | +| `WARREN_AGENT_FALCON_BASE_URL` | `WARREN_FALCON_BASE_URL` | +| `WARREN_AGENT_SLACK_USER_TOKEN` | `WARREN_SLACK_TOOL_USER_TOKEN` | +| `WARREN_AGENT_BIGQUERY_CONFIG` | `WARREN_BIGQUERY_CONFIG` | +| `WARREN_AGENT_BIGQUERY_PROJECT_ID` | `WARREN_BIGQUERY_PROJECT_ID` | +| `WARREN_AGENT_BIGQUERY_SCAN_SIZE_LIMIT` | `WARREN_BIGQUERY_SCAN_LIMIT` | +| `WARREN_AGENT_BIGQUERY_RUNBOOK_DIR` | `WARREN_BIGQUERY_RUNBOOK_PATH` | +| `WARREN_AGENT_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT` | `WARREN_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT` | + +> The old `WARREN_AGENT_*` variables are **no longer read**. Set the new names before deploying v0.18.0, or the corresponding integration will be silently disabled (Warren skips any tool whose required settings are absent). + +### BigQuery notes + +The BigQuery tool is a superset of the former sub-agent: it supports the same runbooks (`WARREN_BIGQUERY_RUNBOOK_PATH` + the `get_runbook_entry` function), service-account impersonation, scan limit, and config files, and additionally supports GCS result storage (`WARREN_BIGQUERY_STORAGE_BUCKET` / `_STORAGE_PREFIX`), a query timeout (`WARREN_BIGQUERY_TIMEOUT`), and explicit credentials (`WARREN_BIGQUERY_CREDENTIALS`). No functionality is lost. + +The config-file format is per-table (one table per `--bigquery-config` file; pass the flag multiple times for several tables). `scan_size_limit` / `query_timeout` are no longer config-file keys — use the `--bigquery-scan-limit` / `--bigquery-timeout` flags instead. See [configuration.md → BigQuery Tool](../reference/configuration.md#bigquery-tool). + +## New: Jira tool + +v0.18.0 adds a read-only Jira Cloud tool (list projects, search issues via JQL, fetch issue content). Configure it with: + +```bash +export WARREN_JIRA_BASE_URL="https://your-domain.atlassian.net" +export WARREN_JIRA_USER_EMAIL="bot@example.com" +export WARREN_JIRA_API_TOKEN="" +``` + +Create the API token at [Atlassian account → Security → API tokens](https://id.atlassian.com/manage-profile/security/api-tokens). See [pkg/tool/jira/README.md](../../pkg/tool/jira/README.md). + +## Dependency updates + +The bundled `github.com/gollem-dev/tools` modules are updated to their latest versions (notably `slack`, `github`, and `webfetch` to v0.2.0, which add functions such as `slack_get_messages`, `github_get_issue`, and `github_get_pull_request`). No operator action is required. + +## Migration Steps + +1. **Rename environment variables** per the table above in your deployment (Cloud Run env vars / Secret Manager bindings / `.env`). +2. **(Optional) Add the Jira tool** settings if you want Jira correlation. +3. **Deploy v0.18.0.** On start-up, confirm each expected integration is active. A tool whose required settings are missing is skipped silently, so a missed rename surfaces as a quietly disabled tool rather than an error. +4. **Remove the old `WARREN_AGENT_*` variables** once the deployment is verified. + +No data migration jobs are needed for this release. diff --git a/doc/operation/alert-investigation.md b/doc/operation/alert-investigation.md index 34927560..18aedd9e 100644 --- a/doc/operation/alert-investigation.md +++ b/doc/operation/alert-investigation.md @@ -41,7 +41,7 @@ Click **[Salvage]** on a ticket to find similar unbound alerts using AI similari ### Talking to Warren -Mention `@warren` in a ticket's Slack thread to start an AI-powered investigation. Warren will analyze the ticket's alerts using available tools and sub-agents. +Mention `@warren` in a ticket's Slack thread to start an AI-powered investigation. Warren will analyze the ticket's alerts using the available tools. ``` @warren Check if the source IPs in this ticket are malicious @@ -142,6 +142,9 @@ Warren's AI agent has access to the following tools during investigation. Tools |------|-------------|---------| | GitHub | Code search, issue search, file content, commit history, file blame | [README](../../pkg/tool/github/README.md) | | Microsoft Intune | Device compliance, sign-in history | [README](../../pkg/tool/intune/README.md) | +| CrowdStrike Falcon | EDR incidents, alerts, behaviors, devices, CrowdScores, telemetry events | [README](../../pkg/tool/falcon/README.md) | +| BigQuery | Query security log data with SQL and runbooks | [configuration](../reference/configuration.md#bigquery-tool) | +| Jira | List projects, search issues via JQL, fetch issue content | [README](../../pkg/tool/jira/README.md) | ### Other Tools @@ -151,16 +154,6 @@ Warren's AI agent has access to the following tools during investigation. Tools | Knowledge | Save/retrieve investigation knowledge | [README](../../pkg/tool/knowledge/README.md) | | WebFetch | Fetch a web page and return Markdown (with indirect-prompt-injection screening) | [README](../../pkg/tool/webfetch/README.md) | -## Sub-Agents - -Sub-agents are specialized AI agents that handle complex, multi-step operations. The main agent delegates to them when needed. - -| Sub-Agent | Description | Details | -|-----------|-------------|---------| -| BigQuery Agent | Query security log data via natural language | [README](../../pkg/agents/bigquery/README.md) | -| CrowdStrike Falcon Agent | Query EDR incidents, alerts, and events | [README](../../pkg/agents/falcon/README.md) | -| Slack Search Agent | Search Slack messages for context | [README](../../pkg/agents/slack/README.md) | - ## Chat Strategy Warren uses the `aster` chat execution strategy by default. It parallelizes independent tasks for faster multi-tool investigations. diff --git a/doc/reference/configuration.md b/doc/reference/configuration.md index d6159628..861db044 100644 --- a/doc/reference/configuration.md +++ b/doc/reference/configuration.md @@ -165,6 +165,42 @@ Tools are automatically enabled when their API key is configured. Missing keys a | `WARREN_INTUNE_CLIENT_SECRET` | `--intune-client-secret` | - | Azure AD client secret | | `WARREN_INTUNE_BASE_URL` | `--intune-base-url` | `https://graph.microsoft.com/v1.0` | Graph API base URL | +### CrowdStrike Falcon + +Read-only CrowdStrike Falcon queries (incidents, alerts, behaviors, devices, CrowdScores, EDR telemetry events). Enabled when both client ID and secret are set. See [pkg/tool/falcon/README.md](../../pkg/tool/falcon/README.md). + +| Environment Variable | CLI Flag | Default | Description | +|---|---|---|---| +| `WARREN_FALCON_CLIENT_ID` | `--falcon-client-id` | - | CrowdStrike API client ID | +| `WARREN_FALCON_CLIENT_SECRET` | `--falcon-client-secret` | - | CrowdStrike API client secret | +| `WARREN_FALCON_BASE_URL` | `--falcon-base-url` | `https://api.crowdstrike.com` | API base URL (match your CrowdStrike cloud region) | + +### Jira + +Read-only Jira Cloud access (list projects, search issues via JQL, fetch issue content). Enabled when all three values are set. See [pkg/tool/jira/README.md](../../pkg/tool/jira/README.md). + +| Environment Variable | CLI Flag | Default | Description | +|---|---|---|---| +| `WARREN_JIRA_BASE_URL` | `--jira-base-url` | - | Jira site URL (`https://.atlassian.net`) | +| `WARREN_JIRA_USER_EMAIL` | `--jira-user-email` | - | Jira account email for Basic auth | +| `WARREN_JIRA_API_TOKEN` | `--jira-api-token` | - | Jira Cloud API token | + +### BigQuery Tool + +Query security log data with SQL, table schema introspection, and SQL runbooks. Enabled when `WARREN_BIGQUERY_PROJECT_ID` is set together with at least one `WARREN_BIGQUERY_CONFIG` file. + +| Environment Variable | CLI Flag | Default | Description | +|---|---|---|---| +| `WARREN_BIGQUERY_PROJECT_ID` | `--bigquery-project-id` | - | Google Cloud project ID | +| `WARREN_BIGQUERY_CONFIG` | `--bigquery-config` | - | Configuration YAML file path(s) (multiple allowed) | +| `WARREN_BIGQUERY_CREDENTIALS` | `--bigquery-credentials` | - | Path to Google Cloud credentials JSON file (optional) | +| `WARREN_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT` | `--bigquery-impersonate-service-account` | - | Service account email for impersonation | +| `WARREN_BIGQUERY_RUNBOOK_PATH` | `--bigquery-runbook-path` | - | SQL runbook file or directory path(s) (multiple allowed) | +| `WARREN_BIGQUERY_STORAGE_BUCKET` | `--bigquery-storage-bucket` | - | GCS bucket for storing query results | +| `WARREN_BIGQUERY_STORAGE_PREFIX` | `--bigquery-storage-prefix` | - | GCS object path prefix for stored results | +| `WARREN_BIGQUERY_TIMEOUT` | `--bigquery-timeout` | `5m` | Query execution timeout | +| `WARREN_BIGQUERY_SCAN_LIMIT` | `--bigquery-scan-limit` | `10GB` | Per-query scan size limit | + ### Slack Message Search (Tool) | Environment Variable | CLI Flag | Default | Description | @@ -238,32 +274,6 @@ LLM analysis disabled (HITL approval gates every call): warren chat # no --webfetch-llm-* flags ``` -## Sub-Agent Configuration - -### BigQuery Agent - -| Environment Variable | CLI Flag | Default | Description | -|---|---|---|---| -| `WARREN_AGENT_BIGQUERY_CONFIG` | `--agent-bigquery-config` | - | YAML config file path (required) | -| `WARREN_AGENT_BIGQUERY_PROJECT_ID` | `--agent-bigquery-project-id` | - | GCP project ID | -| `WARREN_AGENT_BIGQUERY_SCAN_SIZE_LIMIT` | `--agent-bigquery-scan-size-limit` | - | Override scan limit | -| `WARREN_AGENT_BIGQUERY_RUNBOOK_DIR` | `--agent-bigquery-runbook-dir` | - | SQL runbook paths (multiple allowed) | -| `WARREN_AGENT_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT` | `--agent-bigquery-impersonate-service-account` | - | Service account to impersonate | - -### CrowdStrike Falcon Agent - -| Environment Variable | CLI Flag | Default | Description | -|---|---|---|---| -| `WARREN_AGENT_FALCON_CLIENT_ID` | `--agent-falcon-client-id` | - | CrowdStrike API client ID | -| `WARREN_AGENT_FALCON_CLIENT_SECRET` | `--agent-falcon-client-secret` | - | CrowdStrike API client secret | -| `WARREN_AGENT_FALCON_BASE_URL` | `--agent-falcon-base-url` | `https://api.crowdstrike.com` | API base URL | - -### Slack Search Agent - -| Environment Variable | CLI Flag | Default | Description | -|---|---|---|---| -| `WARREN_AGENT_SLACK_USER_TOKEN` | `--agent-slack-user-token` | - | Slack User token (`xoxp-...`) with `search:read` scope | - ## Command-Specific Options ### `warren serve` @@ -305,20 +315,33 @@ local: args: ["--mcp-mode"] ``` -### BigQuery Agent Configuration +### BigQuery Tool Configuration + +Each `--bigquery-config` file describes one table the tool may query. Pass the flag multiple times (or set multiple `WARREN_BIGQUERY_CONFIG` entries) to register several tables. Scan limit and timeout are CLI flags (`--bigquery-scan-limit`, `--bigquery-timeout`), not config-file fields. ```yaml -tables: - - project_id: my-project - dataset_id: security_logs - table_id: events - description: Security event logs - -scan_size_limit: "10GB" -query_timeout: 5m +dataset_id: security_logs +table_id: events +description: Security event logs +columns: + - name: timestamp + description: Event timestamp + value_example: "2024-03-20T12:00:00Z" + type: TIMESTAMP + - name: src_ip + description: Source IP address + value_example: "192.168.1.1" + type: STRING +partitioning: + field: timestamp + type: time + time_unit: daily +# Optional: SQL runbooks (also settable via --bigquery-runbook-path) +runbook_paths: + - ./runbooks/failed_logins.sql ``` -See [BigQuery Agent README](../../pkg/agents/bigquery/README.md) for detailed configuration. +See the [BigQuery Tool](#bigquery-tool) settings above for the full flag/environment-variable list. ## Asynchronous Alert Processing diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index bf93014f..d9cab2fe 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -76,7 +76,7 @@ importers: version: 4.1.0 dompurify: specifier: ^3.4.9 - version: 3.4.10 + version: 3.4.11 graphql: specifier: ^16.13.2 version: 16.13.2 @@ -2093,8 +2093,8 @@ packages: resolution: {integrity: sha512-35mSku4ZXK0vfCuHEDAwt55dg2jNajHZ1odvF+8SSr82EsZY4QmXfuWso8oEd8zRhVObSN18aM0CjSdoBX7zIw==} engines: {node: '>=0.10.0'} - dompurify@3.4.10: - resolution: {integrity: sha512-0xzNv0e7oYC6yyuOGZIABPM4qtg3QxLFniDNPP4ZP90wR8Yq3zgwpRbrNiT4N3IKqDbbYFEJLV+JWEs19aZ//w==} + dompurify@3.4.11: + resolution: {integrity: sha512-zhlUV12GsaRzMsf9q5M254YhA4+VuF0fG+QFqu6aYpoGlKtz+w8//jBcGVYBgQkR5GHjUomejY84AV+/uPbWdw==} dot-case@3.0.4: resolution: {integrity: sha512-Kv5nKlh6yRrdrGvxeJ2e5y2eRUpkUosIW4A2AS38zwSz27zu7ufDwQPi5Jhs3XAlGNetl3bmnGhQsMtkKJnj3w==} @@ -6084,7 +6084,7 @@ snapshots: dependencies: esutils: 2.0.3 - dompurify@3.4.10: + dompurify@3.4.11: optionalDependencies: '@types/trusted-types': 2.0.7 diff --git a/go.mod b/go.mod index 274f68cc..94090e9a 100644 --- a/go.mod +++ b/go.mod @@ -15,15 +15,17 @@ require ( github.com/gollem-dev/gollem v0.26.0 github.com/gollem-dev/tools/abusech v0.1.0 github.com/gollem-dev/tools/bigquery v0.1.0 - github.com/gollem-dev/tools/github v0.1.0 + github.com/gollem-dev/tools/falcon v0.1.0 + github.com/gollem-dev/tools/github v0.2.0 github.com/gollem-dev/tools/intune v0.1.0 github.com/gollem-dev/tools/ipdb v0.1.0 + github.com/gollem-dev/tools/jira v0.1.0 github.com/gollem-dev/tools/otx v0.1.0 github.com/gollem-dev/tools/shodan v0.1.0 - github.com/gollem-dev/tools/slack v0.1.0 + github.com/gollem-dev/tools/slack v0.2.0 github.com/gollem-dev/tools/urlscan v0.1.0 github.com/gollem-dev/tools/vt v0.1.0 - github.com/gollem-dev/tools/webfetch v0.1.0 + github.com/gollem-dev/tools/webfetch v0.2.0 github.com/gollem-dev/tools/whois v0.1.0 github.com/google/go-github/v74 v74.0.0 github.com/google/uuid v1.6.0 @@ -210,7 +212,7 @@ require ( golang.org/x/mod v0.35.0 // indirect golang.org/x/net v0.54.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sync v0.20.0 // indirect + golang.org/x/sync v0.21.0 // indirect golang.org/x/sys v0.44.0 // indirect golang.org/x/telemetry v0.0.0-20260409153401-be6f6cb8b1fa // indirect golang.org/x/term v0.43.0 // indirect diff --git a/go.sum b/go.sum index 42c25e51..1ed3ef95 100644 --- a/go.sum +++ b/go.sum @@ -224,24 +224,28 @@ github.com/gollem-dev/tools/abusech v0.1.0 h1:0EvwiABq6585q7sarWYkpNRyDvMIgIAYxA github.com/gollem-dev/tools/abusech v0.1.0/go.mod h1:4xn+UrOECXLWiBTPSgZoOOVrsqo4P2wZ0u/umtbsRFc= github.com/gollem-dev/tools/bigquery v0.1.0 h1:TL3v6IVrozstX/LbFsID4NFiP/I/OUDXojYzv+H2by8= github.com/gollem-dev/tools/bigquery v0.1.0/go.mod h1:H+sKbHbwEpCP1+FCaToZ5ftiOpMMO1Wbi2DaL3E7AqE= -github.com/gollem-dev/tools/github v0.1.0 h1:HK3CKjYrC62gqhw7v+Oi39ai3yxWbkEqjHlN7hD42X4= -github.com/gollem-dev/tools/github v0.1.0/go.mod h1:LHt/fbIzFTU+eU2U5uXSgTTu6Yoco2Lkzc9DRMaT/zI= +github.com/gollem-dev/tools/falcon v0.1.0 h1:0VTF2KJUg3K3ICPPt+pXTjBF3jmYUUC4zclFrA5qRZQ= +github.com/gollem-dev/tools/falcon v0.1.0/go.mod h1:134+tUzR8Jv+P9J0Yo+ZVdNzoh5ckasS14OYQf1t/tQ= +github.com/gollem-dev/tools/github v0.2.0 h1:Na9LZuaer7HwA3m36ySUuRK1Ru8VFYrWI00yFMzcwSE= +github.com/gollem-dev/tools/github v0.2.0/go.mod h1:LHt/fbIzFTU+eU2U5uXSgTTu6Yoco2Lkzc9DRMaT/zI= github.com/gollem-dev/tools/intune v0.1.0 h1:Qe41jsuRDvQrXsSmk91q15hBCT68bhNacupKIeU2vhg= github.com/gollem-dev/tools/intune v0.1.0/go.mod h1:WdgquLJMs6tPFhzISArygZ+AGr2KvpOPulhtp1gp+oY= github.com/gollem-dev/tools/ipdb v0.1.0 h1:Df1j1PwcIQvVFvqI6LP09EtRqHAj10pLzPnQk8EKfoU= github.com/gollem-dev/tools/ipdb v0.1.0/go.mod h1:Zz0w9yXKpBcimB+ut4I3jlRpXWV+Qs6QQqb2G6MnpuE= +github.com/gollem-dev/tools/jira v0.1.0 h1:IRWE0Mk3Tvv6VB2bZB3GUy616X9U0gkTuVQ6Xjq4oh8= +github.com/gollem-dev/tools/jira v0.1.0/go.mod h1:S6s4X2ejygepdPmC6QDX3SKnwMaR3X4Xf/8Ip0at2mA= github.com/gollem-dev/tools/otx v0.1.0 h1:+mKhHd+gikhh3dvNSmyvbSaiACIIMaZe1UoXwzlk2+8= github.com/gollem-dev/tools/otx v0.1.0/go.mod h1:8cuXN15Bq/ILXN4CPACpCFIabe52pz3O76Z1iePHQEk= github.com/gollem-dev/tools/shodan v0.1.0 h1:eWmb426Kg1koYlHCI11sJ6BnZUxcUKbdd0lEoWdKr4k= github.com/gollem-dev/tools/shodan v0.1.0/go.mod h1:AvuHaiYDPmJMJiQOK7fClkJ00pMH+RscS6VjU38/zos= -github.com/gollem-dev/tools/slack v0.1.0 h1:8dmzH5c2jM1WZvwXeLGdvWE/TBQNkp6sWfM31byWwhQ= -github.com/gollem-dev/tools/slack v0.1.0/go.mod h1:g415ENkXWQRDToVCviuSHQ2BMiz0omrIQsjbmQwOahM= +github.com/gollem-dev/tools/slack v0.2.0 h1:looGoq8rfSVQbFUoa5YWeSFHUj4jpUU6d514ZlxBSMw= +github.com/gollem-dev/tools/slack v0.2.0/go.mod h1:98kLjH8YLKX6Q4VhFBb9x4GiItjgiHP4T/ANYK9BEps= github.com/gollem-dev/tools/urlscan v0.1.0 h1:AxMMuy9bsXFUSCWXaIWBesgKAIVYEGl2/Cjt3BvnTgg= github.com/gollem-dev/tools/urlscan v0.1.0/go.mod h1:cZ7bUqclkpFCL2Z5pFKLWd6KWTWzf+iy69NvnFoewIo= github.com/gollem-dev/tools/vt v0.1.0 h1:X2T0E1HT5aoU4BK6B/ZITwRvOLMiq6ttLAgHdVfPeEc= github.com/gollem-dev/tools/vt v0.1.0/go.mod h1:iV1dpMTWMMQSZp7+dbUVzgb4Rd6eIJHLSxPEa8W8mKo= -github.com/gollem-dev/tools/webfetch v0.1.0 h1:jDcY+oQGH6zoKL/tBSzzgin0R9t60UY6xn3jwFMj/9A= -github.com/gollem-dev/tools/webfetch v0.1.0/go.mod h1:zTCZBmDrF3gFJhzES/gaL+pnYHa6BIAglc8E0GR8rrY= +github.com/gollem-dev/tools/webfetch v0.2.0 h1:DbuKdk1ODPL0fvO+d/bEm2s97W6bRf/MlVhEa9rwaGM= +github.com/gollem-dev/tools/webfetch v0.2.0/go.mod h1:zTCZBmDrF3gFJhzES/gaL+pnYHa6BIAglc8E0GR8rrY= github.com/gollem-dev/tools/whois v0.1.0 h1:YePfmuLx44mHFaGi7IOGKATKme81AE+dHqisdALGrfU= github.com/gollem-dev/tools/whois v0.1.0/go.mod h1:U/iXnDoKtJmWZA7ma9bcrOVV4W2oaLwyuRP+vd+6AKY= github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= @@ -498,8 +502,8 @@ golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w= golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ= golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= diff --git a/pkg/agents/agents.go b/pkg/agents/agents.go deleted file mode 100644 index 06fdd47e..00000000 --- a/pkg/agents/agents.go +++ /dev/null @@ -1,48 +0,0 @@ -package agents - -import ( - "context" - - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/agents/bigquery" - "github.com/secmon-lab/warren/pkg/agents/falcon" - "github.com/secmon-lab/warren/pkg/agents/slack" - "github.com/secmon-lab/warren/pkg/domain/interfaces" - "github.com/urfave/cli/v3" -) - -// All is the list of all available agent factories. -// To add a new agent, simply append its Factory to this list. -var All = []ToolSetFactory{ - &bigquery.Factory{}, - &falcon.Factory{}, - &slack.Factory{}, -} - -// AllFlags returns CLI flags for all registered agents. -func AllFlags() []cli.Flag { - var flags []cli.Flag - for _, factory := range All { - flags = append(flags, factory.Flags()...) - } - return flags -} - -// ConfigureAll initializes all configured agents and returns a slice of ToolSets. -// Agents that are not configured will return nil and be skipped. -func ConfigureAll(ctx context.Context) ([]interfaces.ToolSet, error) { - var toolSets []interfaces.ToolSet - - for _, factory := range All { - ts, err := factory.Configure(ctx) - if err != nil { - return nil, goerr.Wrap(err, "failed to configure agent") - } - - if ts != nil { - toolSets = append(toolSets, ts) - } - } - - return toolSets, nil -} diff --git a/pkg/agents/bigquery/README.md b/pkg/agents/bigquery/README.md deleted file mode 100644 index d5039798..00000000 --- a/pkg/agents/bigquery/README.md +++ /dev/null @@ -1,259 +0,0 @@ -# BigQuery Agent Configuration - -## Configuration File - -Create a YAML file defining BigQuery tables: - -```yaml -tables: - - project_id: my-security-project - dataset_id: cloudtrail_logs - table_id: events - description: CloudTrail events with user activity, API calls, and resource changes - -scan_size_limit: "10GB" # Required: Maximum bytes scanned per query -query_timeout: 5m # Optional: Default 5m -``` - -## Configuration Fields - -### Tables Structure - -```yaml -tables: - - project_id: # Required: GCP project ID - dataset_id: # Required: Dataset ID - table_id: # Required: Table ID - description: # Recommended: Detailed table description -``` - -### Global Settings - -- `scan_size_limit` (required): Maximum bytes scanned per query - - Format: `"1GB"`, `"10GB"`, `"1TB"`, etc. -- `query_timeout` (optional): Query timeout duration - - Format: `"5m"`, `"30s"`, etc. - - Default: `5m` - -## Table Descriptions - -**Good descriptions improve agent performance significantly.** - -Good example: -```yaml -description: | - AWS CloudTrail events containing: - - Authentication events (ConsoleLogin, AssumeRole) - - API calls (eventName, requestParameters, responseElements) - - Failed attempts (errorCode, errorMessage) - - Fields: sourceIPAddress, userAgent, mfaUsed - - Partitioned by event_date (last 90 days) -``` - -Bad example: -```yaml -description: CloudTrail data # Too vague -``` - -## SQL Runbooks - -SQL runbooks are pre-written query templates for common investigation patterns. - -### Runbook File Format - -Create `.sql` files with metadata in comments: - -```sql --- Title: Failed Login Investigation --- Description: Query to find failed login attempts by IP address -SELECT - timestamp, - user_email, - source_ip, - error_code -FROM `project.dataset.auth_logs` -WHERE - event_type = 'login_failed' - AND timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 24 HOUR) -ORDER BY timestamp DESC -LIMIT 100 -``` - -**Metadata Format** (case insensitive): -- `-- Title: ` or `-- title: <title>` -- `-- Description: <description>` or `-- description: <description>` - -If no title is specified, the filename (without `.sql`) is used. - -### How Runbooks Work - -1. **Loading**: Runbooks are loaded at agent initialization -2. **Listing**: Agent's system prompt includes all runbook IDs, titles, and descriptions -3. **Retrieval**: Agent can use `get_runbook` tool to fetch full SQL content -4. **Adaptation**: Agent adapts runbook SQL for specific investigations - -## Deployment - -### Environment Variables - -```bash -export WARREN_AGENT_BIGQUERY_CONFIG="/path/to/config.yaml" -export WARREN_AGENT_BIGQUERY_PROJECT_ID="my-project" -export WARREN_AGENT_BIGQUERY_SCAN_SIZE_LIMIT="10GB" # Optional override -export WARREN_AGENT_BIGQUERY_RUNBOOK_DIR="/path/to/runbooks,/path/to/more" # Optional -``` - -### CLI Flags - -```bash -warren serve \ - --agent-bigquery-config=/path/to/config.yaml \ - --agent-bigquery-project-id=my-project \ - --agent-bigquery-runbook-dir=/path/to/runbooks \ - --agent-bigquery-runbook-dir=/path/to/more/runbooks # Can specify multiple times -``` - -### Cloud Run - -```bash -gcloud run services update warren \ - --set-env-vars="WARREN_AGENT_BIGQUERY_CONFIG=/app/config.yaml" \ - --set-env-vars="WARREN_AGENT_BIGQUERY_PROJECT_ID=my-project" -``` - -Note: Config file must be in the container image or mounted. - -## Authentication - -Uses [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials). - -### Local Development - -```bash -gcloud auth application-default login -``` - -### Service Account Impersonation - -You can configure the BigQuery agent to impersonate a service account for all BigQuery operations. This is useful when: -- Running with different permissions than the default credentials -- Implementing least-privilege access patterns -- Testing with specific service account permissions - -#### Configuration - -Set via CLI flag or environment variable: - -```bash -# CLI flag -warren serve \ - --agent-bigquery-config=/path/to/config.yaml \ - --agent-bigquery-impersonate-service-account=bigquery-reader@my-project.iam.gserviceaccount.com - -# Environment variable -export WARREN_AGENT_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT="bigquery-reader@my-project.iam.gserviceaccount.com" -warren serve --agent-bigquery-config=/path/to/config.yaml -``` - -#### Required Permissions - -The identity running Warren (ADC) must have the `roles/iam.serviceAccountTokenCreator` role on the target service account: - -```bash -gcloud iam service-accounts add-iam-policy-binding \ - bigquery-reader@my-project.iam.gserviceaccount.com \ - --member="<MEMBER>" \ - --role="roles/iam.serviceAccountTokenCreator" -``` - -The impersonated service account needs BigQuery permissions: - -```bash -gcloud projects add-iam-policy-binding $PROJECT_ID \ - --member="serviceAccount:bigquery-reader@my-project.iam.gserviceaccount.com" \ - --role="roles/bigquery.jobUser" - -gcloud projects add-iam-policy-binding $PROJECT_ID \ - --member="serviceAccount:bigquery-reader@my-project.iam.gserviceaccount.com" \ - --role="roles/bigquery.dataViewer" -``` - -### Service Account Permissions - -Required IAM roles for the service account (either default ADC or impersonated): -- `roles/bigquery.jobUser` -- `roles/bigquery.dataViewer` - -```bash -gcloud projects add-iam-policy-binding $PROJECT_ID \ - --member="serviceAccount:warren-service@${PROJECT_ID}.iam.gserviceaccount.com" \ - --role="roles/bigquery.jobUser" - -gcloud projects add-iam-policy-binding $PROJECT_ID \ - --member="serviceAccount:warren-service@${PROJECT_ID}.iam.gserviceaccount.com" \ - --role="roles/bigquery.dataViewer" -``` - -## Example Configurations - -### Minimal - -```yaml -tables: - - project_id: my-project - dataset_id: security_logs - table_id: events - description: Security event logs - -scan_size_limit: "1GB" -``` - -### Multiple Tables - -```yaml -tables: - - project_id: prod-logs - dataset_id: aws_cloudtrail - table_id: events_2024 - description: | - CloudTrail events from 2024: - - Authentication (ConsoleLogin, AssumeRole) - - S3 access (GetObject, PutObject) - - IAM changes - - Partitioned by event_date - - - project_id: prod-logs - dataset_id: gcp_audit - table_id: admin_activity - description: GCP admin activity logs - - - project_id: staging-logs - dataset_id: app_logs - table_id: authentication - description: App authentication events - -scan_size_limit: "5GB" -query_timeout: 3m -``` - -## Migration from Old Format - -If you're upgrading from the previous nested format, convert your configuration: - -**Old format (deprecated):** -```yaml -projects: - - id: my-project - datasets: - - id: security_logs - tables: - - id: events -``` - -**New format:** -```yaml -tables: - - project_id: my-project - dataset_id: security_logs - table_id: events -``` diff --git a/pkg/agents/bigquery/agent.go b/pkg/agents/bigquery/agent.go deleted file mode 100644 index db9c34c6..00000000 --- a/pkg/agents/bigquery/agent.go +++ /dev/null @@ -1,52 +0,0 @@ -package bigquery - -import ( - "context" - "strings" - - "github.com/gollem-dev/gollem" -) - -// toolSet implements interfaces.ToolSet by wrapping the internalTool -// and providing Warren-specific metadata for the planner. -type toolSet struct { - config *Config - tool *internalTool -} - -// ID implements interfaces.ToolSet. -func (ts *toolSet) ID() string { - return "bigquery" -} - -// Description implements interfaces.ToolSet. -// It dynamically includes available table names so the planner -// knows which tables this tool set has access to. -func (ts *toolSet) Description() string { - base := "Retrieve data from BigQuery tables. This tool ONLY extracts data records - it does NOT analyze or interpret the data. After receiving the data, YOU must analyze it yourself and provide a complete answer to the user based on the retrieved data. The tool handles table selection, query construction, and returns raw data records." - - if len(ts.config.Tables) > 0 { - var tables []string - for _, t := range ts.config.Tables { - tables = append(tables, strings.Join([]string{t.ProjectID, t.DatasetID, t.TableID}, ".")) - } - base += " Available tables: " + strings.Join(tables, ", ") + "." - } - - return base -} - -// Prompt implements interfaces.ToolSet. -func (ts *toolSet) Prompt(ctx context.Context) (string, error) { - return buildSystemPrompt(ts.config) -} - -// Specs implements gollem.ToolSet (delegated to internalTool). -func (ts *toolSet) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return ts.tool.Specs(ctx) -} - -// Run implements gollem.ToolSet (delegated to internalTool). -func (ts *toolSet) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - return ts.tool.Run(ctx, name, args) -} diff --git a/pkg/agents/bigquery/agent_test.go b/pkg/agents/bigquery/agent_test.go deleted file mode 100644 index 41952d37..00000000 --- a/pkg/agents/bigquery/agent_test.go +++ /dev/null @@ -1,99 +0,0 @@ -package bigquery_test - -import ( - "context" - "strings" - "testing" - - "github.com/m-mizutani/gt" - bqagent "github.com/secmon-lab/warren/pkg/agents/bigquery" -) - -func TestToolSet_ID(t *testing.T) { - ts := bqagent.NewToolSetForTest(&bqagent.Config{ - Tables: []bqagent.TableConfig{}, - ScanSizeLimit: 1000000, - }, "test-project", "") - - gt.V(t, ts.ID()).Equal("bigquery") -} - -func TestToolSet_Description(t *testing.T) { - t.Run("without tables", func(t *testing.T) { - ts := bqagent.NewToolSetForTest(&bqagent.Config{ - Tables: []bqagent.TableConfig{}, - ScanSizeLimit: 1000000, - }, "test-project", "") - - description := ts.Description() - gt.V(t, description).NotEqual("") - gt.True(t, strings.Contains(description, "BigQuery")) - // Should not contain "Available tables:" when no tables configured - gt.False(t, strings.Contains(description, "Available tables:")) - }) - - t.Run("with tables", func(t *testing.T) { - ts := bqagent.NewToolSetForTest(&bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "test-project", - DatasetID: "test-dataset", - TableID: "test-table", - Description: "Test table", - }, - }, - ScanSizeLimit: 1000000, - }, "test-project", "") - - description := ts.Description() - gt.True(t, strings.Contains(description, "Available tables:")) - gt.True(t, strings.Contains(description, "test-project.test-dataset.test-table")) - }) -} - -func TestToolSet_Prompt(t *testing.T) { - ts := bqagent.NewToolSetForTest(&bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "test-project", - DatasetID: "test-dataset", - TableID: "test-table", - Description: "Test table", - }, - }, - ScanSizeLimit: 1000000, - }, "test-project", "") - - ctx := context.Background() - prompt, err := ts.Prompt(ctx) - gt.NoError(t, err) - gt.True(t, len(prompt) > 0) -} - -func TestToolSet_Specs(t *testing.T) { - ts := bqagent.NewToolSetForTest(&bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "test-project", - DatasetID: "test-dataset", - TableID: "test-table", - Description: "Test table", - }, - }, - ScanSizeLimit: 1000000, - }, "test-project", "") - - ctx := context.Background() - specs, err := ts.Specs(ctx) - gt.NoError(t, err) - // Should have at least bigquery_query and bigquery_schema - gt.N(t, len(specs)).GreaterOrEqual(2) - - // Verify expected tool names - names := make(map[string]bool) - for _, s := range specs { - names[s.Name] = true - } - gt.True(t, names["bigquery_query"]) - gt.True(t, names["bigquery_schema"]) -} diff --git a/pkg/agents/bigquery/config.go b/pkg/agents/bigquery/config.go deleted file mode 100644 index 02fa0f47..00000000 --- a/pkg/agents/bigquery/config.go +++ /dev/null @@ -1,121 +0,0 @@ -package bigquery - -import ( - "context" - "os" - "path/filepath" - "time" - - "github.com/dustin/go-humanize" - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/domain/model/bigquery" - "github.com/secmon-lab/warren/pkg/domain/types" - "gopkg.in/yaml.v3" -) - -// Config represents BigQuery Agent configuration -type Config struct { - Tables []TableConfig `yaml:"tables"` - ScanSizeLimit uint64 `yaml:"-"` // Parsed from ScanSizeLimitStr - ScanSizeLimitStr string `yaml:"scan_size_limit"` - QueryTimeout time.Duration `yaml:"query_timeout"` // Timeout for waiting for BigQuery job completion (default: 5 minutes) - Runbooks map[types.RunbookID]*bigquery.RunbookEntry `yaml:"-"` // Loaded runbooks (not in YAML) -} - -// TableConfig represents a BigQuery table configuration -type TableConfig struct { - ProjectID string `yaml:"project_id"` - DatasetID string `yaml:"dataset_id"` - TableID string `yaml:"table_id"` - Description string `yaml:"description,omitempty"` -} - -// LoadConfig loads BigQuery Agent configuration from a YAML file -func LoadConfig(path string) (*Config, error) { - data, err := os.ReadFile(filepath.Clean(path)) // #nosec G304 - path is provided by user configuration - if err != nil { - return nil, goerr.Wrap(err, "failed to read config file", goerr.V("path", path)) - } - - var cfg Config - if err := yaml.Unmarshal(data, &cfg); err != nil { - return nil, goerr.Wrap(err, "failed to unmarshal config", goerr.V("path", path)) - } - - if err := cfg.setDefaults(); err != nil { - return nil, err - } - return &cfg, nil -} - -// LoadConfigWithRunbooks loads configuration and runbooks -func LoadConfigWithRunbooks(ctx context.Context, path string, runbookPaths []string) (*Config, error) { - cfg, err := LoadConfig(path) - if err != nil { - return nil, err - } - - if len(runbookPaths) > 0 { - if err := cfg.loadRunbooks(ctx, runbookPaths); err != nil { - return nil, goerr.Wrap(err, "failed to load runbooks") - } - } - - return cfg, nil -} - -// loadRunbooks loads runbooks from specified paths -func (c *Config) loadRunbooks(ctx context.Context, paths []string) error { - loader := newRunbookLoader(paths) - entries, err := loader.loadRunbooks(ctx) - if err != nil { - return goerr.Wrap(err, "failed to load runbook entries") - } - - c.Runbooks = make(map[types.RunbookID]*bigquery.RunbookEntry) - for _, entry := range entries { - c.Runbooks[entry.ID] = entry - } - - return nil -} - -// setDefaults sets default values for config fields if not specified -func (c *Config) setDefaults() error { - if c.QueryTimeout == 0 { - c.QueryTimeout = 5 * time.Minute - } - - // Parse scan size limit from string - if c.ScanSizeLimitStr != "" { - limit, err := ParseScanSizeLimit(c.ScanSizeLimitStr) - if err != nil { - return goerr.Wrap(err, "failed to parse scan_size_limit") - } - c.ScanSizeLimit = limit - } - - return nil -} - -// GetQueryTimeout returns the query timeout with fallback to default -func (c *Config) GetQueryTimeout() time.Duration { - if c.QueryTimeout == 0 { - return 5 * time.Minute - } - return c.QueryTimeout -} - -// ParseScanSizeLimit parses human-readable size string (e.g., "10GB") into bytes -func ParseScanSizeLimit(sizeStr string) (uint64, error) { - if sizeStr == "" { - return 0, nil - } - - bytes, err := humanize.ParseBytes(sizeStr) - if err != nil { - return 0, goerr.Wrap(err, "failed to parse scan size limit", goerr.V("size_str", sizeStr)) - } - - return bytes, nil -} diff --git a/pkg/agents/bigquery/config_test.go b/pkg/agents/bigquery/config_test.go deleted file mode 100644 index d4f97ce5..00000000 --- a/pkg/agents/bigquery/config_test.go +++ /dev/null @@ -1,128 +0,0 @@ -package bigquery_test - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/m-mizutani/gt" - "github.com/secmon-lab/warren/pkg/agents/bigquery" -) - -func TestLoadConfigWithRunbooks(t *testing.T) { - ctx := context.Background() - configPath := filepath.Join("testdata", "config_with_runbooks.yaml") - runbookDir := filepath.Join("testdata", "runbooks") - - // Load config with runbooks - cfg, err := bigquery.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - gt.V(t, cfg).NotNil() - - // Verify basic config loaded - gt.V(t, len(cfg.Tables)).Equal(1) - gt.V(t, cfg.Tables[0].TableID).Equal("test_table") - - // Verify runbooks loaded - gt.V(t, len(cfg.Runbooks) > 0).Equal(true) - - // Verify runbook contents - var foundFailedLogins, foundSuspiciousAccess bool - for _, entry := range cfg.Runbooks { - if entry.Title == "Failed Login Investigation" { - foundFailedLogins = true - gt.V(t, strings.Contains(entry.Description, "failed login attempts")).Equal(true) - gt.V(t, strings.Contains(entry.SQLContent, "login_failed")).Equal(true) - } - if entry.Title == "Suspicious Access Pattern" { - foundSuspiciousAccess = true - gt.V(t, strings.Contains(entry.Description, "unusual access patterns")).Equal(true) - gt.V(t, strings.Contains(entry.SQLContent, "access_logs")).Equal(true) - } - } - - gt.V(t, foundFailedLogins).Equal(true) - gt.V(t, foundSuspiciousAccess).Equal(true) -} - -func TestLoadConfigWithAdditionalRunbookPaths(t *testing.T) { - ctx := context.Background() - configPath := filepath.Join("testdata", "config_with_runbooks.yaml") - - // Create a temp runbook file - tmpDir := t.TempDir() - tmpRunbook := filepath.Join(tmpDir, "temp_runbook.sql") - content := `-- Title: Temp Runbook --- Description: Temporary runbook for testing -SELECT * FROM test -` - err := os.WriteFile(tmpRunbook, []byte(content), 0644) - gt.NoError(t, err) - - // Load config with additional paths - cfg, err := bigquery.LoadConfigWithRunbooks(ctx, configPath, []string{tmpDir}) - gt.NoError(t, err) - - // Verify additional runbook loaded - var foundTempRunbook bool - for _, entry := range cfg.Runbooks { - if entry.Title == "Temp Runbook" { - foundTempRunbook = true - break - } - } - gt.V(t, foundTempRunbook).Equal(true) -} - -func TestLoadConfigWithoutRunbooks(t *testing.T) { - ctx := context.Background() - - // Create a temp config without runbooks - tmpDir := t.TempDir() - configPath := filepath.Join(tmpDir, "config_no_runbooks.yaml") - configContent := `tables: - - project_id: test-project - dataset_id: test_dataset - table_id: test_table - description: Test table - -scan_size_limit: "1GB" -` - err := os.WriteFile(configPath, []byte(configContent), 0644) - gt.NoError(t, err) - - // Load config without runbooks - cfg, err := bigquery.LoadConfigWithRunbooks(ctx, configPath, nil) - gt.NoError(t, err) - gt.V(t, cfg).NotNil() - gt.V(t, len(cfg.Runbooks)).Equal(0) -} - -func TestRunbookIDCollisionHandling(t *testing.T) { - ctx := context.Background() - configPath := filepath.Join("testdata", "config_with_runbooks.yaml") - runbookDir := filepath.Join("testdata", "runbooks") - - // Load config with runbooks (including subdirectory) - cfg, err := bigquery.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - - // Verify we have runbooks from both root and subdirectory - // IDs are UUIDs, so we search by title - var hasRootFailedLogins, hasSubdirFailedLogins bool - for _, entry := range cfg.Runbooks { - t.Logf("Runbook ID: %s, Title: %s", entry.ID, entry.Title) - if entry.Title == "Failed Login Investigation" { - hasRootFailedLogins = true - } - if entry.Title == "Failed Login in Subdirectory" { - hasSubdirFailedLogins = true - } - } - - // Both should exist with different UUIDs - gt.V(t, hasRootFailedLogins).Equal(true) - gt.V(t, hasSubdirFailedLogins).Equal(true) -} diff --git a/pkg/agents/bigquery/export_test.go b/pkg/agents/bigquery/export_test.go deleted file mode 100644 index 464e3b67..00000000 --- a/pkg/agents/bigquery/export_test.go +++ /dev/null @@ -1,39 +0,0 @@ -package bigquery - -import ( - "github.com/gollem-dev/gollem" -) - -// NewToolSetForTest creates a toolSet instance for testing with direct configuration -func NewToolSetForTest(config *Config, projectID string, impersonateServiceAccount string) *toolSet { - return &toolSet{ - config: config, - tool: &internalTool{ - config: config, - projectID: projectID, - impersonateServiceAccount: impersonateServiceAccount, - }, - } -} - -// ExportNewInternalTool creates a new internalTool instance for testing -func ExportNewInternalTool(config *Config, projectID string, impersonateServiceAccount string) gollem.ToolSet { - return &internalTool{ - config: config, - projectID: projectID, - impersonateServiceAccount: impersonateServiceAccount, - } -} - -// ToolSpec is exported for testing -type ToolSpec = gollem.ToolSpec - -// ExportedBuildSystemPrompt is exported for testing -func ExportedBuildSystemPrompt(config *Config) (string, error) { - return buildSystemPrompt(config) -} - -// ExportedBuildPromptHint is exported for testing -func ExportedBuildPromptHint(config *Config) (string, error) { - return buildPromptHint(config) -} diff --git a/pkg/agents/bigquery/factory.go b/pkg/agents/bigquery/factory.go deleted file mode 100644 index 04b00f5f..00000000 --- a/pkg/agents/bigquery/factory.go +++ /dev/null @@ -1,109 +0,0 @@ -package bigquery - -import ( - "context" - "strings" - - "github.com/dustin/go-humanize" - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/domain/interfaces" - "github.com/secmon-lab/warren/pkg/utils/logging" - "github.com/urfave/cli/v3" -) - -// Factory implements agents.ToolSetFactory interface. -type Factory struct { - configPath string - projectID string - scanSizeLimitStr string - runbookPaths []string - impersonateServiceAccount string -} - -// Flags implements agents.ToolSetFactory -func (f *Factory) Flags() []cli.Flag { - return []cli.Flag{ - &cli.StringFlag{ - Name: "agent-bigquery-config", - Usage: "Path to BigQuery Agent configuration file (YAML)", - Destination: &f.configPath, - Category: "Agent:BigQuery", - Sources: cli.EnvVars("WARREN_AGENT_BIGQUERY_CONFIG"), - }, - &cli.StringFlag{ - Name: "agent-bigquery-project-id", - Usage: "Google Cloud Project ID for BigQuery operations", - Destination: &f.projectID, - Category: "Agent:BigQuery", - Sources: cli.EnvVars("WARREN_AGENT_BIGQUERY_PROJECT_ID"), - }, - &cli.StringFlag{ - Name: "agent-bigquery-scan-size-limit", - Usage: "Maximum scan size limit for BigQuery queries (e.g., '10GB', '1TB')", - Destination: &f.scanSizeLimitStr, - Category: "Agent:BigQuery", - Sources: cli.EnvVars("WARREN_AGENT_BIGQUERY_SCAN_SIZE_LIMIT"), - }, - &cli.StringSliceFlag{ - Name: "agent-bigquery-runbook-dir", - Usage: "Path to SQL runbook files or directories", - Destination: &f.runbookPaths, - Category: "Agent:BigQuery", - Sources: cli.EnvVars("WARREN_AGENT_BIGQUERY_RUNBOOK_DIR"), - }, - &cli.StringFlag{ - Name: "agent-bigquery-impersonate-service-account", - Usage: "Service account email to impersonate for BigQuery operations", - Destination: &f.impersonateServiceAccount, - Category: "Agent:BigQuery", - Sources: cli.EnvVars("WARREN_AGENT_BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT"), - }, - } -} - -// Configure implements agents.ToolSetFactory -func (f *Factory) Configure(ctx context.Context) (interfaces.ToolSet, error) { - if f.configPath == "" { - return nil, nil - } - - // Load config and runbooks - cfg, err := LoadConfigWithRunbooks(ctx, f.configPath, f.runbookPaths) - if err != nil { - return nil, goerr.Wrap(err, "failed to load BigQuery Agent config") - } - - // Override scan size limit from CLI flag if provided - if f.scanSizeLimitStr != "" { - limit, err := ParseScanSizeLimit(f.scanSizeLimitStr) - if err != nil { - return nil, goerr.Wrap(err, "failed to parse scan size limit") - } - cfg.ScanSizeLimit = limit - } - - // Log configuration - scanLimit := humanize.Bytes(cfg.ScanSizeLimit) - var tables []string - for _, t := range cfg.Tables { - tables = append(tables, strings.Join([]string{ - t.ProjectID, t.DatasetID, t.TableID, - }, ".")) - } - logging.From(ctx).Info("BigQuery Agent configured", - "tables", tables, - "scan_limit", scanLimit, - "runbooks", len(cfg.Runbooks)) - - // Create and return toolSet - ts := &toolSet{ - config: cfg, - tool: &internalTool{ - config: cfg, - projectID: f.projectID, - impersonateServiceAccount: f.impersonateServiceAccount, - }, - } - - return ts, nil -} diff --git a/pkg/agents/bigquery/prompt.go b/pkg/agents/bigquery/prompt.go deleted file mode 100644 index 872e6ab0..00000000 --- a/pkg/agents/bigquery/prompt.go +++ /dev/null @@ -1,103 +0,0 @@ -package bigquery - -import ( - "bytes" - _ "embed" - "text/template" - - "github.com/dustin/go-humanize" - "github.com/m-mizutani/goerr/v2" -) - -//go:embed prompt/base.md -var basePrompt string - -//go:embed prompt/system_with_memories.md -var systemWithMemoriesTemplate string - -//go:embed prompt/runbooks.md -var runbooksTemplate string - -//go:embed prompt/tool_description.md -var toolDescriptionTemplate string - -var systemWithMemoriesTmpl *template.Template -var runbooksTmpl *template.Template -var toolDescriptionTmpl *template.Template - -func init() { - systemWithMemoriesTmpl = template.Must(template.New("system_with_memories").Parse(systemWithMemoriesTemplate)) - runbooksTmpl = template.Must(template.New("runbooks").Parse(runbooksTemplate)) - toolDescriptionTmpl = template.Must(template.New("tool_description").Parse(toolDescriptionTemplate)) -} - -// promptData represents the data for system prompt template -type promptData struct { - Tables []TableConfig - Runbooks map[string]interface{} -} - -// promptHintData represents the data for tool_description.md template -type promptHintData struct { - HasTables bool - Tables []TableConfig - ScanSizeLimit string - QueryTimeout string -} - -// buildPromptHint renders the tool_description.md template with config data. -// The result is intended to be included in the parent agent's system prompt. -func buildPromptHint(config *Config) (string, error) { - data := promptHintData{ - HasTables: len(config.Tables) > 0, - Tables: config.Tables, - } - - if config.ScanSizeLimit > 0 { - data.ScanSizeLimit = humanize.IBytes(config.ScanSizeLimit) - } - if config.QueryTimeout > 0 { - data.QueryTimeout = config.QueryTimeout.String() - } - - var buf bytes.Buffer - if err := toolDescriptionTmpl.Execute(&buf, data); err != nil { - return "", goerr.Wrap(err, "failed to execute tool description template") - } - - return buf.String(), nil -} - -// buildSystemPrompt builds system prompt without memories (for factory) -func buildSystemPrompt(config *Config) (string, error) { - // Build base prompt - var buf bytes.Buffer - buf.WriteString(basePrompt) - buf.WriteString("\n\n") - - // Prepare template data - runbooksData := make(map[string]interface{}) - for id, entry := range config.Runbooks { - runbooksData[id.String()] = entry - } - - data := promptData{ - Tables: config.Tables, - Runbooks: runbooksData, - } - - // Execute main template - if err := systemWithMemoriesTmpl.Execute(&buf, data); err != nil { - return "", goerr.Wrap(err, "failed to execute system prompt template") - } - - // Append runbooks section - if len(config.Runbooks) > 0 { - buf.WriteString("\n\n") - if err := runbooksTmpl.Execute(&buf, data); err != nil { - return "", goerr.Wrap(err, "failed to execute runbooks template") - } - } - - return buf.String(), nil -} diff --git a/pkg/agents/bigquery/prompt/base.md b/pkg/agents/bigquery/prompt/base.md deleted file mode 100644 index 2e2a750b..00000000 --- a/pkg/agents/bigquery/prompt/base.md +++ /dev/null @@ -1,145 +0,0 @@ -# BigQuery Data Analysis Agent - -You are a BigQuery data analysis agent. Your role is to understand natural language queries, construct appropriate SQL queries, and retrieve data from BigQuery tables efficiently. - -## Core Principles - -### Efficiency First -- Construct the minimum necessary queries to achieve the objective -- Avoid exploratory queries unless data structure is unknown -- Use schema information to build queries correctly on the first attempt -- Results are limited to 1000 rows max. If your query returns more than 1000 rows, the results will be truncated - -### Data Fidelity -- Return query results in their raw form without interpretation -- Do not add opinions, insights, or observations -- Only add factual execution metadata (e.g., "No data found for the specified condition", "Results limited to 100 rows") - -### Available Tools -- `bigquery_schema`: Inspect table structure and field definitions -- `bigquery_query`: Execute SQL queries with automatic validation -- Scan size limits and timeouts are automatically enforced - -## Standard Workflow - -### 1. Understand the Request -Parse the natural language query to identify: -- Required data and fields -- Time range and temporal conditions -- Filter criteria and conditions -- Aggregation or grouping needs - -### 2. Schema Verification -Before writing SQL, use `bigquery_schema` to: -- Verify field names and their exact spelling -- Check data types and nested structures -- Understand field semantics and relationships -- Confirm the table contains expected data - -### 3. Query Construction -Build SQL following these requirements: -- Use fully qualified table names: `project.dataset.table` -- Select only necessary fields -- Apply appropriate LIMIT clauses -- Use proper date/time functions for temporal filtering -- Include GROUP BY when using aggregation functions - -### 4. Execution and Result Handling -- Execute the query once if schema verification was done properly -- If zero results are returned, systematically verify the query: - 1. **Check data existence**: Run `SELECT COUNT(*) FROM table` to confirm the table has data - 2. **Verify time range**: Remove or widen temporal filters to check if data exists in other periods - 3. **Inspect actual values**: Use `SELECT DISTINCT field FROM table LIMIT 20` on filter fields to see actual values - 4. **Sample raw data**: Run `SELECT * FROM table LIMIT 10` to understand actual data structure - 5. **Validate assumptions**: Compare expected vs actual field values (case, format, semantics) -- Adjust query based on verification results and re-execute - -## Common Data Field Categories - -Use these field patterns to **identify relevant fields from the schema** when constructing queries. After retrieving the schema, match the user's request to appropriate field categories to determine which fields to use. - -**Example usage**: -- Temporal fields → WHERE clause for time ranges, GROUP BY for time-based aggregation -- Classification fields → WHERE clause for filtering by type/status, GROUP BY for categorization -- Metric fields → Aggregation functions (SUM, AVG, COUNT), numeric filtering - -### Temporal Information -- Timestamps, event times, duration fields -- Creation, modification, expiration dates - -### Classification -- Event types, categories, severity levels -- Status, outcome, result codes - -### Identity and Principal -- User IDs, emails, usernames, account identifiers -- Service accounts, API keys, authentication tokens - -### Network and Location -- IP addresses, ports, hostnames, URLs -- Geographic data: country, region, coordinates -- Network protocols, connection states - -### Operations and Actions -- API calls, methods, commands, operations -- CRUD operations, administrative actions -- Target resources, affected objects - -### Metrics and Measurements -- Counts, sums, averages, percentiles -- Byte counts, request counts, error rates -- Performance metrics, latency, throughput - -### Contextual Attributes -- Request/response data, headers, parameters -- Error messages, stack traces, logs -- Custom metadata, tags, labels - -## Query Optimization Guidelines - -### Minimize Data Scanned -- Filter on partitioned columns (typically timestamp fields) -- Use WHERE clauses to reduce scanned data -- NEVER use SELECT * — always specify only the columns you need -- Select fewer columns to reduce scan cost (BigQuery is columnar storage — each column costs) - -### Leverage Schema Knowledge -- Reference nested fields correctly: `field.subfield` -- Use UNNEST for repeated fields when needed -- Apply appropriate type casting for comparisons - -### Handle Edge Cases -- Use COALESCE for nullable fields -- Apply SAFE_CAST to prevent type errors -- Consider time zone implications for timestamp comparisons - -## Anti-Patterns (MUST AVOID) - -### DO NOT paginate with OFFSET -- NEVER use OFFSET to retrieve additional pages of results -- If results are truncated (1000 rows), switch to COUNT/GROUP BY aggregation or add stricter WHERE filters - -### DO NOT use SELECT * -- Always specify only the columns needed for analysis -- Each additional column increases BigQuery scan cost - -### DO NOT retrieve raw data when aggregation suffices -- If you only need counts, distributions, or trends, use COUNT/GROUP BY/DISTINCT directly -- Only retrieve raw rows when individual record details are specifically required - -## Response Format - -Return results containing: -- The actual data records from the query -- All fields included in the SELECT clause -- Original data types and structures (nested objects, arrays) -- Row counts and query metadata (bytes scanned, execution time) -- Any result limitations (e.g., "Limited to 1000 rows", "No records matched the filter") - -Do not include: -- Interpretations or analysis of the data -- Recommendations or suggestions -- Observations about data patterns -- Security insights or threat assessments - -Present only factual query results and execution information. diff --git a/pkg/agents/bigquery/prompt/runbooks.md b/pkg/agents/bigquery/prompt/runbooks.md deleted file mode 100644 index 77446381..00000000 --- a/pkg/agents/bigquery/prompt/runbooks.md +++ /dev/null @@ -1,13 +0,0 @@ -{{if .Runbooks}} -## SQL Runbooks - -Available SQL runbook templates (use `get_runbook` tool with ID to get full SQL content): - -{{range $id, $entry := .Runbooks}}- **ID**: `{{$id}}` - **Title**: {{$entry.Title}}{{if $entry.Description}} - **Description**: {{$entry.Description}}{{end}} -{{end}} - -These runbooks contain pre-written SQL queries for common investigation patterns. Use the `get_runbook` tool to retrieve the full SQL content, which you can then adapt for your specific investigation needs. - -{{end}} diff --git a/pkg/agents/bigquery/prompt/system_with_memories.md b/pkg/agents/bigquery/prompt/system_with_memories.md deleted file mode 100644 index 97b48477..00000000 --- a/pkg/agents/bigquery/prompt/system_with_memories.md +++ /dev/null @@ -1,7 +0,0 @@ -## Available BigQuery Tables - -You have access to the following BigQuery tables: - -{{range .Tables -}} -- `{{.ProjectID}}.{{.DatasetID}}.{{.TableID}}`{{if .Description}}: {{.Description}}{{end}} -{{end}} diff --git a/pkg/agents/bigquery/prompt/tool_description.md b/pkg/agents/bigquery/prompt/tool_description.md deleted file mode 100644 index 6f6fd10c..00000000 --- a/pkg/agents/bigquery/prompt/tool_description.md +++ /dev/null @@ -1,28 +0,0 @@ -## BigQuery Agent - -You can query BigQuery tables using the `query_bigquery` tool. The agent will automatically check table schemas and construct appropriate queries. - -**Important Guidelines**: -- The agent MUST check table schemas before constructing queries -- Results will be returned as raw data records without summarization -- All query fields will be preserved in the response - -**How to Use**: -- Do NOT specify table names or SQL details in your query -- Focus on describing WHAT information you need, not HOW to get it -- Be clear about the data you want to retrieve (e.g., "login failures in the last 24 hours") -- The agent will automatically select appropriate tables and construct queries - -{{if .HasTables -}} -### Available Tables - -{{range .Tables -}} -- **`{{.ProjectID}}.{{.DatasetID}}.{{.TableID}}`**{{if .Description}}: {{.Description}}{{end}} -{{end}} -{{end -}} -{{if .ScanSizeLimit -}} -**Scan Size Limit**: {{.ScanSizeLimit}} -{{end -}} -{{if .QueryTimeout -}} -**Query Timeout**: {{.QueryTimeout}} -{{end -}} diff --git a/pkg/agents/bigquery/prompt_test.go b/pkg/agents/bigquery/prompt_test.go deleted file mode 100644 index 708006ae..00000000 --- a/pkg/agents/bigquery/prompt_test.go +++ /dev/null @@ -1,117 +0,0 @@ -package bigquery_test - -import ( - "testing" - "time" - - "github.com/m-mizutani/gt" - bqagent "github.com/secmon-lab/warren/pkg/agents/bigquery" -) - -func TestBuildSystemPrompt(t *testing.T) { - config := &bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "test-project", - DatasetID: "test-dataset", - TableID: "test-table", - Description: "Test table", - }, - }, - ScanSizeLimit: 1000000, - } - - prompt, err := bqagent.ExportedBuildSystemPrompt(config) - gt.NoError(t, err) - gt.V(t, prompt).NotEqual("") - gt.True(t, len(prompt) > 0) - gt.S(t, prompt).Contains("test-project") - gt.S(t, prompt).Contains("test-dataset") - gt.S(t, prompt).Contains("test-table") -} - -func TestBuildPromptHint(t *testing.T) { - t.Run("with tables and limits", func(t *testing.T) { - config := &bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "my-project", - DatasetID: "my-dataset", - TableID: "my-table", - Description: "A test table", - }, - { - ProjectID: "my-project", - DatasetID: "my-dataset", - TableID: "other-table", - Description: "Another table", - }, - }, - ScanSizeLimit: 10 * 1024 * 1024 * 1024, // 10 GB - QueryTimeout: 5 * time.Minute, - } - - hint, err := bqagent.ExportedBuildPromptHint(config) - gt.NoError(t, err) - gt.V(t, hint).NotEqual("") - gt.S(t, hint).Contains("my-project.my-dataset.my-table") - gt.S(t, hint).Contains("my-project.my-dataset.other-table") - gt.S(t, hint).Contains("A test table") - gt.S(t, hint).Contains("Another table") - gt.S(t, hint).Contains("10 GiB") - gt.S(t, hint).Contains("5m0s") - }) - - t.Run("without tables", func(t *testing.T) { - config := &bqagent.Config{ - Tables: nil, - } - - hint, err := bqagent.ExportedBuildPromptHint(config) - gt.NoError(t, err) - gt.V(t, hint).NotEqual("") - // Should still contain the base description - gt.S(t, hint).Contains("BigQuery Agent") - }) - - t.Run("without limits", func(t *testing.T) { - config := &bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "proj", - DatasetID: "ds", - TableID: "tbl", - }, - }, - } - - hint, err := bqagent.ExportedBuildPromptHint(config) - gt.NoError(t, err) - gt.S(t, hint).Contains("proj.ds.tbl") - gt.S(t, hint).NotContains("Scan Size Limit") - gt.S(t, hint).NotContains("Query Timeout") - }) -} - -func TestDynamicDescription(t *testing.T) { - config := &bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "proj-a", - DatasetID: "ds-a", - TableID: "tbl-a", - }, - { - ProjectID: "proj-b", - DatasetID: "ds-b", - TableID: "tbl-b", - }, - }, - } - - ts := bqagent.NewToolSetForTest(config, "test-project", "") - desc := ts.Description() - gt.S(t, desc).Contains("proj-a.ds-a.tbl-a") - gt.S(t, desc).Contains("proj-b.ds-b.tbl-b") - gt.S(t, desc).Contains("Available tables:") -} diff --git a/pkg/agents/bigquery/runbook.go b/pkg/agents/bigquery/runbook.go deleted file mode 100644 index 89487031..00000000 --- a/pkg/agents/bigquery/runbook.go +++ /dev/null @@ -1,208 +0,0 @@ -package bigquery - -import ( - "bufio" - "context" - "io/fs" - "os" - "path/filepath" - "regexp" - "strings" - - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/domain/model/bigquery" - "github.com/secmon-lab/warren/pkg/domain/types" -) - -// runbookLoader loads SQL runbook files and processes them -type runbookLoader struct { - paths []string -} - -// newRunbookLoader creates a new runbookLoader -func newRunbookLoader(paths []string) *runbookLoader { - return &runbookLoader{ - paths: paths, - } -} - -// loadRunbooks loads all SQL files from configured paths and returns RunbookEntries -func (r *runbookLoader) loadRunbooks(ctx context.Context) (bigquery.RunbookEntries, error) { - var entries bigquery.RunbookEntries - - for _, path := range r.paths { - pathEntries, err := r.loadFromPath(ctx, path) - if err != nil { - return nil, goerr.Wrap(err, "failed to load runbooks from path", goerr.V("path", path)) - } - entries = append(entries, pathEntries...) - } - - return entries, nil -} - -// loadFromPath loads runbooks from a single path (file or directory) -func (r *runbookLoader) loadFromPath(ctx context.Context, path string) (bigquery.RunbookEntries, error) { - info, err := os.Stat(path) - if err != nil { - return nil, goerr.Wrap(err, "failed to stat path", goerr.V("path", path)) - } - - if info.IsDir() { - return r.loadFromDirectory(ctx, path) - } - - // Single file - entry, err := r.loadFromFile(ctx, path) - if err != nil { - return nil, err - } - - if entry == nil { - return nil, nil // Not a SQL file - } - - return bigquery.RunbookEntries{entry}, nil -} - -// loadFromDirectory loads all SQL files from a directory recursively -func (r *runbookLoader) loadFromDirectory(ctx context.Context, dirPath string) (bigquery.RunbookEntries, error) { - var entries bigquery.RunbookEntries - - // Clean the base directory path - cleanDirPath := filepath.Clean(dirPath) - - err := filepath.WalkDir(cleanDirPath, func(path string, d fs.DirEntry, err error) error { - if err != nil { - return goerr.Wrap(err, "failed to walk directory", goerr.V("path", path)) - } - - if d.IsDir() { - return nil - } - - // Only process .sql files - if !strings.HasSuffix(strings.ToLower(path), ".sql") { - return nil - } - - entry, err := r.loadFromFile(ctx, path) - if err != nil { - return err - } - - if entry != nil { - entries = append(entries, entry) - } - - return nil - }) - - if err != nil { - return nil, err - } - - return entries, nil -} - -// loadFromFile loads a single SQL file and creates a RunbookEntry -func (r *runbookLoader) loadFromFile(_ context.Context, filePath string) (*bigquery.RunbookEntry, error) { - // Only process .sql files - if !strings.HasSuffix(strings.ToLower(filePath), ".sql") { - return nil, nil - } - - // Clean the file path to prevent path traversal - cleanPath := filepath.Clean(filePath) - - // Read file content - content, err := os.ReadFile(cleanPath) - if err != nil { - return nil, goerr.Wrap(err, "failed to read SQL file", goerr.V("path", cleanPath)) - } - - // Use UUID for ID to prevent collisions - id := types.NewRunbookID() - - // Extract filename (base name) for default title - filename := filepath.Base(cleanPath) - - // Extract title and description from comments - title, description := r.extractTitleAndDescription(string(content)) - if title == "" { - // Use filename as title if no title found in comments - title = strings.TrimSuffix(filename, ".sql") - } - - entry := &bigquery.RunbookEntry{ - ID: id, - Title: title, - Description: description, - SQLContent: string(content), - } - - return entry, nil -} - -// extractTitleAndDescription extracts title and description from SQL comments -// Format: -// -- Title: Title of the runbook -// -- Description: Description of the runbook -// -- This can span multiple lines -func (r *runbookLoader) extractTitleAndDescription(content string) (string, string) { - scanner := bufio.NewScanner(strings.NewReader(content)) - - var title string - var description strings.Builder - var inDescription bool - - titleRegex := regexp.MustCompile(`(?i)^--\s*title\s*:\s*(.+)$`) - descRegex := regexp.MustCompile(`(?i)^--\s*description\s*:\s*(.+)$`) - commentRegex := regexp.MustCompile(`^--\s*(.*)$`) - - for scanner.Scan() { - line := strings.TrimSpace(scanner.Text()) - - // Skip empty lines - if line == "" { - continue - } - - // If we hit a non-comment line, stop processing - if !strings.HasPrefix(line, "--") { - break - } - - // Check for title - if matches := titleRegex.FindStringSubmatch(line); matches != nil { - title = strings.TrimSpace(matches[1]) - inDescription = false - continue - } - - // Check for description start - if matches := descRegex.FindStringSubmatch(line); matches != nil { - inDescription = true - if description.Len() > 0 { - description.WriteString(" ") - } - description.WriteString(strings.TrimSpace(matches[1])) - continue - } - - // If we're in description mode, continue collecting comment lines - if inDescription { - if matches := commentRegex.FindStringSubmatch(line); matches != nil { - if description.Len() > 0 { - description.WriteString(" ") - } - description.WriteString(strings.TrimSpace(matches[1])) - } else { - // Non-comment line, stop collecting description - break - } - } - } - - return title, description.String() -} diff --git a/pkg/agents/bigquery/testdata/config_with_runbooks.yaml b/pkg/agents/bigquery/testdata/config_with_runbooks.yaml deleted file mode 100644 index c9a66963..00000000 --- a/pkg/agents/bigquery/testdata/config_with_runbooks.yaml +++ /dev/null @@ -1,8 +0,0 @@ -tables: - - project_id: test-project - dataset_id: test_dataset - table_id: test_table - description: Test table with sample data - -scan_size_limit: "1GB" -query_timeout: 5m diff --git a/pkg/agents/bigquery/testdata/example_config.yaml b/pkg/agents/bigquery/testdata/example_config.yaml deleted file mode 100644 index 35a4cf41..00000000 --- a/pkg/agents/bigquery/testdata/example_config.yaml +++ /dev/null @@ -1,39 +0,0 @@ -# BigQuery Agent Configuration Example -# This file demonstrates the flat table structure for configuring BigQuery tables - -# Minimal configuration -tables: - - project_id: my-project - dataset_id: security_logs - table_id: events - description: Security event logs - -scan_size_limit: "1GB" - -# --- -# Multi-table configuration example -# Uncomment and customize the section below for multiple tables - -# tables: -# - project_id: prod-logs -# dataset_id: aws_cloudtrail -# table_id: events_2024 -# description: | -# CloudTrail events from 2024: -# - Authentication (ConsoleLogin, AssumeRole) -# - S3 access (GetObject, PutObject) -# - IAM changes -# - Partitioned by event_date -# -# - project_id: prod-logs -# dataset_id: gcp_audit -# table_id: admin_activity -# description: GCP admin activity logs -# -# - project_id: staging-logs -# dataset_id: app_logs -# table_id: authentication -# description: App authentication events -# -# scan_size_limit: "5GB" -# query_timeout: 3m diff --git a/pkg/agents/bigquery/testdata/runbooks/failed_logins.sql b/pkg/agents/bigquery/testdata/runbooks/failed_logins.sql deleted file mode 100644 index 9d7e3aa3..00000000 --- a/pkg/agents/bigquery/testdata/runbooks/failed_logins.sql +++ /dev/null @@ -1,14 +0,0 @@ --- Title: Failed Login Investigation --- Description: Query to find failed login attempts by IP address and user -SELECT - timestamp, - user_email, - source_ip, - error_code, - user_agent -FROM `project.dataset.auth_logs` -WHERE - event_type = 'login_failed' - AND timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 24 HOUR) -ORDER BY timestamp DESC -LIMIT 100 diff --git a/pkg/agents/bigquery/testdata/runbooks/subdirectory/failed_logins.sql b/pkg/agents/bigquery/testdata/runbooks/subdirectory/failed_logins.sql deleted file mode 100644 index ed338dbb..00000000 --- a/pkg/agents/bigquery/testdata/runbooks/subdirectory/failed_logins.sql +++ /dev/null @@ -1,4 +0,0 @@ --- Title: Failed Login in Subdirectory --- Description: This is a test for collision handling -SELECT * FROM subdirectory_logs -WHERE failed = true diff --git a/pkg/agents/bigquery/testdata/runbooks/suspicious_access.sql b/pkg/agents/bigquery/testdata/runbooks/suspicious_access.sql deleted file mode 100644 index 76d8f38f..00000000 --- a/pkg/agents/bigquery/testdata/runbooks/suspicious_access.sql +++ /dev/null @@ -1,14 +0,0 @@ --- title: Suspicious Access Pattern --- description: Identify unusual access patterns from specific IP addresses -SELECT - timestamp, - source_ip, - COUNT(*) as access_count, - ARRAY_AGG(DISTINCT user_email) as users -FROM `project.dataset.access_logs` -WHERE - timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 HOUR) -GROUP BY timestamp, source_ip -HAVING access_count > 10 -ORDER BY access_count DESC -LIMIT 50 diff --git a/pkg/agents/bigquery/tool.go b/pkg/agents/bigquery/tool.go deleted file mode 100644 index 99a70459..00000000 --- a/pkg/agents/bigquery/tool.go +++ /dev/null @@ -1,500 +0,0 @@ -package bigquery - -import ( - "context" - "encoding/json" - "fmt" - - "cloud.google.com/go/bigquery" - "github.com/dustin/go-humanize" - "github.com/gollem-dev/gollem" - "github.com/google/uuid" - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/domain/types" - "github.com/secmon-lab/warren/pkg/utils/logging" - "github.com/secmon-lab/warren/pkg/utils/msg" - "google.golang.org/api/impersonate" - "google.golang.org/api/iterator" - "google.golang.org/api/option" -) - -// internalTool implements the actual BigQuery operations -type internalTool struct { - config *Config - projectID string - impersonateServiceAccount string -} - -// ID implements gollem.ToolSet -func (t *internalTool) ID() string { - return "bigquery_agent" -} - -// Specs implements gollem.ToolSet -func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - // Build table descriptions for the prompt - tableDescriptions := "" - for i, table := range t.config.Tables { - tableDescriptions += fmt.Sprintf("\n%d. %s.%s.%s: %s", - i+1, table.ProjectID, table.DatasetID, table.TableID, table.Description) - } - - specs := []gollem.ToolSpec{ - { - Name: "bigquery_query", - Description: fmt.Sprintf(`Execute a BigQuery SQL query to retrieve data. -Available tables:%s - -Important guidelines: -- Always specify the full table name as project.dataset.table -- Results are limited to 1000 rows max. Queries returning more than 1000 rows will be truncated -- Scan size limit: %s (queries exceeding this will fail) -- For time-based queries, use proper date/time functions and partitioning fields when available -- Select only the fields needed for analysis to minimize scan size (DO NOT use SELECT *) -- Use WHERE clauses to filter data efficiently -- DO NOT use OFFSET for pagination. If results are truncated, use COUNT/GROUP BY to aggregate or add stricter WHERE filters - -Best practices: -- Start with schema inspection if unfamiliar with table structure -- Apply time range filters to reduce scan size -- Use LIMIT to restrict rows appropriately`, - tableDescriptions, - humanize.Bytes(t.config.ScanSizeLimit), - ), - Parameters: map[string]*gollem.Parameter{ - "sql": { - Type: gollem.TypeString, - Description: "The SQL query to execute. Must be a valid BigQuery SQL query.", - Required: true, - }, - "description": { - Type: gollem.TypeString, - Description: "Brief description of what this query is trying to achieve (optional, for logging/tracking)", - }, - }, - }, - { - Name: "bigquery_schema", - Description: `Get detailed schema information for a specific BigQuery table. -Use this to understand available fields, data types, and nested structures before constructing queries. -Returns complete schema including nested RECORD fields.`, - Parameters: map[string]*gollem.Parameter{ - "project_id": { - Type: gollem.TypeString, - Description: "The project ID of the table", - Required: true, - }, - "dataset_id": { - Type: gollem.TypeString, - Description: "The dataset ID of the table", - Required: true, - }, - "table_id": { - Type: gollem.TypeString, - Description: "The table ID to get schema for", - Required: true, - }, - }, - }, - } - - // Add get_runbook tool if runbooks are configured - if len(t.config.Runbooks) > 0 { - specs = append(specs, gollem.ToolSpec{ - Name: "get_runbook", - Description: "Get the full SQL content of a runbook by its ID. Use this when you want to see or adapt a pre-written query template.", - Parameters: map[string]*gollem.Parameter{ - "runbook_id": { - Type: gollem.TypeString, - Description: "The ID of the runbook to retrieve", - Required: true, - }, - }, - }) - } - - return specs, nil -} - -// Run implements gollem.ToolSet -func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Internal tool run started", "function", name, "args_keys", getMapKeys(args)) - - switch name { - case "bigquery_query": - return t.executeQuery(ctx, args) - case "bigquery_schema": - return t.getTableSchema(ctx, args) - case "get_runbook": - return t.getRunbook(ctx, args) - default: - log.Debug("Unknown internal tool function", "name", name) - return nil, goerr.New("unknown function", goerr.V("name", name)) - } -} - -func (t *internalTool) executeQuery(ctx context.Context, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Executing BigQuery query") - - sql, ok := args["sql"].(string) - if !ok { - log.Debug("SQL parameter is missing or invalid") - return nil, goerr.New("sql parameter is required") - } - - log.Debug("Query SQL prepared", "sql_length", len(sql)) - msg.Trace(ctx, "📝 Executing query: ```\n%s\n````", sql) - - // Determine project ID: CLI flag takes precedence over config - projectID := t.projectID - if projectID == "" { - // Fallback to first table's project ID from config - if len(t.config.Tables) == 0 { - log.Debug("No project ID specified and no tables configured") - return nil, goerr.New("no project ID specified and no tables configured") - } - projectID = t.config.Tables[0].ProjectID - log.Debug("Using project ID from first table", "project_id", projectID) - } else { - log.Debug("Using project ID from configuration", "project_id", projectID) - } - - // Create BigQuery client - log.Debug("Creating BigQuery client", "project_id", projectID) - client, err := t.newBigQueryClient(ctx, projectID) - if err != nil { - log.Debug("Failed to create BigQuery client", "error", err) - msg.Trace(ctx, "❌ Failed to create BigQuery client: %v", err) - return nil, goerr.Wrap(err, "failed to create BigQuery client") - } - defer func() { - if err := client.Close(); err != nil { - log.Debug("Failed to close BigQuery client", "error", err) - msg.Trace(ctx, "⚠️ Failed to close BigQuery client: %v", err) - } - }() - log.Debug("BigQuery client created successfully") - - // Create query - q := client.Query(sql) - - // Perform dry run to check scan size - log.Debug("Performing dry run to check scan size") - q.DryRun = true - job, err := q.Run(ctx) - if err != nil { - log.Debug("Dry run failed", "error", err) - msg.Trace(ctx, "❌ Dry run failed: %v", err) - return nil, goerr.Wrap(err, "failed to dry run query") - } - - totalBytes := job.LastStatus().Statistics.TotalBytesProcessed - log.Debug("Dry run completed", "total_bytes", totalBytes, "job_id", job.ID()) - - if totalBytes < 0 { - log.Debug("Invalid negative bytes processed", "bytes", totalBytes) - msg.Trace(ctx, "❌ Invalid negative bytes processed: %d", totalBytes) - return nil, goerr.New("invalid negative bytes processed", - goerr.V("bytes_processed", totalBytes)) - } - - // Check scan size - log.Debug("Checking scan size", - "total_bytes", totalBytes, - "scan_limit", t.config.ScanSizeLimit, - "total_bytes_human", humanize.Bytes(uint64(totalBytes)), - "scan_limit_human", humanize.Bytes(t.config.ScanSizeLimit)) - - // Check scan size limit - if totalBytes > 0 && uint64(totalBytes) > t.config.ScanSizeLimit { - log.Debug("Query scan size exceeds limit", - "scan_size", totalBytes, - "scan_limit", t.config.ScanSizeLimit) - msg.Trace(ctx, "❌ Query scan size exceeds limit") - - return nil, goerr.New("query scan size exceeds limit", - goerr.V("scan_size", totalBytes), - goerr.V("scan_limit", t.config.ScanSizeLimit)) - } - - // Execute the actual query - log.Debug("Executing actual query") - q.DryRun = false - job, err = q.Run(ctx) - if err != nil { - log.Debug("Query execution failed", "error", err) - msg.Trace(ctx, "❌ Query execution failed: %v", err) - return nil, goerr.Wrap(err, "failed to run query") - } - log.Debug("Query job submitted", "job_id", job.ID()) - - // Wait for job to complete (with timeout) - timeout := t.config.GetQueryTimeout() - log.Debug("Waiting for job completion", "job_id", job.ID(), "timeout", timeout) - waitCtx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - status, err := job.Wait(waitCtx) - if err != nil { - log.Debug("Job wait failed", "job_id", job.ID(), "error", err) - msg.Trace(ctx, "❌ Job wait failed: %v", err) - return nil, goerr.Wrap(err, "failed to wait for job completion") - } - log.Debug("Job wait completed", "job_id", job.ID()) - - if err := status.Err(); err != nil { - log.Debug("Job execution failed", "job_id", job.ID(), "error", err) - msg.Trace(ctx, "❌ Job execution failed: %v", err) - return nil, goerr.Wrap(err, "job failed") - } - log.Debug("Job completed successfully", "job_id", job.ID()) - - // Read results - log.Debug("Reading query results", "job_id", job.ID()) - it, err := job.Read(ctx) - if err != nil { - log.Debug("Failed to read query results", "job_id", job.ID(), "error", err) - msg.Trace(ctx, "❌ Failed to read query results: %v", err) - return nil, goerr.Wrap(err, "failed to read query results") - } - log.Debug("Result iterator created", "job_id", job.ID(), "schema_fields", len(it.Schema)) - - // Collect rows (limit to 1000 rows as a safety cap) - var rows []map[string]any - limit := 1000 - log.Debug("Collecting rows", "limit", limit) - for len(rows) < limit { - var row []bigquery.Value - err := it.Next(&row) - if err == iterator.Done { - log.Debug("Reached end of results", "rows_collected", len(rows)) - break - } - if err != nil { - log.Debug("Failed to iterate results", "error", err, "rows_collected", len(rows)) - msg.Trace(ctx, "❌ Failed to iterate results: %v", err) - return nil, goerr.Wrap(err, "failed to iterate results") - } - - // Convert row to map - rowMap := make(map[string]any) - for i, field := range it.Schema { - rowMap[field.Name] = convertBigQueryValue(row[i]) - } - rows = append(rows, rowMap) - } - log.Debug("Rows collected", "count", len(rows), "has_more", len(rows) >= limit) - - queryID := uuid.New().String() - log.Debug("Query execution complete", "query_id", queryID, "total_rows", len(rows)) - - // Trace: Query completed successfully - msg.Trace(ctx, "✅ Query completed: %d rows retrieved, scan size: %s (query_id: %s)", - len(rows), humanize.Bytes(uint64(totalBytes)), queryID) - - hasMore := len(rows) >= limit - note := "Query completed successfully." - if hasMore { - note = fmt.Sprintf("Limited to %d rows. There are more results. Do NOT paginate with OFFSET. Instead, use COUNT/GROUP BY to aggregate the data, or add more specific WHERE filters to narrow the results.", limit) - } - - return map[string]any{ - "query_id": queryID, - "rows": rows, - "total_rows": len(rows), - "bytes_processed": totalBytes, - "execution_status": "completed", - "has_more": hasMore, - "note": note, - }, nil -} - -// getTableSchema retrieves detailed schema information for a BigQuery table -func (t *internalTool) getTableSchema(ctx context.Context, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Getting BigQuery table schema") - - projectID, ok := args["project_id"].(string) - if !ok { - log.Debug("project_id parameter is missing or invalid") - return nil, goerr.New("project_id parameter is required") - } - datasetID, ok := args["dataset_id"].(string) - if !ok { - log.Debug("dataset_id parameter is missing or invalid") - return nil, goerr.New("dataset_id parameter is required") - } - tableID, ok := args["table_id"].(string) - if !ok { - log.Debug("table_id parameter is missing or invalid") - return nil, goerr.New("table_id parameter is required") - } - - log.Debug("Schema request parameters", - "project_id", projectID, - "dataset_id", datasetID, - "table_id", tableID) - msg.Trace(ctx, "🧐 Retrieving schema... `%s.%s.%s`", projectID, datasetID, tableID) - - // Create BigQuery client - log.Debug("Creating BigQuery client for schema retrieval", "project_id", projectID) - client, err := t.newBigQueryClient(ctx, projectID) - if err != nil { - log.Debug("Failed to create BigQuery client", "error", err) - return nil, goerr.Wrap(err, "failed to create BigQuery client") - } - defer func() { - if err := client.Close(); err != nil { - log.Debug("Failed to close BigQuery client", "error", err) - msg.Trace(ctx, "⚠️ Failed to close BigQuery client: %v", err) - } - }() - log.Debug("BigQuery client created for schema retrieval") - - // Get table metadata - log.Debug("Fetching table metadata", "dataset", datasetID, "table", tableID) - table := client.Dataset(datasetID).Table(tableID) - metadata, err := table.Metadata(ctx) - if err != nil { - log.Debug("Failed to get table metadata", - "error", err, - "project_id", projectID, - "dataset_id", datasetID, - "table_id", tableID) - return nil, goerr.Wrap(err, "failed to get table metadata", - goerr.V("project_id", projectID), - goerr.V("dataset_id", datasetID), - goerr.V("table_id", tableID)) - } - log.Debug("Table metadata retrieved", "schema_fields", len(metadata.Schema)) - - // Convert schema to JSON - log.Debug("Converting schema to JSON", "field_count", len(metadata.Schema)) - schemaJSON, err := json.Marshal(metadata.Schema) - if err != nil { - log.Debug("Failed to marshal schema to JSON", "error", err) - return nil, goerr.Wrap(err, "failed to marshal schema to JSON") - } - log.Debug("Schema marshaled to JSON", "json_length", len(schemaJSON)) - - var result []map[string]any - if err := json.Unmarshal(schemaJSON, &result); err != nil { - log.Debug("Failed to unmarshal schema from JSON", "error", err) - return nil, goerr.Wrap(err, "failed to unmarshal schema from JSON") - } - log.Debug("Schema successfully retrieved", - "project_id", projectID, - "dataset_id", datasetID, - "table_id", tableID, - "field_count", len(result)) - - return map[string]any{ - "schema": result, - "project_id": projectID, - "dataset_id": datasetID, - "table_id": tableID, - }, nil -} - -// newBigQueryClient creates a new BigQuery client -func (t *internalTool) newBigQueryClient(ctx context.Context, projectID string) (*bigquery.Client, error) { - log := logging.From(ctx) - var opts []option.ClientOption - - // If impersonation is configured, use impersonated credentials - if t.impersonateServiceAccount != "" { - log.Debug("Using service account impersonation for BigQuery", - "service_account", t.impersonateServiceAccount) - - ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ - TargetPrincipal: t.impersonateServiceAccount, - Scopes: []string{ - "https://www.googleapis.com/auth/bigquery", - "https://www.googleapis.com/auth/cloud-platform", - }, - }) - if err != nil { - log.Debug("Failed to create impersonated credentials", - "error", err, - "service_account", t.impersonateServiceAccount) - return nil, goerr.Wrap(err, "failed to create impersonated credentials", - goerr.V("service_account", t.impersonateServiceAccount)) - } - opts = append(opts, option.WithTokenSource(ts)) - } else { - log.Debug("Using Application Default Credentials for BigQuery") - } - - return bigquery.NewClient(ctx, projectID, opts...) -} - -// convertBigQueryValue converts BigQuery values to JSON-safe types -func convertBigQueryValue(val bigquery.Value) any { - if val == nil { - return nil - } - - switch v := val.(type) { - case []bigquery.Value: - // Handle repeated fields (arrays) - result := make([]any, len(v)) - for i, item := range v { - result[i] = convertBigQueryValue(item) - } - return result - case map[string]bigquery.Value: - // Handle nested/struct fields - result := make(map[string]any) - for key, item := range v { - result[key] = convertBigQueryValue(item) - } - return result - default: - // For primitive types, return as-is - return val - } -} - -// getRunbook retrieves a runbook entry by ID -func (t *internalTool) getRunbook(ctx context.Context, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Getting runbook") - - runbookIDStr, ok := args["runbook_id"].(string) - if !ok { - log.Debug("runbook_id parameter is missing or invalid") - return nil, goerr.New("runbook_id parameter is required") - } - - // Parse runbook ID - runbookID := types.RunbookID(runbookIDStr) - log.Debug("Looking up runbook", "runbook_id", runbookID) - - // Find runbook in config - entry, exists := t.config.Runbooks[runbookID] - if !exists { - log.Debug("Runbook not found", "runbook_id", runbookID) - return nil, goerr.New("runbook not found", goerr.V("runbook_id", runbookID)) - } - - log.Debug("Runbook found", "runbook_id", runbookID, "title", entry.Title) - msg.Trace(ctx, "📖 Retrieved runbook: %s", entry.Title) - - return map[string]any{ - "runbook_id": runbookID.String(), - "title": entry.Title, - "description": entry.Description, - "sql": entry.SQLContent, - }, nil -} - -// getMapKeys returns the keys of a map as a slice -func getMapKeys(m map[string]any) []string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return keys -} diff --git a/pkg/agents/bigquery/tool_test.go b/pkg/agents/bigquery/tool_test.go deleted file mode 100644 index 397be1b7..00000000 --- a/pkg/agents/bigquery/tool_test.go +++ /dev/null @@ -1,163 +0,0 @@ -package bigquery_test - -import ( - "context" - "testing" - - "github.com/m-mizutani/gt" - bqagent "github.com/secmon-lab/warren/pkg/agents/bigquery" -) - -// Tests for internal tool implementation - -func TestInternalTool_Specs(t *testing.T) { - ctx := context.Background() - config := &bqagent.Config{ - Tables: []bqagent.TableConfig{ - { - ProjectID: "test-project", - DatasetID: "test-dataset", - TableID: "test-table", - Description: "Test table for unit testing", - }, - }, - ScanSizeLimit: 10 * 1024 * 1024 * 1024, // 10GB - } - - tool := bqagent.ExportNewInternalTool(config, "test-project", "") - - specs, err := tool.Specs(ctx) - gt.NoError(t, err) - gt.V(t, len(specs)).Equal(2) // bigquery_query and bigquery_schema - - // Find query spec - var querySpec *bqagent.ToolSpec - for i := range specs { - if specs[i].Name == "bigquery_query" { - querySpec = &specs[i] - break - } - } - - gt.V(t, querySpec).NotNil() - gt.V(t, querySpec.Name).Equal("bigquery_query") - gt.True(t, querySpec.Parameters["sql"].Required) -} - -func TestInternalTool_GetRunbook(t *testing.T) { - ctx := context.Background() - - t.Run("retrieve runbook by ID", func(t *testing.T) { - // Load config with runbooks - configPath := "testdata/config_with_runbooks.yaml" - runbookDir := "testdata/runbooks" - cfg, err := bqagent.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - - // Find runbook ID by title (IDs are UUIDs) - var targetID string - for id, entry := range cfg.Runbooks { - if entry.Title == "Failed Login Investigation" { - targetID = id.String() - break - } - } - gt.V(t, targetID != "").Equal(true) - - // Call get_runbook through internal tool - tool := bqagent.ExportNewInternalTool(cfg, "test-project", "") - result, err := tool.Run(ctx, "get_runbook", map[string]any{ - "runbook_id": targetID, - }) - gt.NoError(t, err) - gt.V(t, result).NotNil() - - // Verify result structure - gt.V(t, result["runbook_id"]).Equal(targetID) - gt.V(t, result["title"]).Equal("Failed Login Investigation") - gt.V(t, result["description"]).NotNil() - gt.V(t, result["sql"]).NotNil() - - // Verify SQL content matches expected content - sqlContent, ok := result["sql"].(string) - gt.V(t, ok).Equal(true) - gt.V(t, len(sqlContent) > 0).Equal(true) - gt.S(t, sqlContent).Contains("-- Title: Failed Login Investigation") - gt.S(t, sqlContent).Contains("SELECT") - gt.S(t, sqlContent).Contains("FROM `project.dataset.auth_logs`") - gt.S(t, sqlContent).Contains("event_type = 'login_failed'") - }) - - t.Run("runbook not found", func(t *testing.T) { - configPath := "testdata/config_with_runbooks.yaml" - runbookDir := "testdata/runbooks" - cfg, err := bqagent.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - - tool := bqagent.ExportNewInternalTool(cfg, "test-project", "") - _, err = tool.Run(ctx, "get_runbook", map[string]any{ - "runbook_id": "nonexistent", - }) - gt.Error(t, err) - }) - - t.Run("missing runbook_id parameter", func(t *testing.T) { - configPath := "testdata/config_with_runbooks.yaml" - runbookDir := "testdata/runbooks" - cfg, err := bqagent.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - - tool := bqagent.ExportNewInternalTool(cfg, "test-project", "") - _, err = tool.Run(ctx, "get_runbook", map[string]any{}) - gt.Error(t, err) - }) -} - -func TestInternalTool_Specs_WithRunbooks(t *testing.T) { - ctx := context.Background() - - t.Run("specs include get_runbook when runbooks configured", func(t *testing.T) { - configPath := "testdata/config_with_runbooks.yaml" - runbookDir := "testdata/runbooks" - cfg, err := bqagent.LoadConfigWithRunbooks(ctx, configPath, []string{runbookDir}) - gt.NoError(t, err) - - tool := bqagent.ExportNewInternalTool(cfg, "test-project", "") - specs, err := tool.Specs(ctx) - gt.NoError(t, err) - - // Should have bigquery_query, bigquery_schema, and get_runbook - gt.V(t, len(specs)).Equal(3) - - // Find get_runbook spec - var getRunbookSpec *bqagent.ToolSpec - for i := range specs { - if specs[i].Name == "get_runbook" { - getRunbookSpec = &specs[i] - break - } - } - - gt.V(t, getRunbookSpec).NotNil() - gt.V(t, getRunbookSpec.Name).Equal("get_runbook") - gt.True(t, getRunbookSpec.Parameters["runbook_id"].Required) - }) - - t.Run("specs exclude get_runbook when no runbooks", func(t *testing.T) { - configPath := "testdata/config_with_runbooks.yaml" - cfg, err := bqagent.LoadConfig(configPath) - gt.NoError(t, err) - - tool := bqagent.ExportNewInternalTool(cfg, "test-project", "") - specs, err := tool.Specs(ctx) - gt.NoError(t, err) - - // Should only have bigquery_query and bigquery_schema - gt.V(t, len(specs)).Equal(2) - - // Verify get_runbook is not present - for _, spec := range specs { - gt.V(t, spec.Name).NotEqual("get_runbook") - } - }) -} diff --git a/pkg/agents/factory.go b/pkg/agents/factory.go deleted file mode 100644 index 602de663..00000000 --- a/pkg/agents/factory.go +++ /dev/null @@ -1,19 +0,0 @@ -package agents - -import ( - "context" - - "github.com/secmon-lab/warren/pkg/domain/interfaces" - "github.com/urfave/cli/v3" -) - -// ToolSetFactory is the interface that all agent packages must implement. -// This interface provides a unified way for the CLI layer to interact with agents. -type ToolSetFactory interface { - // Flags returns CLI flags for this agent - Flags() []cli.Flag - - // Configure creates and initializes the agent, returning a ToolSet. - // Returns (nil, nil) if the agent is not configured. - Configure(ctx context.Context) (interfaces.ToolSet, error) -} diff --git a/pkg/agents/falcon/README.md b/pkg/agents/falcon/README.md deleted file mode 100644 index 1e395ac0..00000000 --- a/pkg/agents/falcon/README.md +++ /dev/null @@ -1,149 +0,0 @@ -# CrowdStrike Falcon Agent Configuration - -## Overview - -The Falcon agent enables Warren to query CrowdStrike Falcon for security incidents, alerts, behaviors, CrowdScores, and EDR telemetry events. This is a **read-only** agent — it does not modify any data in your Falcon environment. - -## Prerequisites - -### Create API Client Credentials - -1. Log in to the [CrowdStrike Falcon Console](https://falcon.crowdstrike.com/) -2. Navigate to **Support and resources > Resources and tools > API clients and keys** -3. Click **Create API client** -4. Configure the following: - - **Client Name**: A descriptive name (e.g., `warren-readonly`) - - **API Scopes**: Select the scopes listed below -5. Click **Create** -6. **Copy and securely store** the Client ID and Client Secret (the secret cannot be retrieved later) - -### Required API Scopes - -| Scope | Permission | Purpose | -|---|---|---| -| **Incidents** | Read | Search and retrieve incident details, behaviors | -| **Alerts** | Read | Search and retrieve alert details, aggregates | -| **NGSIEM** | Read + Write | Search EDR telemetry events via Next-Gen SIEM API | - -> **Note:** The NGSIEM scope requires both Read and Write permissions because the Search API uses Write to create query jobs and Read to retrieve results. Despite requiring Write permission, this agent only performs read operations (search queries) — it does not modify any data. - -## Cloud Regions - -CrowdStrike operates in multiple cloud regions. Set the base URL matching your Falcon tenant region: - -| Cloud Region | Base URL | -|---|---| -| US-1 (default) | `https://api.crowdstrike.com` | -| US-2 | `https://api.us-2.crowdstrike.com` | -| EU-1 | `https://api.eu-1.crowdstrike.com` | -| US-GOV-1 | `https://api.laggar.gcw.crowdstrike.com` | -| US-GOV-2 | `https://api.us-gov-2.crowdstrike.mil` | - -You can find your region by checking the domain in your Falcon console URL (e.g., `falcon.us-2.crowdstrike.com` means US-2). - -## Deployment - -### Environment Variables - -```bash -export WARREN_AGENT_FALCON_CLIENT_ID="your-client-id" -export WARREN_AGENT_FALCON_CLIENT_SECRET="your-client-secret" -export WARREN_AGENT_FALCON_BASE_URL="https://api.crowdstrike.com" # Optional, defaults to US-1 -``` - -### CLI Flags - -```bash -warren serve \ - --agent-falcon-client-id=your-client-id \ - --agent-falcon-client-secret=your-client-secret \ - --agent-falcon-base-url=https://api.crowdstrike.com -``` - -### Cloud Run - -```bash -gcloud run services update warren \ - --set-env-vars="WARREN_AGENT_FALCON_CLIENT_ID=your-client-id" \ - --set-env-vars="WARREN_AGENT_FALCON_CLIENT_SECRET=your-client-secret" \ - --set-env-vars="WARREN_AGENT_FALCON_BASE_URL=https://api.crowdstrike.com" -``` - -> **Security Note:** For production, use Secret Manager for the client secret rather than plain environment variables: -> ```bash -> gcloud run services update warren \ -> --set-secrets="WARREN_AGENT_FALCON_CLIENT_SECRET=falcon-secret:latest" -> ``` - -## Authentication - -The agent uses [OAuth2 Client Credentials Flow](https://www.falconpy.io/Service-Collections/OAuth2.html) to authenticate with the CrowdStrike API: - -1. Sends `client_id` and `client_secret` to `POST /oauth2/token` -2. Receives a bearer token (valid for 30 minutes) -3. Automatically refreshes the token before expiry - -No manual token management is required. - -## Available Tools - -The agent exposes the following read-only tools to the LLM: - -### Incidents - -| Tool | Description | -|---|---| -| `falcon_search_incidents` | Search for incidents using FQL filters | -| `falcon_get_incidents` | Get detailed incident information by ID | - -### Alerts - -| Tool | Description | -|---|---| -| `falcon_search_alerts` | Search and retrieve alerts (combined endpoint with cursor pagination) | -| `falcon_get_alerts` | Get detailed alert information by composite ID | - -### Behaviors - -| Tool | Description | -|---|---| -| `falcon_search_behaviors` | Search for behaviors using FQL filters | -| `falcon_get_behaviors` | Get detailed behavior information by ID | - -### CrowdScores - -| Tool | Description | -|---|---| -| `falcon_get_crowdscores` | Get CrowdScore values for your environment | - -### Events (EDR Telemetry) - -| Tool | Description | -|---|---| -| `falcon_search_events` | Search raw EDR events using CQL (process executions, network connections, DNS, file writes, etc.) | - -> **Note:** Event search uses the Next-Gen SIEM Search API, which runs queries asynchronously. The tool handles job creation and polling internally — results are returned once the search completes. - -## Troubleshooting - -### Agent not appearing in available tools - -Verify that both `WARREN_AGENT_FALCON_CLIENT_ID` and `WARREN_AGENT_FALCON_CLIENT_SECRET` are set. The agent is skipped if either is empty. - -### 401 Unauthorized errors - -- Verify your Client ID and Client Secret are correct -- Check that the API client has not been revoked in the Falcon console -- Ensure the base URL matches your Falcon tenant region - -### 403 Forbidden errors - -- Verify the API client has the required scopes (`Incidents: Read`, `Alerts: Read`, `NGSIEM: Read + Write`) -- Check that the API client is active and not expired - -### Empty results - -- Verify FQL filter syntax (values must be quoted, e.g., `status:'new'` not `status:new`) -- For event search, verify CQL syntax (e.g., `aid=abc123`, `#event_simpleName=ProcessRollup2`) -- Check the time range — CrowdStrike may have data retention limits -- Use simpler filters first, then refine diff --git a/pkg/agents/falcon/agent.go b/pkg/agents/falcon/agent.go deleted file mode 100644 index 747537a2..00000000 --- a/pkg/agents/falcon/agent.go +++ /dev/null @@ -1,32 +0,0 @@ -package falcon - -import ( - "context" - - "github.com/gollem-dev/gollem" -) - -// toolSet implements interfaces.ToolSet for CrowdStrike Falcon. -type toolSet struct { - internal *internalTool -} - -func (ts *toolSet) ID() string { - return "falcon" -} - -func (ts *toolSet) Description() string { - return "CrowdStrike Falcon EDR queries for incidents, alerts, behaviors, devices, and raw events" -} - -func (ts *toolSet) Prompt(_ context.Context) (string, error) { - return buildSystemPrompt(), nil -} - -func (ts *toolSet) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return ts.internal.Specs(ctx) -} - -func (ts *toolSet) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - return ts.internal.Run(ctx, name, args) -} diff --git a/pkg/agents/falcon/auth.go b/pkg/agents/falcon/auth.go deleted file mode 100644 index 396630e0..00000000 --- a/pkg/agents/falcon/auth.go +++ /dev/null @@ -1,130 +0,0 @@ -package falcon - -import ( - "context" - "encoding/json" - "io" - "net/http" - "net/url" - "strings" - "sync" - "time" - - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/utils/logging" - "github.com/secmon-lab/warren/pkg/utils/safe" -) - -// tokenProvider manages OAuth2 Client Credentials Flow for CrowdStrike Falcon API. -// It handles token acquisition and automatic refresh before expiry. -type tokenProvider struct { - clientID string - clientSecret string - baseURL string - httpClient *http.Client - - mu sync.Mutex - token string - expiry time.Time -} - -// tokenResponse represents the OAuth2 token response from CrowdStrike. -type tokenResponse struct { - AccessToken string `json:"access_token"` - ExpiresIn int `json:"expires_in"` - TokenType string `json:"token_type"` -} - -// newTokenProvider creates a new tokenProvider for CrowdStrike OAuth2 authentication. -func newTokenProvider(clientID, clientSecret, baseURL string) *tokenProvider { - return &tokenProvider{ - clientID: clientID, - clientSecret: clientSecret, - baseURL: baseURL, - httpClient: &http.Client{Timeout: 30 * time.Second}, - } -} - -// getToken returns a valid bearer token, refreshing if necessary. -func (tp *tokenProvider) getToken(ctx context.Context) (string, error) { - tp.mu.Lock() - defer tp.mu.Unlock() - - // Return cached token if still valid (with 30-second buffer) - if tp.token != "" && time.Now().Add(30*time.Second).Before(tp.expiry) { - return tp.token, nil - } - - if err := tp.refreshToken(ctx); err != nil { - return "", err - } - - return tp.token, nil -} - -// clearToken invalidates the cached token, forcing a refresh on next getToken call. -func (tp *tokenProvider) clearToken() { - tp.mu.Lock() - defer tp.mu.Unlock() - tp.token = "" - tp.expiry = time.Time{} -} - -// refreshToken acquires a new OAuth2 token from the CrowdStrike API. -func (tp *tokenProvider) refreshToken(ctx context.Context) error { - log := logging.From(ctx) - log.Debug("Refreshing CrowdStrike OAuth2 token", - "base_url", tp.baseURL, - "client_id", tp.clientID, - ) - - form := url.Values{ - "client_id": {tp.clientID}, - "client_secret": {tp.clientSecret}, - } - - req, err := http.NewRequestWithContext( - ctx, - http.MethodPost, - tp.baseURL+"/oauth2/token", - strings.NewReader(form.Encode()), - ) - if err != nil { - return goerr.Wrap(err, "failed to create token request") - } - - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - - resp, err := tp.httpClient.Do(req) - if err != nil { - return goerr.Wrap(err, "failed to send token request") - } - defer safe.Close(ctx, resp.Body) - - body, err := io.ReadAll(resp.Body) - if err != nil { - return goerr.Wrap(err, "failed to read token response body") - } - - if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { - log.Warn("CrowdStrike OAuth2 token request failed", - "status", resp.StatusCode, - ) - return goerr.New("OAuth2 token request failed", - goerr.V("status", resp.StatusCode), - goerr.V("body", string(body)), - ) - } - - var tokenResp tokenResponse - if err := json.Unmarshal(body, &tokenResp); err != nil { - return goerr.Wrap(err, "failed to parse token response") - } - - tp.token = tokenResp.AccessToken - tp.expiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) - - log.Debug("CrowdStrike OAuth2 token refreshed", "expires_in", tokenResp.ExpiresIn) - - return nil -} diff --git a/pkg/agents/falcon/auth_test.go b/pkg/agents/falcon/auth_test.go deleted file mode 100644 index 2f72d965..00000000 --- a/pkg/agents/falcon/auth_test.go +++ /dev/null @@ -1,164 +0,0 @@ -package falcon_test - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "sync" - "sync/atomic" - "testing" - - "github.com/m-mizutani/gt" - "github.com/secmon-lab/warren/pkg/agents/falcon" -) - -func TestTokenProvider_GetToken(t *testing.T) { - var callCount atomic.Int32 - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - callCount.Add(1) - - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/oauth2/token") - gt.Equal(t, r.Header.Get("Content-Type"), "application/x-www-form-urlencoded") - - err := r.ParseForm() - gt.NoError(t, err) - gt.Equal(t, r.FormValue("client_id"), "test-client-id") - gt.Equal(t, r.FormValue("client_secret"), "test-client-secret") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "access_token": "test-token-123", - "expires_in": 1800, - "token_type": "bearer", - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - })) - defer srv.Close() - - tp := falcon.NewTokenProviderForTest( - "test-client-id", - "test-client-secret", - srv.URL, - ) - - ctx := context.Background() - - // First call should fetch a new token - token, err := tp.GetToken(ctx) - gt.NoError(t, err) - gt.Equal(t, token, "test-token-123") - gt.Equal(t, callCount.Load(), int32(1)) - - // Second call should return cached token - token, err = tp.GetToken(ctx) - gt.NoError(t, err) - gt.Equal(t, token, "test-token-123") - gt.Equal(t, callCount.Load(), int32(1)) // No additional call -} - -func TestTokenProvider_GetToken_ConcurrentAccess(t *testing.T) { - var callCount atomic.Int32 - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - callCount.Add(1) - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "access_token": "concurrent-token", - "expires_in": 1800, - "token_type": "bearer", - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - })) - defer srv.Close() - - tp := falcon.NewTokenProviderForTest( - "test-client-id", - "test-client-secret", - srv.URL, - ) - - ctx := context.Background() - - var wg sync.WaitGroup - for i := 0; i < 10; i++ { - wg.Add(1) - go func() { - defer wg.Done() - token, err := tp.GetToken(ctx) - gt.NoError(t, err) - gt.Equal(t, token, "concurrent-token") - }() - } - wg.Wait() - - // Due to mutex, should have minimal API calls (ideally 1, but could be a few due to timing) - gt.True(t, callCount.Load() <= 10) -} - -func TestTokenProvider_GetToken_ClearAndRefresh(t *testing.T) { - var callCount atomic.Int32 - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - count := callCount.Add(1) - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "access_token": "token-" + string(rune('0'+count)), - "expires_in": 1800, - "token_type": "bearer", - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - })) - defer srv.Close() - - tp := falcon.NewTokenProviderForTest( - "test-client-id", - "test-client-secret", - srv.URL, - ) - - ctx := context.Background() - - // Get initial token - token1, err := tp.GetToken(ctx) - gt.NoError(t, err) - gt.Equal(t, callCount.Load(), int32(1)) - - // Clear token - tp.ClearToken() - - // Should fetch a new token - token2, err := tp.GetToken(ctx) - gt.NoError(t, err) - gt.Equal(t, callCount.Load(), int32(2)) - - // Tokens should be different - gt.True(t, token1 != token2) -} - -func TestTokenProvider_GetToken_ServerError(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - _, err := w.Write([]byte(`{"errors": [{"message": "internal error"}]}`)) - gt.NoError(t, err) - })) - defer srv.Close() - - tp := falcon.NewTokenProviderForTest( - "test-client-id", - "test-client-secret", - srv.URL, - ) - - ctx := context.Background() - - _, err := tp.GetToken(ctx) - gt.Error(t, err) -} diff --git a/pkg/agents/falcon/export_test.go b/pkg/agents/falcon/export_test.go deleted file mode 100644 index acef1bff..00000000 --- a/pkg/agents/falcon/export_test.go +++ /dev/null @@ -1,52 +0,0 @@ -package falcon - -import "context" - -// TokenProviderForTest wraps tokenProvider for testing purposes. -type TokenProviderForTest struct { - tp *tokenProvider -} - -// NewTokenProviderForTest creates a tokenProvider wrapper for testing. -func NewTokenProviderForTest(clientID, clientSecret, baseURL string) *TokenProviderForTest { - return &TokenProviderForTest{ - tp: newTokenProvider(clientID, clientSecret, baseURL), - } -} - -// GetToken returns a valid bearer token. -func (t *TokenProviderForTest) GetToken(ctx context.Context) (string, error) { - return t.tp.getToken(ctx) -} - -// ClearToken invalidates the cached token. -func (t *TokenProviderForTest) ClearToken() { - t.tp.clearToken() -} - -// InternalToolForTest wraps internalTool for testing purposes. -type InternalToolForTest struct { - tool *internalTool -} - -// NewInternalToolForTest creates an internalTool wrapper for testing. -func NewInternalToolForTest(clientID, clientSecret, baseURL string) *InternalToolForTest { - tp := newTokenProvider(clientID, clientSecret, baseURL) - return &InternalToolForTest{ - tool: newInternalTool(tp, baseURL), - } -} - -// Run executes a tool by name with the given arguments. -func (t *InternalToolForTest) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - return t.tool.Run(ctx, name, args) -} - -// SpecCount returns the number of tool specs. -func (t *InternalToolForTest) SpecCount(ctx context.Context) (int, error) { - specs, err := t.tool.Specs(ctx) - if err != nil { - return 0, err - } - return len(specs), nil -} diff --git a/pkg/agents/falcon/extract.go b/pkg/agents/falcon/extract.go deleted file mode 100644 index 8ca1f9ff..00000000 --- a/pkg/agents/falcon/extract.go +++ /dev/null @@ -1 +0,0 @@ -package falcon diff --git a/pkg/agents/falcon/factory.go b/pkg/agents/falcon/factory.go deleted file mode 100644 index 8f4c39ea..00000000 --- a/pkg/agents/falcon/factory.go +++ /dev/null @@ -1,71 +0,0 @@ -package falcon - -import ( - "context" - - "github.com/secmon-lab/warren/pkg/domain/interfaces" - "github.com/secmon-lab/warren/pkg/utils/logging" - "github.com/urfave/cli/v3" -) - -const defaultBaseURL = "https://api.crowdstrike.com" - -// Factory implements agents.ToolSetFactory interface for CrowdStrike Falcon. -type Factory struct { - clientID string - clientSecret string - baseURL string -} - -// Flags implements agents.ToolSetFactory. -func (f *Factory) Flags() []cli.Flag { - return []cli.Flag{ - &cli.StringFlag{ - Name: "agent-falcon-client-id", - Usage: "CrowdStrike Falcon API Client ID", - Destination: &f.clientID, - Category: "Agent:Falcon", - Sources: cli.EnvVars("WARREN_AGENT_FALCON_CLIENT_ID"), - }, - &cli.StringFlag{ - Name: "agent-falcon-client-secret", - Usage: "CrowdStrike Falcon API Client Secret", - Destination: &f.clientSecret, - Category: "Agent:Falcon", - Sources: cli.EnvVars("WARREN_AGENT_FALCON_CLIENT_SECRET"), - }, - &cli.StringFlag{ - Name: "agent-falcon-base-url", - Usage: "CrowdStrike Falcon API Base URL (default: US-1)", - Destination: &f.baseURL, - Category: "Agent:Falcon", - Sources: cli.EnvVars("WARREN_AGENT_FALCON_BASE_URL"), - Value: defaultBaseURL, - }, - } -} - -// Configure implements agents.ToolSetFactory. -// Returns (nil, nil) if client_id or client_secret is not set. -func (f *Factory) Configure(ctx context.Context) (interfaces.ToolSet, error) { - if f.clientID == "" || f.clientSecret == "" { - return nil, nil - } - - baseURL := f.baseURL - if baseURL == "" { - baseURL = defaultBaseURL - } - - tp := newTokenProvider(f.clientID, f.clientSecret, baseURL) - - logging.From(ctx).Info("CrowdStrike Falcon Agent configured", - "base_url", baseURL, - "client_id", f.clientID, - "client_secret_length", len(f.clientSecret), - ) - - return &toolSet{ - internal: newInternalTool(tp, baseURL), - }, nil -} diff --git a/pkg/agents/falcon/prompt.go b/pkg/agents/falcon/prompt.go deleted file mode 100644 index 86cfeea4..00000000 --- a/pkg/agents/falcon/prompt.go +++ /dev/null @@ -1,13 +0,0 @@ -package falcon - -import ( - _ "embed" -) - -//go:embed prompt/system.md -var systemPromptTemplate string - -// buildSystemPrompt returns the system prompt for the Falcon agent. -func buildSystemPrompt() string { - return systemPromptTemplate -} diff --git a/pkg/agents/falcon/prompt/extract.md b/pkg/agents/falcon/prompt/extract.md deleted file mode 100644 index 3db82ea3..00000000 --- a/pkg/agents/falcon/prompt/extract.md +++ /dev/null @@ -1,53 +0,0 @@ -# Task -Extract the raw CrowdStrike Falcon records from the conversation history and return them as a JSON object with a "records" field. - -# Key Principle: Understand User Intent -- Do NOT just extract data literally based on the user's wording -- UNDERSTAND what the user actually wants to know or achieve -- INTERPRET the user's intent from their request and the context -- Select records that FULFILL the user's actual needs, not just match keywords - -# Guidelines -1. First, understand the user's true intent from their original request -2. Look for API response data in the conversation history (e.g., `falcon_search_alerts`, `falcon_get_incidents` tool responses) -3. Parse the data carefully — responses contain nested structures with resources/entities -4. Convert each record into a JSON object preserving the key fields -5. Preserve ALL relevant field names and values exactly as they appear -6. If multiple API calls were made, intelligently combine and deduplicate records -7. **IMPORTANT**: Each record MUST be a non-empty object with the actual field names and values - -# Record Types - -## Incidents -Key fields: incident ID, status, fine_score, start/end time, hosts, users, tactics, techniques - -## Alerts -Key fields: composite_id, status, severity, timestamp, hostname, tactic, technique, filename, cmdline, description - -## Behaviors -Key fields: behavior_id, tactic, technique, severity, pattern_disposition, device info, timestamp - -## CrowdScores -Key fields: id, timestamp, score, adjusted_score - -# Output Format -Return a JSON object with a "records" field: -```json -{ - "records": [ - { - "field1": "value1", - "field2": "value2" - } - ] -} -``` - -# Important -- The output must be a JSON OBJECT with a "records" field -- The "records" field contains an ARRAY of record objects -- Each record object must contain the actual field names and values from the API response -- **DO NOT return empty objects** — each object must contain the actual field values -- If no records match the user's intent, return: `{"records": []}` -- Flatten deeply nested structures where appropriate for readability -- Focus on user's INTENT, not literal wording diff --git a/pkg/agents/falcon/prompt/system.md b/pkg/agents/falcon/prompt/system.md deleted file mode 100644 index f480b665..00000000 --- a/pkg/agents/falcon/prompt/system.md +++ /dev/null @@ -1,207 +0,0 @@ -# CrowdStrike Falcon Investigation Agent - -You are a CrowdStrike Falcon investigation agent. Your role is to query the Falcon API to retrieve security incident, alert, behavior, device/host, and CrowdScore data for threat investigation. - -## Core Principles - -### Efficiency First -- Use the most direct query path to get the needed data -- Prefer `falcon_search_alerts` (combined endpoint) over separate search+get calls -- Start with specific filters and broaden only if no results are found - -### Data Fidelity -- Return query results in their raw form without interpretation -- Do not add threat assessments or security recommendations -- Only add factual metadata (e.g., "No results found", "Showing 50 of 234 total") - -## Available Tools - -### Incidents -- `falcon_search_incidents` — Search for incident IDs using FQL filters -- `falcon_get_incidents` — Get full incident details by ID - -### Alerts -- `falcon_search_alerts` — Search and retrieve alerts in one call (preferred for most queries) -- `falcon_get_alerts` — Get alert details by composite ID - -### Behaviors -- `falcon_search_behaviors` — Search for behavior IDs using FQL filters -- `falcon_get_behaviors` — Get full behavior details by ID - -### Devices (Hosts) -- `falcon_search_devices` — Search for device/host IDs using FQL filters (hostname, IP, OS, platform, last_seen, tags, etc.) -- `falcon_get_devices` — Get full device details by ID (hostname, OS version, IP addresses, sensor version, containment status, tags) - -### CrowdScores -- `falcon_get_crowdscores` — Get environment CrowdScore values - -### Events (EDR Telemetry) -- `falcon_search_events` — Search raw EDR events using CQL (CrowdStrike Query Language). Queries process executions, network connections, file writes, DNS requests, registry changes, and more. The search runs asynchronously but results are returned automatically. - -## FQL (Falcon Query Language) Reference - -**Note:** FQL is used for Incidents, Alerts, Behaviors, and Devices. For raw event search, use CQL (see below). - -### Syntax -- String values must be quoted: `status:'new'` -- Numeric comparisons: `severity:>50`, `severity:<=80` -- Date comparisons: `start:>'2025-01-01'`, `created_on:>'2025-01-01T00:00:00Z'` -- Wildcard: `hostname:'*web*'` -- Negation: `status:!'closed'` -- Logical operators: `+` (AND), `,` (OR) -- Combine: `status:'new'+severity:>50` - -### Common Incident Fields -- `status` — 20: New, 25: Reopened, 30: In Progress, 40: Closed -- `start`, `end` — Incident time range -- `state` — open, closed -- `tags` — User-assigned tags -- `fine_score` — Incident score (0-100) -- `assigned_to_name` — Assigned analyst -- `host_ids` — Host agent IDs associated with the incident - -**Important:** Incidents do NOT support hash-based filters (`sha256`, `file_hash`, `md5`). To find incidents related to a file hash, first search alerts by hash, then use the incident IDs from alert results. - -### Common Alert Fields -- `status` — new, in_progress, closed, reopened -- `severity` — Numeric severity (1-100) -- `type` — Alert type (e.g., ldt for Lateral Movement Detection) -- `tactics` — MITRE ATT&CK tactics -- `techniques` — MITRE ATT&CK techniques -- `timestamp` — Alert timestamp -- `hostname` — Source device hostname -- `filename` — Associated filename -- `sha256` — File hash -- `cmdline` — Command line - -Alerts support the widest range of filter fields including `sha256`, `hostname`, `filename`, `cmdline`, etc. - -### Common Behavior Fields -- `tactic`, `technique` — MITRE ATT&CK mapping -- `severity` — Behavior severity -- `pattern_disposition` — Action taken (e.g., detect, block) -- `behavior_id` — Behavior ID -- `incident_id` — Associated incident ID - -**Important:** Behaviors do NOT support hash-based filters (`sha256`, `md5`). To find behaviors related to a file hash, first search alerts by hash, then use the behavior IDs from the alert's `behaviors` field, or search incidents and retrieve their behaviors. - -### Common Device Fields -- `hostname` — Device hostname -- `platform_name` — OS platform (Windows, Mac, Linux) -- `os_version` — OS version string -- `external_ip` — External IP address -- `local_ip` — Local IP address -- `status` — Device status (normal, containment_pending, contained, lift_containment_pending) -- `last_seen` — Last time the sensor checked in -- `first_seen` — First time the sensor was seen -- `device_id` — Unique device/agent ID -- `tags` — FalconGroupingTags and SensorGroupingTags -- `machine_domain` — Active Directory domain -- `ou` — Organizational unit - -### Filter Field Compatibility - -| Filter Field | Alerts | Incidents | Behaviors | Devices | -|---|---|---|---|---| -| `sha256` | ✅ | ❌ | ❌ | ❌ | -| `hostname` | ✅ | ❌ | ❌ | ✅ | -| `filename` | ✅ | ❌ | ❌ | ❌ | -| `status` | ✅ | ✅ | ❌ | ✅ | -| `severity` | ✅ | ❌ | ✅ | ❌ | -| `host_ids` | ❌ | ✅ | ❌ | ❌ | -| `tactic` | ✅ | ❌ | ✅ | ❌ | -| `platform_name` | ❌ | ❌ | ❌ | ✅ | -| `external_ip` | ❌ | ❌ | ❌ | ✅ | -| `last_seen` | ❌ | ❌ | ❌ | ✅ | -| `tags` | ❌ | ✅ | ❌ | ✅ | - -## CQL (CrowdStrike Query Language) Reference - -CQL is used with `falcon_search_events` to query raw EDR telemetry data. CQL is based on the LogScale Query Language. - -### Basic Syntax -- Field filtering: `aid=abc123`, `#event_simpleName=ProcessRollup2` -- String matching: `FileName="cmd.exe"`, `CommandLine="*powershell*"` -- Logical operators: `AND`, `OR`, `NOT` -- Pipe for transformations: `aid=abc123 | tail(100)` -- Wildcards in values: `FileName="*.exe"` - -### Common Event Types (#event_simpleName) -- `ProcessRollup2` — Process execution events -- `NetworkConnectIP4`, `NetworkConnectIP6` — Network connections -- `DnsRequest` — DNS queries -- `FileWritten` — File write operations -- `RegistryOperationKey`, `RegistryOperationValue` — Registry changes -- `UserLogon`, `UserLogoff` — Authentication events -- `ScriptControlScan` — Script execution monitoring -- `SyntheticProcessRollup2` — Synthetic process events - -### Common Event Fields -- `aid` — Agent/sensor ID -- `ComputerName` — Hostname -- `UserName` — User account -- `FileName` — File or process name -- `FilePath` — Full file path -- `CommandLine` — Command line arguments -- `SHA256HashData` — File SHA256 hash -- `MD5HashData` — File MD5 hash -- `LocalAddressIP4`, `RemoteAddressIP4` — Network addresses -- `RemotePort` — Destination port -- `DomainName` — DNS domain -- `timestamp` — Event timestamp - -### Repositories -- `search-all` — All data (default, recommended) -- `investigate_view` — Falcon EDR endpoint events -- `third-party` — Third-party data sources -- `falcon_for_it_view` — IT Automation data -- `forensics_view` — Forensics triage data - -### CQL Examples -- All process events on a host: `ComputerName="workstation1" AND #event_simpleName=ProcessRollup2` -- Network connections to a specific IP: `RemoteAddressIP4="10.0.0.1" AND #event_simpleName=NetworkConnectIP4` -- DNS queries for a domain: `DomainName="*.malicious.com" AND #event_simpleName=DnsRequest` -- PowerShell executions: `FileName="powershell.exe" AND #event_simpleName=ProcessRollup2 | tail(50)` -- Events by agent ID in last 24h: `aid=abc123` (use start="1d" parameter) - -## Standard Investigation Workflow - -### 1. Understand the Request -- What is being investigated? (incident, alert type, hostname, user, etc.) -- What time frame is relevant? -- What detail level is needed? (IDs only, or full details?) - -### 2. Choose the Right Tool -- **For alerts**: Use `falcon_search_alerts` first (gets details in one call) -- **For incidents**: Use `falcon_search_incidents` then `falcon_get_incidents` -- **For behaviors**: Use `falcon_search_behaviors` then `falcon_get_behaviors` -- **For devices/hosts** (hostname, IP, OS, sensor info): Use `falcon_search_devices` then `falcon_get_devices` -- **For raw EDR events** (process, network, file, DNS, etc.): Use `falcon_search_events` with CQL -- **For overall threat level**: Use `falcon_get_crowdscores` - -### 3. Build Effective FQL Queries -- Start with the most specific filter -- Combine multiple conditions with `+` (AND) -- Use time bounds to limit results -- Apply severity/status filters to focus on actionable items - -### 4. Handle Results -- If search returns IDs, follow up with the corresponding get endpoint -- Note pagination metadata (total count, offset) for large result sets -- Include key identifiers in your response (incident ID, hostname, severity, etc.) - -## Response Format - -Return results containing: -- The actual data records from the API -- Key fields relevant to the investigation -- Pagination info if there are more results -- Any execution notes (e.g., "Filtered to last 7 days") - -Do not include: -- Threat assessments or risk ratings -- Recommendations or remediation steps -- Opinions about the data -- Security insights beyond what the data shows - -Present only factual query results and execution information. diff --git a/pkg/agents/falcon/tool.go b/pkg/agents/falcon/tool.go deleted file mode 100644 index 95f77240..00000000 --- a/pkg/agents/falcon/tool.go +++ /dev/null @@ -1,664 +0,0 @@ -package falcon - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/url" - "strings" - "time" - - "github.com/gollem-dev/gollem" - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/utils/logging" - "github.com/secmon-lab/warren/pkg/utils/msg" - "github.com/secmon-lab/warren/pkg/utils/safe" -) - -// internalTool implements gollem.ToolSet for CrowdStrike Falcon API operations. -type internalTool struct { - tokenProvider *tokenProvider - baseURL string - httpClient *http.Client -} - -// newInternalTool creates a new internalTool for Falcon API calls. -func newInternalTool(tp *tokenProvider, baseURL string) *internalTool { - return &internalTool{ - tokenProvider: tp, - baseURL: baseURL, - httpClient: &http.Client{Timeout: 60 * time.Second}, - } -} - -func (t *internalTool) Specs(_ context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "falcon_search_incidents", - Description: "Search for incident IDs using FQL (Falcon Query Language) filters. Returns a list of incident IDs that can be used with falcon_get_incidents to retrieve full details.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"status:'30'\", \"tags:'critical'\", \"start:>'2025-01-01'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort expression (e.g., \"start.desc\", \"end.asc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 500)", - }, - "offset": { - Type: gollem.TypeNumber, - Description: "Pagination offset", - }, - }, - }, - { - Name: "falcon_get_incidents", - Description: "Get detailed information for specific incidents by their IDs. Returns full incident details including status, tactics, techniques, hosts, and users involved.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated incident IDs (e.g., \"inc:abc123:def456,inc:abc123:ghi789\")", - Required: true, - }, - }, - }, - { - Name: "falcon_search_alerts", - Description: "Search and retrieve alert details in one call using FQL filters with cursor-based pagination. Returns full alert objects including severity, tactic, technique, and device info.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"status:'new'\", \"severity:>50\", \"tactics:'Lateral Movement'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort property (e.g., \"timestamp|desc\", \"severity|asc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of alerts to return (default: 100, max: 1000)", - }, - "after": { - Type: gollem.TypeString, - Description: "Cursor pagination token from a previous response for fetching next page", - }, - }, - }, - { - Name: "falcon_get_alerts", - Description: "Get detailed alert information by composite IDs. Use this when you already have specific alert IDs.", - Parameters: map[string]*gollem.Parameter{ - "composite_ids": { - Type: gollem.TypeString, - Description: "Comma-separated composite alert IDs", - Required: true, - }, - }, - }, - { - Name: "falcon_search_behaviors", - Description: "Search for behavior IDs using FQL filters. Returns behavior IDs that can be used with falcon_get_behaviors for full details.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 500)", - }, - "offset": { - Type: gollem.TypeNumber, - Description: "Pagination offset", - }, - }, - }, - { - Name: "falcon_get_behaviors", - Description: "Get detailed behavior information by IDs. Returns behavior details including tactic, technique, severity, pattern, and associated device info.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated behavior IDs", - Required: true, - }, - }, - }, - { - Name: "falcon_get_crowdscores", - Description: "Get CrowdScore values for the environment. CrowdScore is an overall threat level indicator.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"timestamp:>'2025-01-01'\")", - }, - }, - }, - { - Name: "falcon_search_devices", - Description: "Search for device (host) IDs using FQL filters. Returns a list of device IDs that can be used with falcon_get_devices to retrieve full host details including OS, IP addresses, sensor version, and containment status.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"hostname:'*web*'\", \"platform_name:'Windows'\", \"last_seen:>='2025-01-01'\", \"external_ip:'10.0.0.*'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort expression (e.g., \"hostname.asc\", \"last_seen.desc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 5000)", - }, - "offset": { - Type: gollem.TypeString, - Description: "Pagination offset token from a previous response", - }, - }, - }, - { - Name: "falcon_get_devices", - Description: "Get detailed device (host) information by device IDs. Returns full host details including hostname, OS, IP addresses, sensor version, tags, and containment status.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated device IDs (up to 5000, e.g., \"abc123def456,ghi789jkl012\")", - Required: true, - }, - }, - }, - { - Name: "falcon_search_events", - Description: "Search EDR telemetry events using CrowdStrike Query Language (CQL). This uses the Next-Gen SIEM Search API to query raw event data (process executions, network connections, file writes, DNS requests, etc.). The search runs asynchronously and this tool automatically polls until results are ready.", - Parameters: map[string]*gollem.Parameter{ - "query_string": { - Type: gollem.TypeString, - Description: "CQL query string (e.g., \"aid=abc123\", \"#event_simpleName=ProcessRollup2 AND FileName=cmd.exe\", \"ComputerName=workstation1 | tail(100)\")", - Required: true, - }, - "repository": { - Type: gollem.TypeString, - Description: "Repository to search. Values: \"search-all\" (all data, default), \"investigate_view\" (Falcon EDR), \"third-party\" (third-party data), \"falcon_for_it_view\" (IT Automation), \"forensics_view\" (Forensics)", - }, - "start": { - Type: gollem.TypeString, - Description: "Start time for the search (e.g., \"1d\" for last 1 day, \"24h\" for last 24 hours, \"2025-01-01T00:00:00Z\" for absolute time). Default: \"1d\"", - }, - "end": { - Type: gollem.TypeString, - Description: "End time for the search (e.g., \"now\", \"2025-01-02T00:00:00Z\"). Default: \"now\"", - }, - }, - }, - }, nil -} - -func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case "falcon_search_incidents": - return t.searchIncidents(ctx, args) - case "falcon_get_incidents": - return t.getIncidents(ctx, args) - case "falcon_search_alerts": - return t.searchAlerts(ctx, args) - case "falcon_get_alerts": - return t.getAlerts(ctx, args) - case "falcon_search_behaviors": - return t.searchBehaviors(ctx, args) - case "falcon_get_behaviors": - return t.getBehaviors(ctx, args) - case "falcon_search_devices": - return t.searchDevices(ctx, args) - case "falcon_get_devices": - return t.getDevices(ctx, args) - case "falcon_get_crowdscores": - return t.getCrowdScores(ctx, args) - case "falcon_search_events": - return t.searchEvents(ctx, args) - default: - return nil, goerr.New("unknown tool name", goerr.V("name", name)) - } -} - -// doRequest performs an authenticated HTTP request to the CrowdStrike API. -// On 401 responses, it clears the token and retries once. -func (t *internalTool) doRequest(ctx context.Context, method, path string, body any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Falcon API request", "method", method, "path", path) - - result, err := t.doRequestOnce(ctx, method, path, body) - if err == nil { - return result, nil - } - - // Check if this is an authentication error and retry once - var apiErr *apiError - if errors.As(err, &apiErr) && apiErr.statusCode == http.StatusUnauthorized { - log.Debug("Received 401, clearing token and retrying") - msg.Trace(ctx, "🔄 Received 401, refreshing token and retrying...") - t.tokenProvider.clearToken() - return t.doRequestOnce(ctx, method, path, body) - } - - return nil, err -} - -// apiError represents an HTTP error from the CrowdStrike API. -type apiError struct { - statusCode int - body string -} - -func (e *apiError) Error() string { - return fmt.Sprintf("Falcon API error: status=%d", e.statusCode) -} - -// doRequestOnce performs a single authenticated HTTP request. -func (t *internalTool) doRequestOnce(ctx context.Context, method, path string, body any) (map[string]any, error) { - log := logging.From(ctx) - token, err := t.tokenProvider.getToken(ctx) - if err != nil { - return nil, goerr.Wrap(err, "failed to get auth token") - } - - var reqBody io.Reader - if body != nil { - jsonBody, err := json.Marshal(body) - if err != nil { - return nil, goerr.Wrap(err, "failed to marshal request body") - } - reqBody = bytes.NewReader(jsonBody) - } - - req, err := http.NewRequestWithContext(ctx, method, t.baseURL+path, reqBody) - if err != nil { - return nil, goerr.Wrap(err, "failed to create request") - } - - req.Header.Set("Authorization", "Bearer "+token) - if body != nil { - req.Header.Set("Content-Type", "application/json") - } - req.Header.Set("Accept", "application/json") - - resp, err := t.httpClient.Do(req) - if err != nil { - return nil, goerr.Wrap(err, "failed to send request", goerr.V("path", path)) - } - defer safe.Close(ctx, resp.Body) - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, goerr.Wrap(err, "failed to read response body") - } - - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - log.Warn("Falcon API request failed", - "status", resp.StatusCode, - "path", path, - ) - return nil, goerr.Wrap(&apiError{statusCode: resp.StatusCode, body: string(respBody)}, - "Falcon API request failed", - goerr.V("status", resp.StatusCode), - goerr.V("path", path), - ) - } - - var result map[string]any - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, goerr.Wrap(err, "failed to parse response JSON", goerr.V("body", string(respBody))) - } - - return result, nil -} - -// searchIncidents searches for incident IDs using FQL filters. -func (t *internalTool) searchIncidents(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) - msg.Trace(ctx, "🔍 Searching incidents (filter: `%s`)", filter) - - path := "/incidents/queries/incidents/v1" - params := buildQueryParams(args, "filter", "sort", "limit", "offset") - if params != "" { - path += "?" + params - } - result, err := t.doRequest(ctx, http.MethodGet, path, nil) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Incident search failed (filter: `%s`): %v", filter, err) - return nil, err - } - msg.Trace(ctx, "✅ Incident search completed") - return result, nil -} - -// getIncidents retrieves incident details by IDs. -func (t *internalTool) getIncidents(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { - return nil, goerr.New("ids is required") - } - - msg.Trace(ctx, "📋 Retrieving incident details (ids: `%s`)", ids) - body := map[string]any{ - "ids": splitAndTrim(ids), - } - result, err := t.doRequest(ctx, http.MethodPost, "/incidents/entities/incidents/GET/v1", body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to retrieve incidents (ids: `%s`): %v", ids, err) - return nil, err - } - msg.Trace(ctx, "✅ Retrieved incident details") - return result, nil -} - -// searchAlerts searches and retrieves alert details using FQL filters. -func (t *internalTool) searchAlerts(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) - msg.Trace(ctx, "🔍 Searching alerts (filter: `%s`)", filter) - - body := make(map[string]any) - if filter != "" { - body["filter"] = filter - } - if sort, ok := args["sort"].(string); ok && sort != "" { - body["sort"] = sort - } - if limit, ok := args["limit"].(float64); ok { - body["limit"] = int(limit) - } - if after, ok := args["after"].(string); ok && after != "" { - body["after"] = after - } - result, err := t.doRequest(ctx, http.MethodPost, "/alerts/combined/alerts/v1", body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Alert search failed (filter: `%s`): %v", filter, err) - return nil, err - } - msg.Trace(ctx, "✅ Alert search completed") - return result, nil -} - -// getAlerts retrieves alert details by composite IDs. -func (t *internalTool) getAlerts(ctx context.Context, args map[string]any) (map[string]any, error) { - compositeIDs, ok := args["composite_ids"].(string) - if !ok || compositeIDs == "" { - return nil, goerr.New("composite_ids is required") - } - - msg.Trace(ctx, "📋 Retrieving alert details (ids: `%s`)", compositeIDs) - body := map[string]any{ - "composite_ids": splitAndTrim(compositeIDs), - } - result, err := t.doRequest(ctx, http.MethodPost, "/alerts/entities/alerts/v2", body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to retrieve alerts (ids: `%s`): %v", compositeIDs, err) - return nil, err - } - msg.Trace(ctx, "✅ Retrieved alert details") - return result, nil -} - -// searchBehaviors searches for behavior IDs using FQL filters. -func (t *internalTool) searchBehaviors(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) - msg.Trace(ctx, "🔍 Searching behaviors (filter: `%s`)", filter) - - path := "/incidents/queries/behaviors/v1" - params := buildQueryParams(args, "filter", "limit", "offset") - if params != "" { - path += "?" + params - } - result, err := t.doRequest(ctx, http.MethodGet, path, nil) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Behavior search failed (filter: `%s`): %v", filter, err) - return nil, err - } - msg.Trace(ctx, "✅ Behavior search completed") - return result, nil -} - -// getBehaviors retrieves behavior details by IDs. -func (t *internalTool) getBehaviors(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { - return nil, goerr.New("ids is required") - } - - msg.Trace(ctx, "📋 Retrieving behavior details (ids: `%s`)", ids) - body := map[string]any{ - "ids": splitAndTrim(ids), - } - result, err := t.doRequest(ctx, http.MethodPost, "/incidents/entities/behaviors/GET/v1", body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to retrieve behaviors (ids: `%s`): %v", ids, err) - return nil, err - } - msg.Trace(ctx, "✅ Retrieved behavior details") - return result, nil -} - -// searchDevices searches for device (host) IDs using FQL filters. -func (t *internalTool) searchDevices(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) - msg.Trace(ctx, "🔍 Searching devices (filter: `%s`)", filter) - - path := "/devices/queries/devices-scroll/v1" - params := buildQueryParams(args, "filter", "sort", "limit", "offset") - if params != "" { - path += "?" + params - } - result, err := t.doRequest(ctx, http.MethodGet, path, nil) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Device search failed (filter: `%s`): %v", filter, err) - return nil, err - } - msg.Trace(ctx, "✅ Device search completed") - return result, nil -} - -// getDevices retrieves device (host) details by IDs. -func (t *internalTool) getDevices(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { - return nil, goerr.New("ids is required") - } - - msg.Trace(ctx, "📋 Retrieving device details (ids: `%s`)", ids) - body := map[string]any{ - "ids": splitAndTrim(ids), - } - result, err := t.doRequest(ctx, http.MethodPost, "/devices/entities/devices/v2", body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to retrieve devices (ids: `%s`): %v", ids, err) - return nil, err - } - msg.Trace(ctx, "✅ Retrieved device details") - return result, nil -} - -// getCrowdScores retrieves CrowdScore values. -func (t *internalTool) getCrowdScores(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) - msg.Trace(ctx, "📊 Retrieving CrowdScores (filter: `%s`)", filter) - - path := "/incidents/combined/crowdscores/v1" - params := buildQueryParams(args, "filter") - if params != "" { - path += "?" + params - } - result, err := t.doRequest(ctx, http.MethodGet, path, nil) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* CrowdScore retrieval failed (filter: `%s`): %v", filter, err) - return nil, err - } - msg.Trace(ctx, "✅ CrowdScores retrieved") - return result, nil -} - -// searchEvents runs a CQL query via the Next-Gen SIEM Search API. -// It creates a query job and polls until the job completes, returning all events. -func (t *internalTool) searchEvents(ctx context.Context, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - - queryString, ok := args["query_string"].(string) - if !ok || queryString == "" { - return nil, goerr.New("query_string is required") - } - - repository := "search-all" - if repo, ok := args["repository"].(string); ok && repo != "" { - repository = repo - } - - body := map[string]any{ - "queryString": queryString, - } - if start, ok := args["start"].(string); ok && start != "" { - body["start"] = start - } else { - body["start"] = "1d" - } - if end, ok := args["end"].(string); ok && end != "" { - body["end"] = end - } else { - body["end"] = "now" - } - - // Step 1: Create query job - msg.Trace(ctx, "🔍 Searching events (query: `%s`, repo: `%s`)", queryString, repository) - jobPath := fmt.Sprintf("/humio/api/v1/repositories/%s/queryjobs", repository) - jobResp, err := t.doRequest(ctx, http.MethodPost, jobPath, body) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to create event search job (query: `%s`): %v", queryString, err) - return nil, goerr.Wrap(err, "failed to create event search query job", - goerr.V("repository", repository), - goerr.V("query", queryString), - ) - } - - jobID, ok := jobResp["id"].(string) - if !ok || jobID == "" { - msg.Warn(ctx, "⚠️ *[Falcon]* No job ID returned from event search (query: `%s`)", queryString) - return nil, goerr.New("no job ID returned from query job creation", - goerr.V("response", jobResp), - ) - } - - log.Debug("Event search query job created", "job_id", jobID, "repository", repository) - msg.Trace(ctx, "⏳ Event search job created (job_id: `%s`), polling for results...", jobID) - - // Step 2: Poll for results until done - resultPath := fmt.Sprintf("/humio/api/v1/repositories/%s/queryjobs/%s", repository, jobID) - const ( - maxPolls = 60 - pollInterval = 2 * time.Second - ) - - var allEvents []any - - for i := range maxPolls { - if i > 0 { - select { - case <-ctx.Done(): - return nil, goerr.Wrap(ctx.Err(), "context canceled while polling event search") - case <-time.After(pollInterval): - } - } - - pollResp, err := t.doRequest(ctx, http.MethodGet, resultPath, nil) - if err != nil { - msg.Warn(ctx, "⚠️ *[Falcon]* Failed to poll event search results (job: `%s`, attempt %d): %v", jobID, i+1, err) - return nil, goerr.Wrap(err, "failed to poll event search results", - goerr.V("job_id", jobID), - goerr.V("poll_attempt", i+1), - ) - } - - // Collect events from this poll - if events, ok := pollResp["events"].([]any); ok { - allEvents = append(allEvents, events...) - } - - // Check if done - if done, ok := pollResp["done"].(bool); ok && done { - log.Debug("Event search completed", - "job_id", jobID, - "total_events", len(allEvents), - "polls", i+1, - ) - msg.Trace(ctx, "✅ Event search completed: %d events retrieved", len(allEvents)) - - result := map[string]any{ - "done": true, - "events": allEvents, - "repository": repository, - } - - // Include metadata if available - if meta, ok := pollResp["metadataResult"]; ok { - result["metadata"] = meta - } - - return result, nil - } - - log.Debug("Event search still running, polling...", - "job_id", jobID, - "poll_attempt", i+1, - "events_so_far", len(allEvents), - ) - } - - // Return partial results if max polls reached - log.Warn("Event search reached max poll limit, returning partial results", - "job_id", jobID, - "total_events", len(allEvents), - ) - msg.Trace(ctx, "⚠️ Event search reached poll limit, returning %d partial results", len(allEvents)) - - return map[string]any{ - "done": false, - "events": allEvents, - "repository": repository, - "warning": "Search did not complete within the polling limit. Partial results returned.", - }, nil -} - -// buildQueryParams constructs URL query parameters from tool arguments. -func buildQueryParams(args map[string]any, keys ...string) string { - params := url.Values{} - for _, key := range keys { - if val, ok := args[key]; ok { - switch v := val.(type) { - case string: - if v != "" { - params.Set(key, v) - } - case float64: - params.Set(key, fmt.Sprintf("%d", int(v))) - } - } - } - return params.Encode() -} - -// splitAndTrim splits a comma-separated string and trims whitespace from each element. -func splitAndTrim(s string) []string { - parts := strings.Split(s, ",") - result := make([]string, 0, len(parts)) - for _, p := range parts { - trimmed := strings.TrimSpace(p) - if trimmed != "" { - result = append(result, trimmed) - } - } - return result -} diff --git a/pkg/agents/falcon/tool_test.go b/pkg/agents/falcon/tool_test.go deleted file mode 100644 index 7db1c467..00000000 --- a/pkg/agents/falcon/tool_test.go +++ /dev/null @@ -1,674 +0,0 @@ -package falcon_test - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "net/http/httptest" - "os" - "strings" - "testing" - - "github.com/m-mizutani/gt" - "github.com/secmon-lab/warren/pkg/agents/falcon" -) - -// newTestServer creates a test server that handles both OAuth2 token and API requests. -// The apiHandler is called for all non-token requests. -func newTestServer(t *testing.T, apiHandler http.HandlerFunc) *httptest.Server { - t.Helper() - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Handle OAuth2 token request - if r.URL.Path == "/oauth2/token" { - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "access_token": "test-token", - "expires_in": 1800, - "token_type": "bearer", - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - return - } - - // Verify authorization header - auth := r.Header.Get("Authorization") - gt.Equal(t, auth, "Bearer test-token") - - apiHandler(w, r) - })) -} - -func TestInternalTool_SpecCount(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - count, err := tool.SpecCount(context.Background()) - gt.NoError(t, err) - gt.Equal(t, count, 10) -} - -func TestInternalTool_SearchIncidents(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodGet) - gt.True(t, r.URL.Path == "/incidents/queries/incidents/v1") - gt.Equal(t, r.URL.Query().Get("filter"), "status:'30'") - gt.Equal(t, r.URL.Query().Get("limit"), "10") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []string{"inc:abc:123", "inc:abc:456"}, - "meta": map[string]any{ - "pagination": map[string]any{ - "total": 2, - }, - }, - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_incidents", map[string]any{ - "filter": "status:'30'", - "limit": float64(10), - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 2) - gt.Equal(t, resources[0], "inc:abc:123") - gt.Equal(t, resources[1], "inc:abc:456") -} - -func TestInternalTool_GetIncidents(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/incidents/entities/incidents/GET/v1") - gt.Equal(t, r.Header.Get("Content-Type"), "application/json") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - - ids, ok := body["ids"].([]any) - gt.True(t, ok) - gt.Equal(t, len(ids), 2) - gt.Equal(t, ids[0], "inc:abc:123") - gt.Equal(t, ids[1], "inc:abc:456") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - {"incident_id": "inc:abc:123", "status": 30}, - {"incident_id": "inc:abc:456", "status": 20}, - }, - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_get_incidents", map[string]any{ - "ids": "inc:abc:123, inc:abc:456", - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 2) -} - -func TestInternalTool_SearchAlerts(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/alerts/combined/alerts/v1") - gt.Equal(t, r.Header.Get("Content-Type"), "application/json") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - - gt.Equal(t, body["filter"], "severity:>50") - limit, ok := body["limit"].(float64) - gt.True(t, ok) - gt.Equal(t, limit, float64(100)) - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - {"composite_id": "alert:1", "severity": 80}, - }, - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_alerts", map[string]any{ - "filter": "severity:>50", - "limit": float64(100), - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) -} - -func TestInternalTool_GetAlerts(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/alerts/entities/alerts/v2") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - - compositeIDs, ok := body["composite_ids"].([]any) - gt.True(t, ok) - gt.Equal(t, len(compositeIDs), 1) - gt.Equal(t, compositeIDs[0], "alert:1") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - {"composite_id": "alert:1", "severity": 80, "status": "new"}, - }, - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_get_alerts", map[string]any{ - "composite_ids": "alert:1", - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) -} - -func TestInternalTool_SearchBehaviors(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodGet) - gt.True(t, r.URL.Path == "/incidents/queries/behaviors/v1") - gt.Equal(t, r.URL.Query().Get("filter"), "tactic:'Lateral Movement'") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []string{"beh:abc:123"}, - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_behaviors", map[string]any{ - "filter": "tactic:'Lateral Movement'", - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) - gt.Equal(t, resources[0], "beh:abc:123") -} - -func TestInternalTool_GetBehaviors(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/incidents/entities/behaviors/GET/v1") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - - ids, ok := body["ids"].([]any) - gt.True(t, ok) - gt.Equal(t, len(ids), 1) - gt.Equal(t, ids[0], "beh:abc:123") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - {"behavior_id": "beh:abc:123", "tactic": "Lateral Movement"}, - }, - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_get_behaviors", map[string]any{ - "ids": "beh:abc:123", - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) -} - -func TestInternalTool_GetCrowdScores(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodGet) - gt.True(t, r.URL.Path == "/incidents/combined/crowdscores/v1") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - {"id": "score-1", "score": 42, "adjusted_score": 45}, - }, - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_get_crowdscores", map[string]any{}) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) -} - -func TestInternalTool_UnknownTool(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_unknown", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_GetIncidents_MissingIDs(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_get_incidents", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_GetAlerts_MissingCompositeIDs(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_get_alerts", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_GetBehaviors_MissingIDs(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_get_behaviors", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_TokenRetryOn401(t *testing.T) { - var callCount int - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/oauth2/token" { - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "access_token": "new-token", - "expires_in": 1800, - "token_type": "bearer", - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - return - } - - callCount++ - if callCount == 1 { - // First API call returns 401 - w.WriteHeader(http.StatusUnauthorized) - _, err := w.Write([]byte(`{"errors": [{"message": "access denied"}]}`)) - gt.NoError(t, err) - return - } - - // Second API call succeeds - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []string{"inc:retry:123"}, - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - })) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_incidents", map[string]any{}) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 1) - gt.Equal(t, resources[0], "inc:retry:123") - gt.Equal(t, callCount, 2) -} - -func TestInternalTool_SearchEvents_Immediate(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/queryjobs") { - // Create query job - gt.Equal(t, r.Header.Get("Content-Type"), "application/json") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - gt.Equal(t, body["queryString"], "aid=test123") - gt.Equal(t, body["start"], "1d") - gt.Equal(t, body["end"], "now") - - w.Header().Set("Content-Type", "application/json") - err = json.NewEncoder(w).Encode(map[string]any{ - "id": "job-123", - "hashedQueryOnView": "abc", - }) - gt.NoError(t, err) - return - } - - if r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/queryjobs/job-123") { - // Get results — immediately done - w.Header().Set("Content-Type", "application/json") - err := json.NewEncoder(w).Encode(map[string]any{ - "done": true, - "cancelled": false, - "events": []map[string]any{ - { - "timestamp": "1736264422005", - "#event_simpleName": "ProcessRollup2", - "aid": "test123", - "ComputerName": "workstation1", - "FileName": "cmd.exe", - }, - }, - }) - gt.NoError(t, err) - return - } - - t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_events", map[string]any{ - "query_string": "aid=test123", - }) - gt.NoError(t, err) - - done, ok := result["done"].(bool) - gt.True(t, ok) - gt.True(t, done) - - events, ok := result["events"].([]any) - gt.True(t, ok) - gt.Equal(t, len(events), 1) -} - -func TestInternalTool_SearchEvents_WithPolling(t *testing.T) { - var pollCount int - - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/queryjobs") { - w.Header().Set("Content-Type", "application/json") - err := json.NewEncoder(w).Encode(map[string]any{ - "id": "job-poll", - }) - gt.NoError(t, err) - return - } - - if r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/queryjobs/job-poll") { - pollCount++ - w.Header().Set("Content-Type", "application/json") - - if pollCount < 3 { - // Not done yet, return partial results - err := json.NewEncoder(w).Encode(map[string]any{ - "done": false, - "cancelled": false, - "events": []map[string]any{ - {"aid": "test", "event": fmt.Sprintf("event-%d", pollCount)}, - }, - }) - gt.NoError(t, err) - return - } - - // Done on third poll - err := json.NewEncoder(w).Encode(map[string]any{ - "done": true, - "cancelled": false, - "events": []map[string]any{ - {"aid": "test", "event": "event-final"}, - }, - }) - gt.NoError(t, err) - return - } - - t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_events", map[string]any{ - "query_string": "aid=test", - "repository": "investigate_view", - "start": "7d", - "end": "now", - }) - gt.NoError(t, err) - - done, ok := result["done"].(bool) - gt.True(t, ok) - gt.True(t, done) - - // Should have accumulated events from all polls - events, ok := result["events"].([]any) - gt.True(t, ok) - gt.Equal(t, len(events), 3) // 1 from poll 1 + 1 from poll 2 + 1 from final - gt.Equal(t, pollCount, 3) - - repo, ok := result["repository"].(string) - gt.True(t, ok) - gt.Equal(t, repo, "investigate_view") -} - -func TestInternalTool_SearchEvents_MissingQueryString(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_search_events", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_SearchDevices(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodGet) - gt.True(t, r.URL.Path == "/devices/queries/devices-scroll/v1") - gt.Equal(t, r.URL.Query().Get("filter"), "platform_name:'Windows'") - gt.Equal(t, r.URL.Query().Get("limit"), "5") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []string{"device-abc123", "device-def456"}, - "meta": map[string]any{ - "pagination": map[string]any{ - "total": 2, - "offset": "", - }, - }, - } - err := json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_search_devices", map[string]any{ - "filter": "platform_name:'Windows'", - "limit": float64(5), - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 2) - gt.Equal(t, resources[0], "device-abc123") - gt.Equal(t, resources[1], "device-def456") -} - -func TestInternalTool_GetDevices(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - gt.Equal(t, r.Method, http.MethodPost) - gt.Equal(t, r.URL.Path, "/devices/entities/devices/v2") - gt.Equal(t, r.Header.Get("Content-Type"), "application/json") - - var body map[string]any - err := json.NewDecoder(r.Body).Decode(&body) - gt.NoError(t, err) - - ids, ok := body["ids"].([]any) - gt.True(t, ok) - gt.Equal(t, len(ids), 2) - gt.Equal(t, ids[0], "device-abc123") - gt.Equal(t, ids[1], "device-def456") - - w.Header().Set("Content-Type", "application/json") - resp := map[string]any{ - "resources": []map[string]any{ - { - "device_id": "device-abc123", - "hostname": "web-server-01", - "platform_name": "Windows", - "os_version": "Windows 11", - "external_ip": "203.0.113.10", - "local_ip": "10.0.0.5", - "status": "normal", - "agent_version": "7.10.18207.0", - }, - { - "device_id": "device-def456", - "hostname": "db-server-02", - "platform_name": "Linux", - "os_version": "RHEL 9.2", - "external_ip": "203.0.113.11", - "local_ip": "10.0.0.6", - "status": "normal", - "agent_version": "7.10.18207.0", - }, - }, - } - err = json.NewEncoder(w).Encode(resp) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - result, err := tool.Run(context.Background(), "falcon_get_devices", map[string]any{ - "ids": "device-abc123, device-def456", - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(resources), 2) -} - -func TestInternalTool_GetDevices_MissingIDs(t *testing.T) { - tool := falcon.NewInternalToolForTest("id", "secret", "http://localhost") - _, err := tool.Run(context.Background(), "falcon_get_devices", map[string]any{}) - gt.Error(t, err) -} - -func TestInternalTool_APIError(t *testing.T) { - srv := newTestServer(t, func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusTooManyRequests) - _, err := w.Write([]byte(`{"errors": [{"message": "rate limit exceeded"}]}`)) - gt.NoError(t, err) - }) - defer srv.Close() - - tool := falcon.NewInternalToolForTest("id", "secret", srv.URL) - _, err := tool.Run(context.Background(), "falcon_search_incidents", map[string]any{}) - gt.Error(t, err) -} - -func newE2EInternalTool(t *testing.T) *falcon.InternalToolForTest { - t.Helper() - - clientID := os.Getenv("TEST_AGENT_FALCON_CLIENT_ID") - clientSecret := os.Getenv("TEST_AGENT_FALCON_CLIENT_SECRET") - if clientID == "" || clientSecret == "" { - t.Skip("TEST_AGENT_FALCON_CLIENT_ID and TEST_AGENT_FALCON_CLIENT_SECRET are not set") - } - - baseURL := os.Getenv("TEST_AGENT_FALCON_BASE_URL") - if baseURL == "" { - baseURL = "https://api.crowdstrike.com" - } - - return falcon.NewInternalToolForTest(clientID, clientSecret, baseURL) -} - -func TestInternalTool_E2E_SearchDevices(t *testing.T) { - tool := newE2EInternalTool(t) - - result, err := tool.Run(context.Background(), "falcon_search_devices", map[string]any{ - "limit": float64(5), - }) - gt.NoError(t, err) - - resources, ok := result["resources"].([]any) - gt.True(t, ok) - gt.True(t, len(resources) > 0) -} - -func TestInternalTool_E2E_GetDevices(t *testing.T) { - tool := newE2EInternalTool(t) - - // First, search for device IDs - searchResult, err := tool.Run(context.Background(), "falcon_search_devices", map[string]any{ - "limit": float64(2), - }) - gt.NoError(t, err) - - resources, ok := searchResult["resources"].([]any) - gt.True(t, ok) - gt.True(t, len(resources) > 0) - - // Build comma-separated IDs from search results - var ids []string - for _, r := range resources { - if id, ok := r.(string); ok { - ids = append(ids, id) - } - } - gt.True(t, len(ids) > 0) - - // Get device details - getResult, err := tool.Run(context.Background(), "falcon_get_devices", map[string]any{ - "ids": strings.Join(ids, ","), - }) - gt.NoError(t, err) - - devices, ok := getResult["resources"].([]any) - gt.True(t, ok) - gt.Equal(t, len(devices), len(ids)) - - // Verify first device has expected fields - device, ok := devices[0].(map[string]any) - gt.True(t, ok) - _, hasHostname := device["hostname"] - gt.True(t, hasHostname) - _, hasDeviceID := device["device_id"] - gt.True(t, hasDeviceID) -} diff --git a/pkg/agents/slack/README.md b/pkg/agents/slack/README.md deleted file mode 100644 index 211d459f..00000000 --- a/pkg/agents/slack/README.md +++ /dev/null @@ -1,31 +0,0 @@ -# Slack Search Agent - -Sub-agent for searching Slack messages to provide investigative context during alert analysis. Useful for finding related discussions, previous incidents, and contextual information from team communications. - -## Configuration - -| Environment Variable | CLI Flag | Required | Description | -|---|---|---|---| -| `WARREN_AGENT_SLACK_USER_TOKEN` | `--agent-slack-user-token` | Yes | Slack **User** token with `search:read` scope | - -> **Note:** The Slack search API requires a **User token** (`xoxp-...`), not a Bot token (`xoxb-...`). - -## Capabilities - -- Search Slack messages with full query syntax -- Sort results by relevance score or timestamp -- Paginate through large result sets - -## Setup - -1. In your Slack workspace, create or use an existing Slack App -2. Navigate to **OAuth & Permissions** -3. Under **User Token Scopes**, add `search:read` -4. Install the app and copy the **User OAuth Token** (`xoxp-...`) -5. Set `WARREN_AGENT_SLACK_USER_TOKEN` environment variable - -## Agent Memory - -The Slack agent uses agent memory (ID: `slack_search`) to learn effective search patterns and remember useful query strategies from previous investigations. - -The agent is automatically registered when the user token is configured. diff --git a/pkg/agents/slack/agent.go b/pkg/agents/slack/agent.go deleted file mode 100644 index 9738ce62..00000000 --- a/pkg/agents/slack/agent.go +++ /dev/null @@ -1,32 +0,0 @@ -package slack - -import ( - "context" - - "github.com/gollem-dev/gollem" -) - -// toolSet implements interfaces.ToolSet by wrapping the internalTool. -type toolSet struct { - tool *internalTool -} - -func (ts *toolSet) ID() string { - return "slack_agent" -} - -func (ts *toolSet) Description() string { - return "Search for messages in Slack workspace, retrieve thread replies, and get context around specific messages." -} - -func (ts *toolSet) Prompt(_ context.Context) (string, error) { - return buildSystemPrompt() -} - -func (ts *toolSet) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return ts.tool.Specs(ctx) -} - -func (ts *toolSet) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - return ts.tool.Run(ctx, name, args) -} diff --git a/pkg/agents/slack/agent_test.go b/pkg/agents/slack/agent_test.go deleted file mode 100644 index beaa9f89..00000000 --- a/pkg/agents/slack/agent_test.go +++ /dev/null @@ -1,47 +0,0 @@ -package slack_test - -import ( - "context" - "testing" - - "github.com/m-mizutani/gt" - - slackagent "github.com/secmon-lab/warren/pkg/agents/slack" - domainmock "github.com/secmon-lab/warren/pkg/domain/mock" -) - -func TestToolSet_ID(t *testing.T) { - slackClient := &domainmock.SlackClientMock{} - ts := slackagent.NewToolSetForTest(slackClient) - - gt.V(t, ts.ID()).Equal("slack_agent") -} - -func TestToolSet_Description(t *testing.T) { - slackClient := &domainmock.SlackClientMock{} - ts := slackagent.NewToolSetForTest(slackClient) - - description := ts.Description() - gt.V(t, description).NotEqual("") - gt.True(t, len(description) > 0) -} - -func TestToolSet_Prompt(t *testing.T) { - slackClient := &domainmock.SlackClientMock{} - ts := slackagent.NewToolSetForTest(slackClient) - - ctx := context.Background() - prompt, err := ts.Prompt(ctx) - gt.NoError(t, err) - gt.V(t, prompt).NotEqual("") -} - -func TestToolSet_Specs(t *testing.T) { - slackClient := &domainmock.SlackClientMock{} - ts := slackagent.NewToolSetForTest(slackClient) - - ctx := context.Background() - specs, err := ts.Specs(ctx) - gt.NoError(t, err) - gt.N(t, len(specs)).Equal(3) // slack_search_messages, slack_get_thread_messages, slack_get_context_messages -} diff --git a/pkg/agents/slack/export_test.go b/pkg/agents/slack/export_test.go deleted file mode 100644 index 4408e1b9..00000000 --- a/pkg/agents/slack/export_test.go +++ /dev/null @@ -1,35 +0,0 @@ -package slack - -import ( - "github.com/gollem-dev/gollem" - "github.com/secmon-lab/warren/pkg/domain/interfaces" -) - -// Export for testing - -type InternalTool = internalTool - -// NewInternalToolForTest creates an internalTool for testing -func NewInternalToolForTest(slackClient interfaces.SlackClient, maxLimit int) *internalTool { - return &internalTool{ - slackClient: slackClient, - maxLimit: maxLimit, - } -} - -// NewToolSetForTest creates a toolSet instance for testing -func NewToolSetForTest(slackClient interfaces.SlackClient) *toolSet { - return &toolSet{ - tool: &internalTool{slackClient: slackClient}, - } -} - -// ExportedBuildSystemPrompt is exported for testing -func ExportedBuildSystemPrompt() (string, error) { - return buildSystemPrompt() -} - -// ExportedNewPromptTemplate is exported for testing -func ExportedNewPromptTemplate() (*gollem.PromptTemplate, error) { - return newPromptTemplate() -} diff --git a/pkg/agents/slack/extract.go b/pkg/agents/slack/extract.go deleted file mode 100644 index 578351f3..00000000 --- a/pkg/agents/slack/extract.go +++ /dev/null @@ -1 +0,0 @@ -package slack diff --git a/pkg/agents/slack/factory.go b/pkg/agents/slack/factory.go deleted file mode 100644 index df756e23..00000000 --- a/pkg/agents/slack/factory.go +++ /dev/null @@ -1,46 +0,0 @@ -package slack - -import ( - "context" - - "github.com/secmon-lab/warren/pkg/domain/interfaces" - "github.com/secmon-lab/warren/pkg/utils/logging" - slackSDK "github.com/slack-go/slack" - "github.com/urfave/cli/v3" -) - -// Factory implements agents.ToolSetFactory interface. -type Factory struct { - oauthToken string -} - -// Flags implements agents.ToolSetFactory -func (f *Factory) Flags() []cli.Flag { - return []cli.Flag{ - &cli.StringFlag{ - Name: "agent-slack-user-token", - Usage: "Slack User OAuth Token for message search (requires search:read scope)", - Destination: &f.oauthToken, - Category: "Agent:Slack", - Sources: cli.EnvVars("WARREN_AGENT_SLACK_USER_TOKEN"), - }, - } -} - -// Configure implements agents.ToolSetFactory. -// Returns (nil, nil) if the OAuth token is not set. -func (f *Factory) Configure(ctx context.Context) (interfaces.ToolSet, error) { - if f.oauthToken == "" { - return nil, nil - } - - slackClient := slackSDK.New(f.oauthToken) - - ts := &toolSet{ - tool: &internalTool{slackClient: slackClient}, - } - - logging.From(ctx).Info("Slack Search Agent configured") - - return ts, nil -} diff --git a/pkg/agents/slack/prompt.go b/pkg/agents/slack/prompt.go deleted file mode 100644 index 0f1e26fb..00000000 --- a/pkg/agents/slack/prompt.go +++ /dev/null @@ -1,33 +0,0 @@ -package slack - -import ( - _ "embed" - - "github.com/gollem-dev/gollem" -) - -//go:embed prompt/system.md -var systemPromptTemplate string - -// buildSystemPrompt builds system prompt (for factory) -func buildSystemPrompt() (string, error) { - return systemPromptTemplate, nil -} - -// newPromptTemplate creates a PromptTemplate for the SubAgent -func newPromptTemplate() (*gollem.PromptTemplate, error) { - return gollem.NewPromptTemplate( - "{{if ._slack_context}}## Current Slack Context\n{{._slack_context}}\nYou are being invoked from within this Slack context. When the user refers to \"this channel\" or \"here\", they mean the channel above. You can use the channel ID to scope your searches, for example: `in:C12345678`.\n\n{{end}}{{.request}}", - map[string]*gollem.Parameter{ - "request": { - Type: gollem.TypeString, - Description: "DO NOT specify search keywords or terms. Describe ONLY the concept/situation in natural language. The agent will determine all search keywords and variations. ✗ BAD: 'search for authentication keyword', 'messages containing auth error', 'find keyword login' ✓ GOOD: 'people having authentication problems', 'discussions about performance issues', 'error reports in #security-alerts channel'. Include: (1) What concept/situation to find (NOT keywords), (2) Time period if relevant, (3) Channel/user scope if relevant. The Slack agent handles all keyword selection, variations, and multilingual terms automatically.", - Required: true, - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of messages to return in the response (default: 50, max: 200). Use this to control response size.", - }, - }, - ) -} diff --git a/pkg/agents/slack/prompt/extract.md b/pkg/agents/slack/prompt/extract.md deleted file mode 100644 index c92ec7a0..00000000 --- a/pkg/agents/slack/prompt/extract.md +++ /dev/null @@ -1,41 +0,0 @@ -# Task -Extract the raw Slack messages from the conversation history and return them as a JSON array. - -# Key Principle: Understand User Intent -- Do NOT just extract messages literally based on the user's wording -- UNDERSTAND what the user actually wants to find or learn -- INTERPRET the user's intent from their request and the context -- Select messages that FULFILL the user's actual information needs, not just match keywords - -# Guidelines -1. First, understand the user's true intent from their original request -2. Look for tool execution results in the conversation history (e.g., `slack_search_messages` tool responses) -3. Extract the message data that answers the user's actual question or need -4. Do NOT summarize or paraphrase message content - include full original text -5. For each message, include: text, user, channel, timestamp -6. If multiple searches were performed, intelligently select messages that fulfill the user's intent -7. Return ONLY the messages array - no wrapper object, no additional fields - -# Examples of Intent Understanding -- User asks "people having authentication problems" → Extract messages about auth issues, not just containing "authentication" -- User asks "what did X say about Y" → Extract X's messages related to topic Y -- User asks "recent discussions on Z" → Extract conversational messages about Z, not just mentions - -# Output Format -Return a JSON array directly: -[ - { - "text": "full message text here", - "user": "user_id or name", - "channel": "channel_id or name", - "timestamp": "timestamp" - }, - ... -] - -# Important -- The output must be a JSON ARRAY at the top level -- Each element must have: text, user, channel, timestamp fields -- Preserve original message text exactly as it appears -- If no messages match the user's intent, return an empty array: [] -- Focus on user's INTENT, not literal wording diff --git a/pkg/agents/slack/prompt/system.md b/pkg/agents/slack/prompt/system.md deleted file mode 100644 index b47c5dbc..00000000 --- a/pkg/agents/slack/prompt/system.md +++ /dev/null @@ -1,154 +0,0 @@ -# Slack Message Search Agent - -You are a Slack search agent. **Understand the main agent's request, search comprehensively, and return raw message data organized to fulfill their specific need.** - -## Core Mission - -1. **Understand WHY** - What question are they trying to answer? (Who/When/What/Did...) -2. **Search smart** - Use related keywords ONLY for abstract concepts (not for specific names/IDs/unique terms) -3. **Return raw data** - Show actual message content (user, channel, time, text) organized to answer the question -4. **Verify fulfillment** - Does your response actually answer what was asked? - -## Available Tools - -- `slack_search_messages` - Search messages (supports Slack search syntax) -- `slack_get_thread_messages` - Get thread replies -- `slack_get_context_messages` - Get surrounding messages - -## Slack Search Syntax - -- `from:@user` - Messages from user -- `in:#channel` - Messages in channel -- `after:YYYY-MM-DD` / `before:YYYY-MM-DD` - Date filters -- Combine: `from:@user in:#channel error after:2024-01-01` - -## Response Limit - -Maximum messages per search: **{{ .limit }}** - -## Past Learnings - -{{ if .memories }} -You have access to insights learned from past executions. Each insight is self-contained domain knowledge: - -{{ range .memories }} -- {{ .Claim }} -{{ end }} -{{ else }} -No past learnings yet. -{{ end }} - -## Instructions - -### 1. Understand the Request - -- **Extract the CONCEPT**: If the request contains specific keywords, identify the underlying concept (e.g., authentication problems, database issues) -- **What question needs answering?** (e.g., "Who discussed X?", "When did Y happen?", "Are there people with Z problems?") -- **What data would help?** (users? timestamps? message content?) -- **Time constraints?** (last week, more than a month ago, etc.) -- **Specific scope?** (channel, user, etc.) - -### 2. Plan Search Strategy - -**IMPORTANT**: Even if the request contains specific keywords, identify if they represent an abstract concept and plan variations. - -**For ABSTRACT concepts** (e.g., "authentication issues", "performance problems"): -- Identify the core concept -- Think of related terms in multiple languages if needed -- Search with multiple keyword variations to ensure comprehensive results - -**For SPECIFIC terms** (e.g., "sqldef", user IDs, ticket numbers, unique names): -- Search directly - do NOT search for synonyms or related words -- These are already unique and specific - -### 3. Execute and Return - -- Search using Slack syntax -- **If search returns no results**: Try a few related keywords or alternative terms (even for specific terms, if the initial search was empty) -- Use threads/context tools if needed -- **Format response to answer the question**: - - Include raw data: user_name, channel_name, timestamp, formatted_time, text - - Organize by what was asked (group by user if "who?", by time if "when?", etc.) - - You may add brief context, but **raw data is primary** - - **Verify your answer addresses the original question** - -### 4. What NOT to Do - -- Don't search synonyms for specific/unique terms (names, IDs, specific tool names, etc.) UNLESS the initial search returned no results -- Don't dump messages without organization -- Don't make vague summaries without showing actual data -- Don't omit key info (users, times, message content) - -## Examples - -**Request**: "I need to find who discussed the sqldef tool more than a month ago" -- **Understand**: WHO discussed SPECIFIC TOOL WHEN? -- **Search**: - 1. `sqldef before:2024-10-23` (sqldef is a specific tool name) - 2. If no results, try: `schema definition before:2024-10-23`, `DDL before:2024-10-23` -- **Response format**: -``` -Found messages about sqldef from more than a month ago: - -User: john_doe, Channel: #engineering -Time: 2024-10-15T14:23:45+09:00 -Message: "We should consider using sqldef for schema management..." - -User: jane_smith, Channel: #database -Time: 2024-10-10T09:15:30+09:00 -Message: "I've been testing sqldef..." -``` - -**Request**: "Find authentication problems in #security from last week" -- **Understand**: WHAT PROBLEMS occurred WHEN in CHANNEL? -- **Search**: Authentication is abstract - search multiple terms: - - `in:#security authentication after:2024-11-16` - - `in:#security login after:2024-11-16` - - `in:#security auth after:2024-11-16` -- **Response format**: -``` -Authentication-related issues in #security from last week: - -[2024-11-20 10:45] user: admin -"Login timeout errors reported by 15 users..." - -[2024-11-21 15:30] user: security-bot -"Authentication service degraded..." -``` - -**Request**: "Did @john mention ticket-12345 recently?" -- **Understand**: DID user mention SPECIFIC TICKET? -- **Search**: `from:@john ticket-12345` (ticket-12345 is specific - no synonyms) -- **Response format**: -``` -Messages from @john about ticket-12345: - -[Found 2 messages] -Time: 2024-11-22T16:20:00+09:00, Channel: #incidents -Message: "Updated ticket-12345 with root cause analysis..." - -Time: 2024-11-21T09:15:00+09:00, Channel: #incidents -Message: "Working on ticket-12345 reproduction steps..." -``` - -**Request**: "search for authentication keyword" (contains keyword specification) -- **Understand**: This looks like specific keywords BUT represents ABSTRACT CONCEPT (authentication problems) -- **Extract concept**: People having authentication/login troubles -- **Search with variations**: - - `authentication` - - `login` - - `auth` - - `access` - - Combine with problem indicators: `error`, `issue`, `problem`, `trouble` -- **Response format**: -``` -Found people discussing authentication problems: - -User: alice, Channel: #support -Time: 2024-11-22T15:30:00+09:00 -Message: "Can't login to the system..." - -User: bob, Channel: #help -Time: 2024-11-21T10:15:00+09:00 -Message: "Authentication error keeps happening..." -``` diff --git a/pkg/agents/slack/prompt_test.go b/pkg/agents/slack/prompt_test.go deleted file mode 100644 index a87bc33e..00000000 --- a/pkg/agents/slack/prompt_test.go +++ /dev/null @@ -1,67 +0,0 @@ -package slack_test - -import ( - "testing" - - "github.com/m-mizutani/gt" - slackagent "github.com/secmon-lab/warren/pkg/agents/slack" -) - -func TestBuildSystemPrompt(t *testing.T) { - prompt, err := slackagent.ExportedBuildSystemPrompt() - gt.NoError(t, err) - gt.V(t, prompt).NotEqual("") - gt.True(t, len(prompt) > 0) -} - -func TestNewPromptTemplate(t *testing.T) { - template, err := slackagent.ExportedNewPromptTemplate() - gt.NoError(t, err) - gt.V(t, template).NotNil() - - // Check that template has expected parameters - params := template.Parameters() - gt.V(t, len(params)).NotEqual(0) - - // Check request parameter exists and is required - requestParam, hasRequest := params["request"] - gt.True(t, hasRequest) - gt.V(t, requestParam).NotNil() - gt.True(t, requestParam.Required) - gt.V(t, requestParam.Type).Equal("string") - - // Check limit parameter exists and is optional - limitParam, hasLimit := params["limit"] - gt.True(t, hasLimit) - gt.V(t, limitParam).NotNil() - gt.False(t, limitParam.Required) - gt.V(t, limitParam.Type).Equal("number") - - // Check that _memory_context is NOT in parameters (internal only) - _, hasMemoryContext := params["_memory_context"] - gt.False(t, hasMemoryContext) - - // Check that _slack_context is NOT in parameters (internal only) - _, hasSlackContext := params["_slack_context"] - gt.False(t, hasSlackContext) - - t.Run("render with slack context", func(t *testing.T) { - rendered, err := template.Render(map[string]any{ - "request": "find messages about security alerts", - "_slack_context": "Current Slack context: channel_id=C67890, thread_ts=123.456, team_id=T12345", - }) - gt.NoError(t, err) - gt.S(t, rendered).Contains("Current Slack Context") - gt.S(t, rendered).Contains("channel_id=C67890") - gt.S(t, rendered).Contains("find messages about security alerts") - }) - - t.Run("render without slack context", func(t *testing.T) { - rendered, err := template.Render(map[string]any{ - "request": "find messages about security alerts", - }) - gt.NoError(t, err) - gt.S(t, rendered).NotContains("Current Slack Context") - gt.S(t, rendered).Contains("find messages about security alerts") - }) -} diff --git a/pkg/agents/slack/tool.go b/pkg/agents/slack/tool.go deleted file mode 100644 index 1da9e4f4..00000000 --- a/pkg/agents/slack/tool.go +++ /dev/null @@ -1,344 +0,0 @@ -package slack - -import ( - "context" - "fmt" - "strconv" - "strings" - "time" - - "github.com/gollem-dev/gollem" - "github.com/m-mizutani/goerr/v2" - "github.com/secmon-lab/warren/pkg/domain/interfaces" - slackSDK "github.com/slack-go/slack" -) - -type internalTool struct { - slackClient interfaces.SlackClient - maxLimit int // Maximum number of results allowed from search -} - -// parseSlackTimestamp parses a Slack timestamp string (e.g., "1234567890.123456") -// and returns a time.Time with sub-second precision preserved. -func parseSlackTimestamp(tsStr string) time.Time { - parts := strings.Split(tsStr, ".") - if len(parts) == 0 { - return time.Time{} - } - - // Parse seconds - sec, err := strconv.ParseInt(parts[0], 10, 64) - if err != nil { - return time.Time{} - } - - // Parse nanoseconds from fractional part - var nsec int64 - if len(parts) > 1 { - // Pad or truncate to 9 digits (nanoseconds) - fracStr := parts[1] - if len(fracStr) > 9 { - fracStr = fracStr[:9] - } else { - // Pad with zeros to get 9 digits - fracStr = fracStr + strings.Repeat("0", 9-len(fracStr)) - } - nsec, err = strconv.ParseInt(fracStr, 10, 64) - if err != nil { - // If fractional part is invalid, just use seconds - nsec = 0 - } - } - - return time.Unix(sec, nsec) -} - -func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "slack_search_messages", - Description: "Search for messages in Slack workspace using the search.messages API", - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "The search query (e.g., 'from:@user', 'in:general', 'has:link')", - Required: true, - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort order: 'score' (relevance) or 'timestamp' (newest first)", - }, - "sort_dir": { - Type: gollem.TypeString, - Description: "Sort direction: 'asc' or 'desc'", - }, - "count": { - Type: gollem.TypeNumber, - Description: "Number of results to return (default: 20, max: 100)", - }, - "page": { - Type: gollem.TypeNumber, - Description: "Page number for pagination (default: 1)", - }, - "highlight": { - Type: gollem.TypeBoolean, - Description: "Enable highlighting of search terms in results", - }, - }, - }, - { - Name: "slack_get_thread_messages", - Description: "Get all messages in a thread", - Parameters: map[string]*gollem.Parameter{ - "channel": { - Type: gollem.TypeString, - Description: "Channel ID", - Required: true, - }, - "thread_ts": { - Type: gollem.TypeString, - Description: "Thread timestamp (ts of the parent message)", - Required: true, - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of messages to return (default: 50, max: 200)", - }, - }, - }, - { - Name: "slack_get_context_messages", - Description: "Get messages before and after a specific message timestamp", - Parameters: map[string]*gollem.Parameter{ - "channel": { - Type: gollem.TypeString, - Description: "Channel ID", - Required: true, - }, - "around_ts": { - Type: gollem.TypeString, - Description: "Timestamp of the message to get context around", - Required: true, - }, - "before": { - Type: gollem.TypeNumber, - Description: "Number of messages before the timestamp (default: 10)", - }, - "after": { - Type: gollem.TypeNumber, - Description: "Number of messages after the timestamp (default: 10)", - }, - }, - }, - }, nil -} - -func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case "slack_search_messages": - return t.searchMessages(ctx, args) - case "slack_get_thread_messages": - return t.getThreadMessages(ctx, args) - case "slack_get_context_messages": - return t.getContextMessages(ctx, args) - default: - return nil, goerr.New("unknown tool name", goerr.V("name", name)) - } -} - -func (t *internalTool) searchMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - query, ok := args["query"].(string) - if !ok || query == "" { - return nil, goerr.New("query is required") - } - - params := slackSDK.SearchParameters{ - Count: 20, - Page: 1, - } - - if sort, ok := args["sort"].(string); ok { - params.Sort = sort - } - if sortDir, ok := args["sort_dir"].(string); ok { - params.SortDirection = sortDir - } - if count, ok := args["count"].(float64); ok { - params.Count = int(count) - } - if page, ok := args["page"].(float64); ok { - params.Page = int(page) - } - if highlight, ok := args["highlight"].(bool); ok { - params.Highlight = highlight - } - - resp, err := t.slackClient.SearchMessagesContext(ctx, query, params) - if err != nil { - return nil, goerr.Wrap(err, "failed to search messages") - } - - messages := make([]any, 0, len(resp.Matches)) - for i, msg := range resp.Matches { - // Enforce maxLimit from parent agent - if t.maxLimit > 0 && i >= t.maxLimit { - break - } - - formattedTime := parseSlackTimestamp(msg.Timestamp) - - // Thread timestamp - if this message is in a thread, use msg.Timestamp as thread_ts - // Note: Slack's SearchMessage doesn't directly expose thread_ts, but for thread replies, - // the timestamp can be used with slack_get_thread_messages - threadTS := "" // Empty if not a thread message - // If the message has Previous context, it might be a thread reply - // For simplicity, we include timestamp which can be used for thread operations - if msg.Previous.Text != "" || msg.Previous2.Text != "" { - threadTS = msg.Timestamp - } - - item := map[string]any{ - "channel_id": msg.Channel.ID, - "channel_name": msg.Channel.Name, - "user_name": msg.Username, - "text": msg.Text, - "timestamp": msg.Timestamp, - "thread_ts": threadTS, - "formatted_time": formattedTime.Format(time.RFC3339), - } - - messages = append(messages, item) - } - - return map[string]any{ - "total": float64(resp.Total), - "messages": messages, - }, nil -} - -func (t *internalTool) getThreadMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - channel, ok := args["channel"].(string) - if !ok || channel == "" { - return nil, goerr.New("channel is required") - } - - threadTS, ok := args["thread_ts"].(string) - if !ok || threadTS == "" { - return nil, goerr.New("thread_ts is required") - } - - limit := 50 - if l, ok := args["limit"].(float64); ok { - limit = int(l) - } - if limit > 200 { - limit = 200 - } - - params := &slackSDK.GetConversationRepliesParameters{ - ChannelID: channel, - Timestamp: threadTS, - Limit: limit, - } - - msgs, _, _, err := t.slackClient.GetConversationRepliesContext(ctx, params) - if err != nil { - return nil, goerr.Wrap(err, "failed to get thread messages") - } - - messages := make([]any, 0, len(msgs)) - for _, msg := range msgs { - formattedTime := parseSlackTimestamp(msg.Timestamp) - - messages = append(messages, map[string]any{ - "user_name": msg.Username, - "text": msg.Text, - "timestamp": msg.Timestamp, - "formatted_time": formattedTime.Format(time.RFC3339), - }) - } - - return map[string]any{ - "messages": messages, - }, nil -} - -func (t *internalTool) getContextMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - channel, ok := args["channel"].(string) - if !ok || channel == "" { - return nil, goerr.New("channel is required") - } - - aroundTS, ok := args["around_ts"].(string) - if !ok || aroundTS == "" { - return nil, goerr.New("around_ts is required") - } - - before := 10 - if b, ok := args["before"].(float64); ok { - before = int(b) - } - - after := 10 - if a, ok := args["after"].(float64); ok { - after = int(a) - } - - // Parse the timestamp - ts, err := strconv.ParseFloat(aroundTS, 64) - if err != nil { - return nil, goerr.Wrap(err, "invalid timestamp format") - } - - // Get messages before the timestamp - var beforeMessages []slackSDK.Message - if before > 0 { - beforeParams := &slackSDK.GetConversationHistoryParameters{ - ChannelID: channel, - Latest: fmt.Sprintf("%.6f", ts), - Inclusive: false, - Limit: before, - } - beforeResp, err := t.slackClient.GetConversationHistoryContext(ctx, beforeParams) - if err != nil { - return nil, goerr.Wrap(err, "failed to get messages before timestamp") - } - beforeMessages = beforeResp.Messages - } - - // Get messages after the timestamp - var afterMessages []slackSDK.Message - if after > 0 { - afterParams := &slackSDK.GetConversationHistoryParameters{ - ChannelID: channel, - Oldest: fmt.Sprintf("%.6f", ts), - Inclusive: false, - Limit: after, - } - afterResp, err := t.slackClient.GetConversationHistoryContext(ctx, afterParams) - if err != nil { - return nil, goerr.Wrap(err, "failed to get messages after timestamp") - } - afterMessages = afterResp.Messages - } - - formatMessages := func(msgs []slackSDK.Message) []any { - result := make([]any, 0, len(msgs)) - for _, msg := range msgs { - formattedTime := parseSlackTimestamp(msg.Timestamp) - - result = append(result, map[string]any{ - "user_name": msg.Username, - "text": msg.Text, - "timestamp": msg.Timestamp, - "formatted_time": formattedTime.Format(time.RFC3339), - }) - } - return result - } - - return map[string]any{ - "before_messages": formatMessages(beforeMessages), - "after_messages": formatMessages(afterMessages), - }, nil -} diff --git a/pkg/agents/slack/tool_helper_test.go b/pkg/agents/slack/tool_helper_test.go deleted file mode 100644 index e008d3ba..00000000 --- a/pkg/agents/slack/tool_helper_test.go +++ /dev/null @@ -1,61 +0,0 @@ -package slack - -import ( - "testing" - "time" - - "github.com/m-mizutani/gt" -) - -func TestParseSlackTimestamp(t *testing.T) { - t.Run("parses timestamp with sub-second precision", func(t *testing.T) { - // Slack timestamp: 1234567890.123456 - ts := parseSlackTimestamp("1234567890.123456") - - // Expected: 2009-02-13 23:31:30.123456 UTC - expected := time.Unix(1234567890, 123456000) - - gt.V(t, ts.Unix()).Equal(expected.Unix()) - gt.V(t, ts.Nanosecond()).Equal(expected.Nanosecond()) - }) - - t.Run("parses timestamp without sub-second precision", func(t *testing.T) { - // Slack timestamp: 1234567890 - ts := parseSlackTimestamp("1234567890") - - expected := time.Unix(1234567890, 0) - - gt.V(t, ts.Unix()).Equal(expected.Unix()) - gt.V(t, ts.Nanosecond()).Equal(0) - }) - - t.Run("returns zero time for invalid timestamp", func(t *testing.T) { - ts := parseSlackTimestamp("invalid") - - gt.True(t, ts.IsZero()) - }) - - t.Run("returns zero time for empty timestamp", func(t *testing.T) { - ts := parseSlackTimestamp("") - - gt.True(t, ts.IsZero()) - }) - - t.Run("preserves microsecond precision", func(t *testing.T) { - // Test various microsecond values - testCases := []struct { - input string - expected int64 // nanoseconds - }{ - {"1234567890.000001", 1000}, // 1 microsecond - {"1234567890.000100", 100000}, // 100 microseconds - {"1234567890.100000", 100000000}, // 0.1 second - {"1234567890.999999", 999999000}, // 999999 microseconds - } - - for _, tc := range testCases { - ts := parseSlackTimestamp(tc.input) - gt.V(t, ts.Nanosecond()).Equal(int(tc.expected)) - } - }) -} diff --git a/pkg/agents/slack/tool_test.go b/pkg/agents/slack/tool_test.go deleted file mode 100644 index f015e629..00000000 --- a/pkg/agents/slack/tool_test.go +++ /dev/null @@ -1,276 +0,0 @@ -package slack_test - -import ( - "context" - "testing" - - "github.com/m-mizutani/gt" - slackSDK "github.com/slack-go/slack" - - slackagent "github.com/secmon-lab/warren/pkg/agents/slack" - domainmock "github.com/secmon-lab/warren/pkg/domain/mock" -) - -func TestInternalTool_SearchMessages(t *testing.T) { - ctx := context.Background() - - slackClient := &domainmock.SlackClientMock{ - SearchMessagesContextFunc: func(ctx context.Context, query string, params slackSDK.SearchParameters) (*slackSDK.SearchMessages, error) { - // Verify query parameter - gt.V(t, query).Equal("test search") - - return &slackSDK.SearchMessages{ - Total: 2, - Matches: []slackSDK.SearchMessage{ - { - Type: "message", - Timestamp: "1234567890.123456", - Text: "Test message 1", - Username: "testuser", - Channel: slackSDK.CtxChannel{ - ID: "C123456", - Name: "general", - }, - }, - { - Type: "message", - Timestamp: "1234567891.654321", - Text: "Test message 2", - Username: "testuser2", - Channel: slackSDK.CtxChannel{ - ID: "C123456", - Name: "general", - }, - }, - }, - }, nil - }, - } - - tool := slackagent.NewInternalToolForTest(slackClient, 0) - - // Call searchMessages directly - result, err := tool.Run(ctx, "slack_search_messages", map[string]any{ - "query": "test search", - }) - - gt.NoError(t, err) - gt.V(t, result).NotNil() - - // Verify response structure - total, ok := result["total"].(float64) - gt.True(t, ok) - gt.V(t, total).Equal(2.0) - - messages, ok := result["messages"].([]any) - gt.True(t, ok) - gt.V(t, len(messages)).Equal(2) - - // Verify first message - msg1, ok := messages[0].(map[string]any) - gt.True(t, ok) - gt.V(t, msg1["text"]).Equal("Test message 1") - gt.V(t, msg1["user_name"]).Equal("testuser") - gt.V(t, msg1["channel_id"]).Equal("C123456") - gt.V(t, msg1["channel_name"]).Equal("general") - gt.V(t, msg1["timestamp"]).Equal("1234567890.123456") - - // Verify formatted time has sub-second precision - formattedTime, ok := msg1["formatted_time"].(string) - gt.True(t, ok) - gt.V(t, formattedTime).NotEqual("") -} - -func TestInternalTool_GetThreadMessages(t *testing.T) { - ctx := context.Background() - - slackClient := &domainmock.SlackClientMock{ - GetConversationRepliesContextFunc: func(ctx context.Context, params *slackSDK.GetConversationRepliesParameters) ([]slackSDK.Message, bool, string, error) { - // Verify parameters - gt.V(t, params.ChannelID).Equal("C123456") - gt.V(t, params.Timestamp).Equal("1234567890.123456") - - return []slackSDK.Message{ - { - Msg: slackSDK.Msg{ - Timestamp: "1234567890.123456", - Text: "Parent message", - User: "U123", - Username: "user1", - }, - }, - { - Msg: slackSDK.Msg{ - Timestamp: "1234567891.654321", - ThreadTimestamp: "1234567890.123456", - Text: "Reply message", - User: "U456", - Username: "user2", - }, - }, - }, false, "", nil - }, - } - - tool := slackagent.NewInternalToolForTest(slackClient, 0) - - // Call getThreadMessages directly - result, err := tool.Run(ctx, "slack_get_thread_messages", map[string]any{ - "channel": "C123456", - "thread_ts": "1234567890.123456", - }) - - gt.NoError(t, err) - gt.V(t, result).NotNil() - - // Verify response structure - messages, ok := result["messages"].([]any) - gt.True(t, ok) - gt.V(t, len(messages)).Equal(2) - - // Verify parent message - msg1, ok := messages[0].(map[string]any) - gt.True(t, ok) - gt.V(t, msg1["text"]).Equal("Parent message") - gt.V(t, msg1["user_name"]).Equal("user1") - gt.V(t, msg1["timestamp"]).Equal("1234567890.123456") - - // Verify reply message - msg2, ok := messages[1].(map[string]any) - gt.True(t, ok) - gt.V(t, msg2["text"]).Equal("Reply message") - gt.V(t, msg2["user_name"]).Equal("user2") - gt.V(t, msg2["timestamp"]).Equal("1234567891.654321") -} - -func TestInternalTool_GetContextMessages(t *testing.T) { - ctx := context.Background() - - slackClient := &domainmock.SlackClientMock{ - GetConversationHistoryContextFunc: func(ctx context.Context, params *slackSDK.GetConversationHistoryParameters) (*slackSDK.GetConversationHistoryResponse, error) { - // Verify channel parameter - gt.V(t, params.ChannelID).Equal("C123456") - - // Return different messages based on Latest/Oldest - if params.Latest != "" { - // Before messages - gt.V(t, params.Latest).Equal("1234567890.000000") - gt.V(t, params.Inclusive).Equal(false) - return &slackSDK.GetConversationHistoryResponse{ - Messages: []slackSDK.Message{ - { - Msg: slackSDK.Msg{ - Timestamp: "1234567888.123456", - Text: "Message before", - Username: "user1", - }, - }, - }, - }, nil - } else if params.Oldest != "" { - // After messages - gt.V(t, params.Oldest).Equal("1234567890.000000") - gt.V(t, params.Inclusive).Equal(false) - return &slackSDK.GetConversationHistoryResponse{ - Messages: []slackSDK.Message{ - { - Msg: slackSDK.Msg{ - Timestamp: "1234567892.654321", - Text: "Message after", - Username: "user2", - }, - }, - }, - }, nil - } - return &slackSDK.GetConversationHistoryResponse{ - Messages: []slackSDK.Message{}, - }, nil - }, - } - - tool := slackagent.NewInternalToolForTest(slackClient, 0) - - // Call getContextMessages directly - result, err := tool.Run(ctx, "slack_get_context_messages", map[string]any{ - "channel": "C123456", - "around_ts": "1234567890", - "before": float64(1), - "after": float64(1), - }) - - gt.NoError(t, err) - gt.V(t, result).NotNil() - - // Verify response structure - beforeMessages, ok := result["before_messages"].([]any) - gt.True(t, ok) - gt.V(t, len(beforeMessages)).Equal(1) - - afterMessages, ok := result["after_messages"].([]any) - gt.True(t, ok) - gt.V(t, len(afterMessages)).Equal(1) - - // Verify before message - beforeMsg, ok := beforeMessages[0].(map[string]any) - gt.True(t, ok) - gt.V(t, beforeMsg["text"]).Equal("Message before") - gt.V(t, beforeMsg["user_name"]).Equal("user1") - - // Verify after message - afterMsg, ok := afterMessages[0].(map[string]any) - gt.True(t, ok) - gt.V(t, afterMsg["text"]).Equal("Message after") - gt.V(t, afterMsg["user_name"]).Equal("user2") -} - -func TestInternalTool_DirectLimitEnforcement(t *testing.T) { - ctx := context.Background() - - // Create 250 mock messages (exceeds limit) - matches := make([]slackSDK.SearchMessage, 250) - for i := 0; i < 250; i++ { - matches[i] = slackSDK.SearchMessage{ - Type: "message", - Timestamp: "1234567890.123456", - Text: "Test message", - Username: "user", - Channel: slackSDK.CtxChannel{ - ID: "C123", - Name: "general", - }, - } - } - - slackClient := &domainmock.SlackClientMock{ - SearchMessagesContextFunc: func(ctx context.Context, query string, params slackSDK.SearchParameters) (*slackSDK.SearchMessages, error) { - return &slackSDK.SearchMessages{ - Total: 250, - Matches: matches, - }, nil - }, - } - - // Create internal tool directly with maxLimit set - tool := slackagent.NewInternalToolForTest(slackClient, 50) - - // Call searchMessages directly - result, err := tool.Run(ctx, "slack_search_messages", map[string]any{ - "query": "test", - }) - - gt.NoError(t, err) - gt.V(t, result).NotNil() - - // Verify that result has messages - messages, ok := result["messages"].([]any) - gt.True(t, ok) - - // Verify that only 50 messages were returned (enforced by maxLimit) - gt.V(t, len(messages)).Equal(50) - - // Verify total count is still 250 (from Slack API) - total, ok := result["total"].(float64) - gt.True(t, ok) - gt.V(t, total).Equal(250.0) -} diff --git a/pkg/cli/chat.go b/pkg/cli/chat.go index 3ce7f308..f7ad3fc8 100644 --- a/pkg/cli/chat.go +++ b/pkg/cli/chat.go @@ -11,7 +11,6 @@ import ( "github.com/m-mizutani/goerr/v2" "github.com/secmon-lab/warren/pkg/adapter/trace" - "github.com/secmon-lab/warren/pkg/agents" "github.com/secmon-lab/warren/pkg/cli/config" "github.com/secmon-lab/warren/pkg/domain/model/ticket" "github.com/secmon-lab/warren/pkg/domain/types" @@ -68,7 +67,6 @@ func cmdChat() *cli.Command { storageCfg.Flags(), tools.Flags(), mcpCfg.Flags(), - agents.AllFlags(), traceCfg.Flags(), userSystemPromptCfg.Flags(), ) @@ -134,13 +132,6 @@ func cmdChat() *cli.Command { return goerr.Wrap(err, "failed to get tool sets") } - // Initialize all configured agents and merge into tool sets - agentToolSets, err := agents.ConfigureAll(ctx) - if err != nil { - return goerr.Wrap(err, "failed to configure agents") - } - allToolSets = append(allToolSets, agentToolSets...) - // Show ticket information fmt.Printf("\n🎫 Ticket Information:\n") fmt.Printf(" 📝 ID: %s\n", ticket.ID) diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go index e09166f1..0cbe54a7 100644 --- a/pkg/cli/serve.go +++ b/pkg/cli/serve.go @@ -17,7 +17,6 @@ import ( "github.com/secmon-lab/warren/pkg/adapter/storage" traceAdapter "github.com/secmon-lab/warren/pkg/adapter/trace" - "github.com/secmon-lab/warren/pkg/agents" "github.com/secmon-lab/warren/pkg/cli/config" server "github.com/secmon-lab/warren/pkg/controller/http" websocket_controller "github.com/secmon-lab/warren/pkg/controller/websocket" @@ -175,7 +174,6 @@ func cmdServe() *cli.Command { storageCfg.Flags(), mcpCfg.Flags(), asyncCfg.Flags(), - agents.AllFlags(), traceCfg.Flags(), userSystemPromptCfg.Flags(), strategyPromptsCfg.Flags(), @@ -322,13 +320,6 @@ func cmdServe() *cli.Command { // Create tag service tagService := tag.New(repo) - // Initialize all configured agents and merge into tool sets - agentToolSets, err := agents.ConfigureAll(ctx) - if err != nil { - return goerr.Wrap(err, "failed to configure agents") - } - toolSets = append(toolSets, agentToolSets...) - // Configure user system prompt userSystemPrompt, err := userSystemPromptCfg.Configure() if err != nil { diff --git a/pkg/cli/utils.go b/pkg/cli/utils.go index 15371abb..dd13e6ab 100644 --- a/pkg/cli/utils.go +++ b/pkg/cli/utils.go @@ -9,9 +9,11 @@ import ( "github.com/secmon-lab/warren/pkg/domain/interfaces" "github.com/secmon-lab/warren/pkg/tool/abusech" "github.com/secmon-lab/warren/pkg/tool/bigquery" + "github.com/secmon-lab/warren/pkg/tool/falcon" "github.com/secmon-lab/warren/pkg/tool/github" "github.com/secmon-lab/warren/pkg/tool/intune" "github.com/secmon-lab/warren/pkg/tool/ipdb" + "github.com/secmon-lab/warren/pkg/tool/jira" "github.com/secmon-lab/warren/pkg/tool/otx" "github.com/secmon-lab/warren/pkg/tool/shodan" "github.com/secmon-lab/warren/pkg/tool/slack" @@ -47,6 +49,8 @@ var tools = toolList{ &whois.Action{}, &intune.Action{}, &webfetch.Action{}, + &falcon.Action{}, + &jira.Action{}, } // InjectDependencies injects repository and embedding client into tools that support them diff --git a/pkg/tool/falcon/README.md b/pkg/tool/falcon/README.md new file mode 100644 index 00000000..4d5e8ffa --- /dev/null +++ b/pkg/tool/falcon/README.md @@ -0,0 +1,52 @@ +# CrowdStrike Falcon Tool Configuration + +## Overview + +The Falcon tool enables Warren to query CrowdStrike Falcon for incidents, alerts, behaviors, devices, CrowdScores, and raw EDR telemetry events (Next-Gen SIEM). Security analysts can pivot from alert context into endpoint detection data during an investigation. + +This is a **read-only** tool — it does not modify any data in your CrowdStrike environment. It is backed by [`github.com/gollem-dev/tools/falcon`](https://github.com/gollem-dev/tools). + +> **Note:** This replaces the former CrowdStrike Falcon *sub-agent* (`WARREN_AGENT_FALCON_*`). See [doc/migration/v0.18.0.md](../../../doc/migration/v0.18.0.md). + +## Configuration + +| Environment Variable | CLI Flag | Default | Required | +|---|---|---|---| +| `WARREN_FALCON_CLIENT_ID` | `--falcon-client-id` | - | Yes | +| `WARREN_FALCON_CLIENT_SECRET` | `--falcon-client-secret` | - | Yes | +| `WARREN_FALCON_BASE_URL` | `--falcon-base-url` | `https://api.crowdstrike.com` | No | + +The tool is enabled only when both `WARREN_FALCON_CLIENT_ID` and `WARREN_FALCON_CLIENT_SECRET` are set; otherwise it is silently skipped. + +## Setup + +1. Sign in to the [CrowdStrike Falcon console](https://falcon.crowdstrike.com/). +2. Go to **Support and resources > API clients and keys**. +3. Click **Create API client** and grant the following **read** scopes: + - Incidents (read) + - Alerts (read) + - Hosts / Devices (read) + - Event search / Next-Gen SIEM (read), if EDR event search is needed +4. Copy the **Client ID** and **Client Secret**. +5. Set the **Base URL** matching your CrowdStrike cloud region (US-1 is the default; US-2, EU-1, US-GOV-1 differ). + +## Available Functions + +| Function | Description | +|---|---| +| `falcon_search_incidents` | Search incidents via FQL | +| `falcon_get_incidents` | Fetch incident details by IDs | +| `falcon_search_alerts` | Search alerts via FQL | +| `falcon_get_alerts` | Fetch alert details by composite IDs | +| `falcon_search_behaviors` | Search behaviors via FQL | +| `falcon_get_behaviors` | Fetch behavior details by IDs | +| `falcon_search_devices` | Search devices (hosts) via FQL | +| `falcon_get_devices` | Fetch device details by IDs | +| `falcon_get_crowdscores` | Retrieve CrowdScore environment risk values | +| `falcon_search_events` | Search raw EDR telemetry events (Next-Gen SIEM) | + +To avoid overwhelming the LLM, search functions return a bounded number of records per call and keep the overflow in an in-memory page store; the LLM fetches subsequent pages by passing the returned `page_token` back to the same function. + +## Authentication + +The tool uses the CrowdStrike OAuth2 client credentials flow (`POST {base_url}/oauth2/token`), caches the bearer token, and retries once on a 401 by clearing the cached token. No manual token management is required. diff --git a/pkg/tool/falcon/action.go b/pkg/tool/falcon/action.go new file mode 100644 index 00000000..14e1daf3 --- /dev/null +++ b/pkg/tool/falcon/action.go @@ -0,0 +1,113 @@ +package falcon + +import ( + "context" + "log/slog" + + "github.com/gollem-dev/gollem" + extfalcon "github.com/gollem-dev/tools/falcon" + "github.com/m-mizutani/goerr/v2" + "github.com/secmon-lab/warren/pkg/domain/interfaces" + "github.com/secmon-lab/warren/pkg/utils/errutil" + "github.com/urfave/cli/v3" +) + +const defaultBaseURL = "https://api.crowdstrike.com" + +// Action is the warren-side wrapper around github.com/gollem-dev/tools/falcon. +// It implements interfaces.Tool, binding CLI flags and warren-specific planner +// metadata onto the external gollem.ToolSet that carries the read-only +// CrowdStrike Falcon query logic (incidents, alerts, behaviors, devices, +// CrowdScores, and EDR telemetry events). +type Action struct { + clientID string + clientSecret string + baseURL string + + inner gollem.ToolSet +} + +var _ interfaces.Tool = &Action{} + +func (x *Action) ID() string { + return "falcon" +} + +func (x *Action) Description() string { + return "CrowdStrike Falcon read-only query (incidents, alerts, behaviors, devices, CrowdScores, EDR events)" +} + +func (x *Action) Helper() *cli.Command { + return nil +} + +func (x *Action) Flags() []cli.Flag { + return []cli.Flag{ + &cli.StringFlag{ + Name: "falcon-client-id", + Usage: "CrowdStrike Falcon API client ID", + Destination: &x.clientID, + Category: "Tool", + Sources: cli.EnvVars("WARREN_FALCON_CLIENT_ID"), + }, + &cli.StringFlag{ + Name: "falcon-client-secret", + Usage: "CrowdStrike Falcon API client secret", + Destination: &x.clientSecret, + Category: "Tool", + Sources: cli.EnvVars("WARREN_FALCON_CLIENT_SECRET"), + }, + &cli.StringFlag{ + Name: "falcon-base-url", + Usage: "CrowdStrike Falcon API base URL", + Destination: &x.baseURL, + Category: "Tool", + Value: defaultBaseURL, + Sources: cli.EnvVars("WARREN_FALCON_BASE_URL"), + }, + } +} + +func (x *Action) Configure(_ context.Context) error { + if x.clientID == "" || x.clientSecret == "" { + return errutil.ErrActionUnavailable + } + + var opts []extfalcon.Option + if x.baseURL != "" { + opts = append(opts, extfalcon.WithBaseURL(x.baseURL)) + } + + ts, err := extfalcon.New(x.clientID, x.clientSecret, opts...) + if err != nil { + return goerr.Wrap(err, "failed to configure Falcon tool") + } + x.inner = ts + return nil +} + +func (x *Action) LogValue() slog.Value { + return slog.GroupValue( + slog.Int("client_id.len", len(x.clientID)), + slog.Int("client_secret.len", len(x.clientSecret)), + slog.String("base_url", x.baseURL), + ) +} + +func (x *Action) Prompt(_ context.Context) (string, error) { + return "", nil +} + +func (x *Action) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { + if x.inner == nil { + return nil, goerr.New("Falcon tool is not configured") + } + return x.inner.Specs(ctx) +} + +func (x *Action) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { + if x.inner == nil { + return nil, goerr.New("Falcon tool is not configured") + } + return x.inner.Run(ctx, name, args) +} diff --git a/pkg/tool/falcon/action_test.go b/pkg/tool/falcon/action_test.go new file mode 100644 index 00000000..5faf7ec7 --- /dev/null +++ b/pkg/tool/falcon/action_test.go @@ -0,0 +1,172 @@ +package falcon_test + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/m-mizutani/gt" + "github.com/secmon-lab/warren/pkg/tool/falcon" + "github.com/secmon-lab/warren/pkg/utils/errutil" + "github.com/secmon-lab/warren/pkg/utils/test" + "github.com/urfave/cli/v3" +) + +// newStubServer returns an httptest server that answers the Falcon OAuth2 token +// endpoint and returns an empty (but well-formed) search response for any other +// path. capturedAuth receives the Authorization header of the last API call. +func newStubServer(t *testing.T, capturedAuth *string) *httptest.Server { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/oauth2/token" { + _, _ = w.Write([]byte(`{"access_token":"test-token","expires_in":1800,"token_type":"bearer"}`)) + return + } + if capturedAuth != nil { + *capturedAuth = r.Header.Get("Authorization") + } + _, _ = w.Write([]byte(`{"resources":[],"meta":{"pagination":{"total":0}}}`)) + })) + t.Cleanup(server.Close) + return server +} + +func newConfiguredAction(t *testing.T, serverURL string) *falcon.Action { + t.Helper() + action := &falcon.Action{} + action.SetCredentials("test-client-id", "test-client-secret") + action.SetTestURL(serverURL) + gt.NoError(t, action.Configure(context.Background())) + return action +} + +func TestFalcon_Configure(t *testing.T) { + t.Run("with credentials", func(t *testing.T) { + action := &falcon.Action{} + action.SetCredentials("id", "secret") + gt.NoError(t, action.Configure(context.Background())) + }) + + t.Run("missing client id", func(t *testing.T) { + action := &falcon.Action{} + action.SetCredentials("", "secret") + gt.Equal(t, action.Configure(context.Background()), errutil.ErrActionUnavailable) + }) + + t.Run("missing client secret", func(t *testing.T) { + action := &falcon.Action{} + action.SetCredentials("id", "") + gt.Equal(t, action.Configure(context.Background()), errutil.ErrActionUnavailable) + }) +} + +func TestFalcon_Specs(t *testing.T) { + server := newStubServer(t, nil) + action := newConfiguredAction(t, server.URL) + + specs, err := action.Specs(context.Background()) + gt.NoError(t, err) + gt.A(t, specs).Length(10) + + got := map[string]bool{} + for _, s := range specs { + got[s.Name] = true + } + for _, name := range []string{ + "falcon_search_incidents", + "falcon_get_incidents", + "falcon_search_alerts", + "falcon_get_alerts", + "falcon_search_behaviors", + "falcon_get_behaviors", + "falcon_search_devices", + "falcon_get_devices", + "falcon_get_crowdscores", + "falcon_search_events", + } { + gt.True(t, got[name]) + } +} + +func TestFalcon_RunDelegation(t *testing.T) { + var gotAuth string + server := newStubServer(t, &gotAuth) + action := newConfiguredAction(t, server.URL) + + result, err := action.Run(context.Background(), "falcon_search_incidents", map[string]any{}) + gt.NoError(t, err) + gt.NotNil(t, result) + // The token fetched from /oauth2/token must be sent as a Bearer credential. + gt.Value(t, gotAuth).Equal("Bearer test-token") +} + +func TestFalcon_RunNotConfigured(t *testing.T) { + action := &falcon.Action{} + _, err := action.Run(context.Background(), "falcon_search_incidents", map[string]any{}) + gt.Error(t, err) + + _, err = action.Specs(context.Background()) + gt.Error(t, err) +} + +func TestFalcon_Identity(t *testing.T) { + action := &falcon.Action{} + gt.Value(t, action.ID()).Equal("falcon") + gt.True(t, action.Description() != "") + prompt, err := action.Prompt(context.Background()) + gt.NoError(t, err) + gt.Value(t, prompt).Equal("") +} + +// TestFalcon_FlagBinding verifies that --falcon-client-id / --falcon-client-secret +// bind via cli.Command parsing and that the resulting tool uses them: the stub +// server is reached and answers a search call. +func TestFalcon_FlagBinding(t *testing.T) { + var gotAuth string + server := newStubServer(t, &gotAuth) + + action := &falcon.Action{} + cmd := cli.Command{ + Name: "falcon", + Flags: action.Flags(), + Action: func(ctx context.Context, _ *cli.Command) error { + // Point at the stub before configuring; the parsed flags supply credentials. + action.SetTestURL(server.URL) + gt.NoError(t, action.Configure(ctx)) + _, err := action.Run(ctx, "falcon_search_devices", map[string]any{}) + gt.NoError(t, err) + return nil + }, + } + gt.NoError(t, cmd.Run(context.Background(), []string{ + "falcon", + "--falcon-client-id", "flag-client-id", + "--falcon-client-secret", "flag-client-secret", + })) + + gt.Value(t, gotAuth).Equal("Bearer test-token") +} + +func TestFalcon_Integration(t *testing.T) { + vars := test.NewEnvVars(t, "TEST_FALCON_CLIENT_ID", "TEST_FALCON_CLIENT_SECRET") + + action := &falcon.Action{} + cmd := cli.Command{ + Name: "falcon", + Flags: action.Flags(), + Action: func(ctx context.Context, _ *cli.Command) error { + gt.NoError(t, action.Configure(ctx)) + resp, err := action.Run(ctx, "falcon_search_incidents", map[string]any{}) + gt.NoError(t, err) + gt.NotNil(t, resp) + return nil + }, + } + gt.NoError(t, cmd.Run(context.Background(), []string{ + "falcon", + "--falcon-client-id", vars.Get("TEST_FALCON_CLIENT_ID"), + "--falcon-client-secret", vars.Get("TEST_FALCON_CLIENT_SECRET"), + })) +} diff --git a/pkg/tool/falcon/export_test.go b/pkg/tool/falcon/export_test.go new file mode 100644 index 00000000..4b2d2a6c --- /dev/null +++ b/pkg/tool/falcon/export_test.go @@ -0,0 +1,12 @@ +package falcon + +// SetCredentials sets the client credentials directly (test-only). +func (x *Action) SetCredentials(clientID, clientSecret string) { + x.clientID = clientID + x.clientSecret = clientSecret +} + +// SetTestURL points the underlying toolset at a stub server (test-only). +func (x *Action) SetTestURL(url string) { + x.baseURL = url +} diff --git a/pkg/tool/github/action_test.go b/pkg/tool/github/action_test.go index 13b1d7f9..f8e87f0d 100644 --- a/pkg/tool/github/action_test.go +++ b/pkg/tool/github/action_test.go @@ -138,7 +138,7 @@ func TestGitHubSpecsDelegation(t *testing.T) { specs, err := action.Specs(context.Background()) gt.NoError(t, err) - gt.A(t, specs).Length(5) + gt.A(t, specs).Length(7) names := map[string]bool{} for _, s := range specs { @@ -150,6 +150,8 @@ func TestGitHubSpecsDelegation(t *testing.T) { "github_get_content", "github_list_commits", "github_get_blame", + "github_get_issue", + "github_get_pull_request", } { gt.Value(t, names[want]).Equal(true) } diff --git a/pkg/tool/jira/README.md b/pkg/tool/jira/README.md new file mode 100644 index 00000000..19df4660 --- /dev/null +++ b/pkg/tool/jira/README.md @@ -0,0 +1,41 @@ +# Jira Tool Configuration + +## Overview + +The Jira tool enables Warren to query Jira Cloud: listing accessible projects, searching issues with JQL, and fetching the content of one or more issues (with the description rendered to Markdown). Security analysts can correlate alerts with tracked tickets during an investigation. + +This is a **read-only** tool — it does not create or modify any Jira data. It is backed by [`github.com/gollem-dev/tools/jira`](https://github.com/gollem-dev/tools). + +## Configuration + +| Environment Variable | CLI Flag | Default | Required | +|---|---|---|---| +| `WARREN_JIRA_BASE_URL` | `--jira-base-url` | - | Yes | +| `WARREN_JIRA_USER_EMAIL` | `--jira-user-email` | - | Yes | +| `WARREN_JIRA_API_TOKEN` | `--jira-api-token` | - | Yes | + +The tool is enabled only when all three values are set; otherwise it is silently skipped. The base URL is your Jira site, e.g. `https://your-domain.atlassian.net`. + +## Setup + +1. Sign in to Jira Cloud with the account Warren should authenticate as. +2. Go to [**Atlassian account > Security > API tokens**](https://id.atlassian.com/manage-profile/security/api-tokens). +3. Click **Create API token**, give it a label, and copy the generated token. +4. Configure Warren with: + - `WARREN_JIRA_BASE_URL` — your site URL (`https://<your-domain>.atlassian.net`) + - `WARREN_JIRA_USER_EMAIL` — the account email + - `WARREN_JIRA_API_TOKEN` — the token created above + +The account's existing permissions bound what the tool can read; no extra scopes are configured on the token itself. + +## Available Functions + +| Function | Description | +|---|---| +| `jira_list_projects` | List Jira projects accessible to the authenticated account | +| `jira_search_issues` | Search issues using a JQL query | +| `jira_get_issues` | Fetch the content of one or more issues by key/ID (description rendered to Markdown) | + +## Authentication + +The tool talks to the Jira Cloud REST API v3 directly over HTTP using Basic authentication (`base64(email:apiToken)`). Because each Jira site lives on its own tenant domain, the base URL is a required setting rather than a fixed constant. diff --git a/pkg/tool/jira/action.go b/pkg/tool/jira/action.go new file mode 100644 index 00000000..cc59bc96 --- /dev/null +++ b/pkg/tool/jira/action.go @@ -0,0 +1,104 @@ +package jira + +import ( + "context" + "log/slog" + + "github.com/gollem-dev/gollem" + extjira "github.com/gollem-dev/tools/jira" + "github.com/m-mizutani/goerr/v2" + "github.com/secmon-lab/warren/pkg/domain/interfaces" + "github.com/secmon-lab/warren/pkg/utils/errutil" + "github.com/urfave/cli/v3" +) + +// Action is the warren-side wrapper around github.com/gollem-dev/tools/jira. +// It implements interfaces.Tool, binding CLI flags and warren-specific planner +// metadata onto the external gollem.ToolSet that carries the read-only Jira +// Cloud query logic (list projects, search issues via JQL, fetch issues). +type Action struct { + baseURL string + email string + apiToken string + + inner gollem.ToolSet +} + +var _ interfaces.Tool = &Action{} + +func (x *Action) ID() string { + return "jira" +} + +func (x *Action) Description() string { + return "Jira Cloud read-only query (list projects, search issues via JQL, fetch issue content)" +} + +func (x *Action) Helper() *cli.Command { + return nil +} + +func (x *Action) Flags() []cli.Flag { + return []cli.Flag{ + &cli.StringFlag{ + Name: "jira-base-url", + Usage: "Jira Cloud site URL (e.g. https://your-domain.atlassian.net)", + Destination: &x.baseURL, + Category: "Tool", + Sources: cli.EnvVars("WARREN_JIRA_BASE_URL"), + }, + &cli.StringFlag{ + Name: "jira-user-email", + Usage: "Jira account email for Basic authentication", + Destination: &x.email, + Category: "Tool", + Sources: cli.EnvVars("WARREN_JIRA_USER_EMAIL"), + }, + &cli.StringFlag{ + Name: "jira-api-token", + Usage: "Jira Cloud API token (paired with the account email)", + Destination: &x.apiToken, + Category: "Tool", + Sources: cli.EnvVars("WARREN_JIRA_API_TOKEN"), + }, + } +} + +func (x *Action) Configure(_ context.Context) error { + if x.baseURL == "" || x.email == "" || x.apiToken == "" { + return errutil.ErrActionUnavailable + } + + ts, err := extjira.New(x.baseURL, x.email, x.apiToken) + if err != nil { + return goerr.Wrap(err, "failed to configure Jira tool") + } + x.inner = ts + return nil +} + +func (x *Action) LogValue() slog.Value { + return slog.GroupValue( + slog.String("base_url", x.baseURL), + slog.String("email", x.email), + slog.Int("api_token.len", len(x.apiToken)), + ) +} + +func (x *Action) Prompt(_ context.Context) (string, error) { + return "", nil +} + +func (x *Action) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { + if x.inner == nil { + return nil, goerr.New("Jira tool is not configured") + } + return x.inner.Specs(ctx) +} + +func (x *Action) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { + if x.inner == nil { + return nil, goerr.New("Jira tool is not configured") + } + return x.inner.Run(ctx, name, args) +} diff --git a/pkg/tool/jira/action_test.go b/pkg/tool/jira/action_test.go new file mode 100644 index 00000000..967aecd8 --- /dev/null +++ b/pkg/tool/jira/action_test.go @@ -0,0 +1,173 @@ +package jira_test + +import ( + "context" + "encoding/base64" + "net/http" + "net/http/httptest" + "testing" + + "github.com/m-mizutani/gt" + "github.com/secmon-lab/warren/pkg/tool/jira" + "github.com/secmon-lab/warren/pkg/utils/errutil" + "github.com/secmon-lab/warren/pkg/utils/test" + "github.com/urfave/cli/v3" +) + +const ( + testEmail = "user@example.com" + testToken = "test-api-token" +) + +// newStubServer returns an httptest server that answers the Jira project search +// endpoint with an empty result and records the Authorization header. +func newStubServer(t *testing.T, capturedAuth *string) *httptest.Server { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if capturedAuth != nil { + *capturedAuth = r.Header.Get("Authorization") + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"isLast":true,"total":0,"values":[]}`)) + })) + t.Cleanup(server.Close) + return server +} + +func newConfiguredAction(t *testing.T, serverURL string) *jira.Action { + t.Helper() + action := &jira.Action{} + action.SetConfig(serverURL, testEmail, testToken) + gt.NoError(t, action.Configure(context.Background())) + return action +} + +func TestJira_Configure(t *testing.T) { + t.Run("with all settings", func(t *testing.T) { + action := &jira.Action{} + action.SetConfig("https://example.atlassian.net", testEmail, testToken) + gt.NoError(t, action.Configure(context.Background())) + }) + + t.Run("missing base url", func(t *testing.T) { + action := &jira.Action{} + action.SetConfig("", testEmail, testToken) + gt.Equal(t, action.Configure(context.Background()), errutil.ErrActionUnavailable) + }) + + t.Run("missing email", func(t *testing.T) { + action := &jira.Action{} + action.SetConfig("https://example.atlassian.net", "", testToken) + gt.Equal(t, action.Configure(context.Background()), errutil.ErrActionUnavailable) + }) + + t.Run("missing api token", func(t *testing.T) { + action := &jira.Action{} + action.SetConfig("https://example.atlassian.net", testEmail, "") + gt.Equal(t, action.Configure(context.Background()), errutil.ErrActionUnavailable) + }) +} + +func TestJira_Specs(t *testing.T) { + server := newStubServer(t, nil) + action := newConfiguredAction(t, server.URL) + + specs, err := action.Specs(context.Background()) + gt.NoError(t, err) + gt.A(t, specs).Length(3) + + got := map[string]bool{} + for _, s := range specs { + got[s.Name] = true + } + for _, name := range []string{ + "jira_list_projects", + "jira_search_issues", + "jira_get_issues", + } { + gt.True(t, got[name]) + } +} + +func TestJira_RunDelegation(t *testing.T) { + var gotAuth string + server := newStubServer(t, &gotAuth) + action := newConfiguredAction(t, server.URL) + + result, err := action.Run(context.Background(), "jira_list_projects", map[string]any{}) + gt.NoError(t, err) + gt.NotNil(t, result) + + // Basic auth header must be base64("email:apiToken"). + wantCred := base64.StdEncoding.EncodeToString([]byte(testEmail + ":" + testToken)) + gt.Value(t, gotAuth).Equal("Basic " + wantCred) +} + +func TestJira_RunNotConfigured(t *testing.T) { + action := &jira.Action{} + _, err := action.Run(context.Background(), "jira_list_projects", map[string]any{}) + gt.Error(t, err) + + _, err = action.Specs(context.Background()) + gt.Error(t, err) +} + +func TestJira_Identity(t *testing.T) { + action := &jira.Action{} + gt.Value(t, action.ID()).Equal("jira") + gt.True(t, action.Description() != "") + prompt, err := action.Prompt(context.Background()) + gt.NoError(t, err) + gt.Value(t, prompt).Equal("") +} + +// TestJira_FlagBinding verifies that --jira-base-url / --jira-user-email / +// --jira-api-token bind via cli.Command parsing and reach the configured tool. +func TestJira_FlagBinding(t *testing.T) { + var gotAuth string + server := newStubServer(t, &gotAuth) + + action := &jira.Action{} + cmd := cli.Command{ + Name: "jira", + Flags: action.Flags(), + Action: func(ctx context.Context, _ *cli.Command) error { + gt.NoError(t, action.Configure(ctx)) + _, err := action.Run(ctx, "jira_list_projects", map[string]any{}) + gt.NoError(t, err) + return nil + }, + } + gt.NoError(t, cmd.Run(context.Background(), []string{ + "jira", + "--jira-base-url", server.URL, + "--jira-user-email", testEmail, + "--jira-api-token", testToken, + })) + + wantCred := base64.StdEncoding.EncodeToString([]byte(testEmail + ":" + testToken)) + gt.Value(t, gotAuth).Equal("Basic " + wantCred) +} + +func TestJira_Integration(t *testing.T) { + vars := test.NewEnvVars(t, "TEST_JIRA_BASE_URL", "TEST_JIRA_USER_EMAIL", "TEST_JIRA_API_TOKEN") + + action := &jira.Action{} + cmd := cli.Command{ + Name: "jira", + Flags: action.Flags(), + Action: func(ctx context.Context, _ *cli.Command) error { + gt.NoError(t, action.Configure(ctx)) + resp, err := action.Run(ctx, "jira_list_projects", map[string]any{}) + gt.NoError(t, err) + gt.NotNil(t, resp) + return nil + }, + } + gt.NoError(t, cmd.Run(context.Background(), []string{ + "jira", + "--jira-base-url", vars.Get("TEST_JIRA_BASE_URL"), + "--jira-user-email", vars.Get("TEST_JIRA_USER_EMAIL"), + "--jira-api-token", vars.Get("TEST_JIRA_API_TOKEN"), + })) +} diff --git a/pkg/tool/jira/export_test.go b/pkg/tool/jira/export_test.go new file mode 100644 index 00000000..a72a750e --- /dev/null +++ b/pkg/tool/jira/export_test.go @@ -0,0 +1,8 @@ +package jira + +// SetConfig sets the connection parameters directly (test-only). +func (x *Action) SetConfig(baseURL, email, apiToken string) { + x.baseURL = baseURL + x.email = email + x.apiToken = apiToken +} diff --git a/pkg/tool/slack/action_test.go b/pkg/tool/slack/action_test.go index bbada230..80dee091 100644 --- a/pkg/tool/slack/action_test.go +++ b/pkg/tool/slack/action_test.go @@ -7,6 +7,7 @@ import ( "os" "testing" + "github.com/gollem-dev/gollem" "github.com/m-mizutani/gt" "github.com/secmon-lab/warren/pkg/tool/slack" "github.com/secmon-lab/warren/pkg/utils/errutil" @@ -93,9 +94,17 @@ func TestSlackMessageSearch_Specs(t *testing.T) { action := newAction(t, server.URL) specs, err := action.Specs(context.Background()) gt.NoError(t, err) - gt.A(t, specs).Length(1) - gt.Value(t, specs[0].Name).Equal("slack_message_search") - gt.Map(t, specs[0].Parameters).HasKey("query") + gt.A(t, specs).Length(2) + + byName := map[string]gollem.ToolSpec{} + for _, s := range specs { + byName[s.Name] = s + } + search, ok := byName["slack_message_search"] + gt.True(t, ok) + gt.Map(t, search.Parameters).HasKey("query") + _, ok = byName["slack_get_messages"] + gt.True(t, ok) } func TestSlackConfigure(t *testing.T) { diff --git a/pkg/usecase/chat/aster/budget.go b/pkg/usecase/chat/aster/budget.go index db479319..71d0a870 100644 --- a/pkg/usecase/chat/aster/budget.go +++ b/pkg/usecase/chat/aster/budget.go @@ -189,47 +189,6 @@ func (t *BudgetTracker) GenerateHandoverInfo() string { return b.String() } -// budgetTrackerCtxKeyType is a context key type for budget tracker. -type budgetTrackerCtxKeyType struct{} - -// withBudgetTracker stores a BudgetTracker in the context. -func withBudgetTracker(ctx context.Context, tracker *BudgetTracker) context.Context { - return context.WithValue(ctx, budgetTrackerCtxKeyType{}, tracker) -} - -// budgetTrackerFrom retrieves a BudgetTracker from the context. -// Returns nil if no tracker is present. -func budgetTrackerFrom(ctx context.Context) *BudgetTracker { - v, _ := ctx.Value(budgetTrackerCtxKeyType{}).(*BudgetTracker) - return v -} - -// newContextAwareBudgetMiddleware creates a gollem.ToolMiddleware that reads -// the BudgetTracker from the context on each invocation. This allows a single -// middleware instance to be shared across tasks while each task provides its -// own tracker via the context. -// If no tracker is found in the context, the middleware passes through. -func newContextAwareBudgetMiddleware() gollem.ToolMiddleware { - return func(next gollem.ToolHandler) gollem.ToolHandler { - return func(ctx context.Context, req *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - tracker := budgetTrackerFrom(ctx) - if tracker == nil { - // No tracker in context — pass through (budget disabled) - return next(ctx, req) - } - return executeBudgetedToolCall(ctx, req, tracker, next) - } - } -} - -// subAgentToolNames contains tool names for sub-agent invocations. -// These have zero cost because their internal tool calls are tracked individually. -var subAgentToolNames = map[string]bool{ - "query_bigquery": true, - "query_falcon": true, - "query_slack": true, -} - // DefaultBudgetStrategy implements BudgetStrategy with sensible defaults. // Target: ~16 tool calls or ~3.5 minutes per task. type DefaultBudgetStrategy struct{} @@ -248,12 +207,10 @@ func (s *DefaultBudgetStrategy) InitialBudget() float64 { // Cost = tool fixed cost + time cost delta. func (s *DefaultBudgetStrategy) BeforeToolCall(ctx ToolCallContext) float64 { var toolCost float64 - switch { - case subAgentToolNames[ctx.ToolName]: - toolCost = 0 - case ctx.ToolName == "bigquery_query": + switch ctx.ToolName { + case "bigquery_query": toolCost = 15.0 - case ctx.ToolName == "bigquery_list_datasets" || ctx.ToolName == "bigquery_get_table_schema": + case "bigquery_list_datasets", "bigquery_get_table_schema": toolCost = 3.0 default: toolCost = 6.25 diff --git a/pkg/usecase/chat/aster/budget_test.go b/pkg/usecase/chat/aster/budget_test.go index e2d1dc6b..5f512a8f 100644 --- a/pkg/usecase/chat/aster/budget_test.go +++ b/pkg/usecase/chat/aster/budget_test.go @@ -33,14 +33,6 @@ func TestDefaultBudgetStrategy_InitialBudget(t *testing.T) { func TestDefaultBudgetStrategy_BeforeToolCall(t *testing.T) { s := aster.NewDefaultBudgetStrategy() - t.Run("sub-agent tool costs 0", func(t *testing.T) { - cost := s.BeforeToolCall(aster.ToolCallContext{ - ToolName: "query_bigquery", - CallCount: 1, - }) - gt.Equal(t, cost, 0.0) - }) - t.Run("bigquery_query costs 15", func(t *testing.T) { cost := s.BeforeToolCall(aster.ToolCallContext{ ToolName: "bigquery_query", @@ -287,201 +279,6 @@ func TestBudgetToolMiddleware_SoftLimit(t *testing.T) { gt.V(t, resp.Result["_budget_warning"]).NotNil() } -func TestBudgetToolMiddleware_SharedTrackerAcrossParentAndSubAgent(t *testing.T) { - strategy := &mockBudgetStrategy{ - initialBudget: 100.0, - beforeCost: 25.0, - afterCost: 0.0, - hardLimitMargin: 1, - } - tracker := aster.NewBudgetTracker(strategy) - - // Create two middleware instances from the same tracker, - // simulating the pattern in exec.go where: - // - parentMW is added to the parent agent (line 171) - // - subAgentMW is injected into sub-agents' child agents (line 131) - parentMW := aster.NewBudgetToolMiddleware(tracker) - subAgentMW := aster.NewBudgetToolMiddleware(tracker) - - passthrough := func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - return &gollem.ToolExecResponse{ - Result: map[string]any{"data": "ok"}, - }, nil - } - - parentHandler := parentMW(passthrough) - subAgentHandler := subAgentMW(passthrough) - - ctx := context.Background() - - // Parent tool call: 100 - 25 = 75 - resp, err := parentHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "parent_tool"}, - }) - gt.NoError(t, err) - gt.V(t, resp.Result["_budget_info"]).NotNil() - gt.Equal(t, tracker.Remaining(), 75.0) - - // Sub-agent internal tool call: 75 - 25 = 50 (shared budget) - _, err = subAgentHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "bigquery_query"}, - }) - gt.NoError(t, err) - gt.Equal(t, tracker.Remaining(), 50.0) - - // Another parent tool call: 50 - 25 = 25 - _, err = parentHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "parent_tool"}, - }) - gt.NoError(t, err) - gt.Equal(t, tracker.Remaining(), 25.0) - - // Sub-agent tool call: 25 - 25 = 0 → SoftLimit - resp, err = subAgentHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "bigquery_query"}, - }) - gt.NoError(t, err) - gt.V(t, resp.Result["_budget_warning"]).NotNil() // soft limit warning - gt.Equal(t, tracker.Remaining(), 0.0) - - // One more sub-agent call after soft: still allowed (margin=1) - _, err = subAgentHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "bigquery_query"}, - }) - gt.NoError(t, err) - - // Next call triggers hard limit regardless of caller (parent or sub-agent) - subAgentCalled := false - hardLimitHandler := subAgentMW(func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - subAgentCalled = true - return &gollem.ToolExecResponse{Result: map[string]any{}}, nil - }) - resp, err = hardLimitHandler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "bigquery_query"}, - }) - gt.NoError(t, err) - gt.B(t, subAgentCalled).False() // tool should be blocked - gt.V(t, resp.Result["error"]).NotNil() -} - -func TestContextAwareBudgetMiddleware_ReadsTrackerFromContext(t *testing.T) { - strategy := &mockBudgetStrategy{ - initialBudget: 100.0, - beforeCost: 10.0, - afterCost: 0.0, - hardLimitMargin: 3, - } - tracker := aster.NewBudgetTracker(strategy) - mw := aster.NewContextAwareBudgetMiddleware() - - handler := mw(func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - return &gollem.ToolExecResponse{ - Result: map[string]any{"data": "test"}, - }, nil - }) - - ctx := aster.WithBudgetTracker(context.Background(), tracker) - resp, err := handler(ctx, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "test_tool"}, - }) - - gt.NoError(t, err) - gt.V(t, resp.Result["_budget_info"]).NotNil() - gt.Equal(t, tracker.Remaining(), 90.0) -} - -func TestContextAwareBudgetMiddleware_NoTrackerInContext(t *testing.T) { - mw := aster.NewContextAwareBudgetMiddleware() - - toolCalled := false - handler := mw(func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - toolCalled = true - return &gollem.ToolExecResponse{ - Result: map[string]any{"data": "test"}, - }, nil - }) - - // No tracker in context — should pass through - resp, err := handler(context.Background(), &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "test_tool"}, - }) - - gt.NoError(t, err) - gt.B(t, toolCalled).True() - // No budget info should be added - gt.V(t, resp.Result["_budget_info"]).Nil() - gt.V(t, resp.Result["_budget_warning"]).Nil() -} - -func TestContextAwareBudgetMiddleware_DifferentTrackersPerTask(t *testing.T) { - strategy := &mockBudgetStrategy{ - initialBudget: 20.0, - beforeCost: 10.0, - afterCost: 0.0, - hardLimitMargin: 0, // hard limit immediately after soft - } - - mw := aster.NewContextAwareBudgetMiddleware() - - passthrough := func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - return &gollem.ToolExecResponse{ - Result: map[string]any{"data": "ok"}, - }, nil - } - handler := mw(passthrough) - - // === Task 1: exhaust the budget === - tracker1 := aster.NewBudgetTracker(strategy) - ctx1 := aster.WithBudgetTracker(context.Background(), tracker1) - - // Call 1: 20 - 10 = 10 (OK) - resp, err := handler(ctx1, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "tool1"}, - }) - gt.NoError(t, err) - gt.V(t, resp.Result["_budget_info"]).NotNil() - - // Call 2: 10 - 10 = 0 (SoftLimit) - resp, err = handler(ctx1, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "tool1"}, - }) - gt.NoError(t, err) - gt.V(t, resp.Result["_budget_warning"]).NotNil() - - // Call 3: HardLimit (margin=0, 1st call after soft) - toolCalled := false - hardHandler := mw(func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - toolCalled = true - return &gollem.ToolExecResponse{Result: map[string]any{}}, nil - }) - resp, err = hardHandler(ctx1, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "tool1"}, - }) - gt.NoError(t, err) - gt.B(t, toolCalled).False() // blocked by hard limit - gt.V(t, resp.Result["error"]).NotNil() - - // === Task 2: new tracker, same middleware — should work fine === - tracker2 := aster.NewBudgetTracker(strategy) - ctx2 := aster.WithBudgetTracker(context.Background(), tracker2) - - // This call must succeed with the fresh tracker, NOT be blocked by tracker1 - toolCalled = false - freshHandler := mw(func(_ context.Context, _ *gollem.ToolExecRequest) (*gollem.ToolExecResponse, error) { - toolCalled = true - return &gollem.ToolExecResponse{ - Result: map[string]any{"data": "fresh"}, - }, nil - }) - resp, err = freshHandler(ctx2, &gollem.ToolExecRequest{ - Tool: &gollem.FunctionCall{Name: "tool1"}, - }) - gt.NoError(t, err) - gt.B(t, toolCalled).True() // tool should execute successfully - gt.V(t, resp.Result["_budget_info"]).NotNil() - gt.Equal(t, tracker2.Remaining(), 10.0) // 20 - 10 = 10 -} - func TestBudgetToolMiddleware_HardLimit(t *testing.T) { strategy := &mockBudgetStrategy{ initialBudget: 10.0, diff --git a/pkg/usecase/chat/aster/exec.go b/pkg/usecase/chat/aster/exec.go index 623fbd58..aac732a3 100644 --- a/pkg/usecase/chat/aster/exec.go +++ b/pkg/usecase/chat/aster/exec.go @@ -169,7 +169,6 @@ func (c *AsterChat) executeTask(ctx context.Context, task TaskPlan, chatCtx *cha var tracker *BudgetTracker if c.budgetStrategy != nil { tracker = newBudgetTracker(c.budgetStrategy) - taskCtx = withBudgetTracker(taskCtx, tracker) } // Build agent options diff --git a/pkg/usecase/chat/aster/export_test.go b/pkg/usecase/chat/aster/export_test.go index 16ce8307..3353aa57 100644 --- a/pkg/usecase/chat/aster/export_test.go +++ b/pkg/usecase/chat/aster/export_test.go @@ -40,16 +40,6 @@ func NewBudgetToolMiddleware(tracker *BudgetTracker) gollem.ToolMiddleware { return newBudgetToolMiddleware(tracker) } -// NewContextAwareBudgetMiddleware exposes newContextAwareBudgetMiddleware for testing. -func NewContextAwareBudgetMiddleware() gollem.ToolMiddleware { - return newContextAwareBudgetMiddleware() -} - -// WithBudgetTracker exposes withBudgetTracker for testing. -func WithBudgetTracker(ctx context.Context, tracker *BudgetTracker) context.Context { - return withBudgetTracker(ctx, tracker) -} - // FilterToolSets exposes filterToolSets for testing. func FilterToolSets(allTools []interfaces.ToolSet, allowedIDs []string) []interfaces.ToolSet { return filterToolSets(allTools, allowedIDs) diff --git a/pkg/usecase/chat/bluebell/budget.go b/pkg/usecase/chat/bluebell/budget.go index bd02b83b..8d2e8efe 100644 --- a/pkg/usecase/chat/bluebell/budget.go +++ b/pkg/usecase/chat/bluebell/budget.go @@ -189,22 +189,6 @@ func (t *BudgetTracker) GenerateHandoverInfo() string { return b.String() } -// budgetTrackerCtxKeyType is a context key type for budget tracker. -type budgetTrackerCtxKeyType struct{} - -// withBudgetTracker stores a BudgetTracker in the context. -func withBudgetTracker(ctx context.Context, tracker *BudgetTracker) context.Context { - return context.WithValue(ctx, budgetTrackerCtxKeyType{}, tracker) -} - -// subAgentToolNames contains tool names for sub-agent invocations. -// These have zero cost because their internal tool calls are tracked individually. -var subAgentToolNames = map[string]bool{ - "query_bigquery": true, - "query_falcon": true, - "query_slack": true, -} - // DefaultBudgetStrategy implements BudgetStrategy with sensible defaults. // Target: ~16 tool calls or ~3.5 minutes per task. type DefaultBudgetStrategy struct{} @@ -223,12 +207,10 @@ func (s *DefaultBudgetStrategy) InitialBudget() float64 { // Cost = tool fixed cost + time cost delta. func (s *DefaultBudgetStrategy) BeforeToolCall(ctx ToolCallContext) float64 { var toolCost float64 - switch { - case subAgentToolNames[ctx.ToolName]: - toolCost = 0 - case ctx.ToolName == "bigquery_query": + switch ctx.ToolName { + case "bigquery_query": toolCost = 15.0 - case ctx.ToolName == "bigquery_list_datasets" || ctx.ToolName == "bigquery_get_table_schema": + case "bigquery_list_datasets", "bigquery_get_table_schema": toolCost = 3.0 default: toolCost = 6.25 diff --git a/pkg/usecase/chat/bluebell/exec.go b/pkg/usecase/chat/bluebell/exec.go index 31543da4..de29b9ca 100644 --- a/pkg/usecase/chat/bluebell/exec.go +++ b/pkg/usecase/chat/bluebell/exec.go @@ -166,7 +166,6 @@ func (c *BluebellChat) executeTask(ctx context.Context, task TaskPlan, chatCtx * var tracker *BudgetTracker if c.budgetStrategy != nil { tracker = newBudgetTracker(c.budgetStrategy) - taskCtx = withBudgetTracker(taskCtx, tracker) } // Build agent options