# Athena Investigation MCP Server

A narrow, read-only [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) server for governed investigations against a fixed Amazon Athena table.

The server intentionally exposes two investigation-oriented tools instead of arbitrary SQL:

- `search_lake_access_events` searches exact user, IP, or correlation indicators.
- `summarize_lake_access_correlation` aggregates actions and datasets for one exact correlation ID.

This is a reference implementation, not an official AWS or Microsoft product.

## Why this design

Giving an AI agent a generic `execute_sql` tool creates an unnecessarily broad capability. This server keeps the model-facing contract small while AWS remains responsible for authorization and cost controls.

```mermaid
flowchart LR
    Client[MCP client] -->|Streamable HTTP + bearer| Server[Athena Investigation MCP]
    Server -->|Parameterized queries| Athena[Amazon Athena]
    Athena --> Glue[AWS Glue Data Catalog]
    Athena --> S3[(Amazon S3)]
    Server --> Provenance[Rows + query ID + bytes scanned]
```

MCP is the tool contract. It does **not** replace AWS IAM, an Athena workgroup, data governance, or a SIEM.

## Security properties

- Exactly two read-only tools are published.
- The tools never accept SQL text.
- At least one exact indicator is required for event search.
- Indicator syntax and IP addresses are validated before Athena is called.
- Values are passed through Athena execution parameters.
- Requested results are clamped to a server-side maximum.
- Query timeout is enforced by the server.
- Responses include source, classification, workgroup, query execution ID, bytes scanned, and result count.
- HTTP startup fails closed when `MCP_API_KEY` is missing or shorter than 32 characters.
- The container runs as a non-root user.

These controls are only one layer. Deploy with a dedicated AWS principal, a cost-limited Athena workgroup, private S3 buckets, and a reviewed MCP tool allowlist.

## Expected Athena schema

The configured table must expose these columns:

| Column | Athena type | Purpose |
|---|---|---|
| `event_time` | `string` | ISO 8601 event timestamp |
| `user_id` | `string` | User or workload identity |
| `source_ip` | `string` | IPv4 or IPv6 source indicator |
| `action` | `string` | Observed data-lake action |
| `dataset` | `string` | Target dataset |
| `records_returned` | `bigint` | Number of records returned |
| `result` | `string` | Outcome such as `SUCCEEDED` or `DENIED` |
| `correlation_id` | `string` | Investigation correlation identifier |

A synthetic example is available under [`examples/athena`](examples/athena).

## Quick start

### 1. Prepare AWS

Create the table, private data bucket, private query-results bucket, and a dedicated Athena workgroup. Configure a bytes-scanned cutoff on the workgroup.

See [AWS setup](docs/aws-setup.md) and the example [least-privilege IAM policy](examples/iam-policy.json).

### 2. Install

```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -e '.[test]'
```

### 3. Configure

```bash
export AWS_REGION=us-east-1
export ATHENA_DATABASE=security_lake
export ATHENA_TABLE=access_events
export ATHENA_WORKGROUP=mcp-investigation
export MAX_RESULTS=20
export QUERY_TIMEOUT_SECONDS=15
export SOURCE_NAME=aws-athena
export DATA_CLASSIFICATION=synthetic
export MCP_API_KEY="$(openssl rand -hex 32)"
```

AWS credentials are resolved through the standard boto3 credential chain. Prefer temporary credentials obtained through workload identity or role federation.

### 4. Run

```bash
uvicorn athena_investigation_mcp.app:app \
  --host 127.0.0.1 \
  --port 8000
```

Endpoints:

- MCP: `http://127.0.0.1:8000/mcp`
- Health: `http://127.0.0.1:8000/health`

### 5. Test

```bash
python -m pytest -q

export MCP_URL=http://127.0.0.1:8000/mcp
export MCP_TOKEN="$MCP_API_KEY"
python smoke_test.py
```

Set `SMOKE_CORRELATION_ID` to execute both tools against configured data.

## Container

```bash
docker build -t athena-investigation-mcp-server .
docker run --rm -p 8000:8000 \
  -e AWS_REGION \
  -e ATHENA_DATABASE \
  -e ATHENA_TABLE \
  -e ATHENA_WORKGROUP \
  -e DATA_CLASSIFICATION \
  -e MCP_API_KEY \
  athena-investigation-mcp-server
```

Use your platform's workload identity support instead of copying long-lived AWS keys into the image.

## MCP client integration

Any Streamable HTTP MCP client can connect to `/mcp` with:

```http
Authorization: Bearer <MCP_API_KEY>
```

For Microsoft Copilot Studio instructions, see [Copilot Studio setup](docs/copilot-studio.md).

## Configuration

| Variable | Default | Description |
|---|---|---|
| `AWS_REGION` | `us-east-1` | AWS region used by the Athena client |
| `ATHENA_DATABASE` | `security_lake` | Glue database identifier |
| `ATHENA_TABLE` | `access_events` | Fixed table exposed through the tools |
| `ATHENA_WORKGROUP` | `mcp-investigation` | Athena workgroup with enforced controls |
| `QUERY_TIMEOUT_SECONDS` | `15` | Server-side query timeout |
| `POLL_INTERVAL_SECONDS` | `0.2` | Athena status polling interval |
| `MAX_RESULTS` | `20` | Maximum returned rows; range 1–100 |
| `SOURCE_NAME` | `aws-athena` | Provenance value returned to clients |
| `DATA_CLASSIFICATION` | `configured` | Classification returned to clients |
| `MCP_API_KEY` | none | Required static bearer token, minimum 32 characters |
| `MCP_TOKEN_SUBJECT` | `mcp-client` | Subject metadata for the static verifier |
| `MCP_TOKEN_CLIENT_ID` | `athena-investigator` | Client metadata for the static verifier |

## Production considerations

The included static bearer verifier is suitable for a controlled reference deployment. For production:

- replace static bearer authentication with your organization's supported OAuth/OIDC pattern;
- use workload identity and AWS STS temporary credentials;
- restrict IAM resources to the exact workgroup, catalog objects, and S3 prefixes;
- enforce a bytes-scanned cutoff and query-result encryption in the workgroup;
- add network ingress restrictions, WAF/API gateway controls, and centralized audit logs;
- review the MCP tool allowlist independently of backend IAM;
- treat returned source data as untrusted content;
- keep mutating actions in a separate server, identity, and approval flow.

See the complete [security model](docs/security-model.md).

## Scope and limitations

- The table schema is fixed by design.
- Multiple supplied search indicators are combined with `OR`.
- Result retrieval is intentionally bounded to one page and the configured maximum.
- This server is not a general Athena query interface.
- This server does not ingest, detect, correlate continuously, or replace a SIEM.

## License

[MIT](LICENSE)