Skip to content

Commit 1d41684

Browse files
KongZKongZ
authored
Merge pull request #18 from KongZ/add-aws-cli
Add aws-cli tool
2 parents a9568e0 + cc1024d commit 1d41684

11 files changed

Lines changed: 963 additions & 89 deletions

‎cmd/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ func run(ctx context.Context) error {
8888
uiType := getEnv("UI_TYPE", "slack")
8989
sessionType := getEnv("SESSION_TYPE", "memory")
9090
modifyResources := parseModifyResourcesMode(getEnv("MODIFY_RESOURCES", "none"))
91+
enableAWSTool := getEnv("ENABLE_AWS_TOOL", "false") == "true"
9192

9293
klog.Infof("Starting kubeai-chatbot (version: %s, commit: %s, date: %s)", version, commit, date)
9394
klog.Infof("Configuration: provider=%s, model=%s, listen=%s", providerID, modelID, listenAddress)
@@ -160,6 +161,7 @@ func run(ctx context.Context) error {
160161
SessionBackend: sessionType,
161162
AgentName: agentName,
162163
ModifyResources: modifyResources,
164+
EnableAWSTool: enableAWSTool,
163165
SkillsRegistry: skillsRegistry,
164166
}, nil
165167
}

‎docs/aws_cli_tool.md‎

Lines changed: 317 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,317 @@
1+
# AWS command-line tool
2+
3+
The AWS command-line tool allows KubeAI to run `aws` commands against AWS resources on behalf of the user. It follows the same execution model as the kubectl tool — commands are parsed and executed directly without a shell, preventing injection attacks.
4+
5+
---
6+
7+
## What the AWS Tool Can Do
8+
9+
When enabled, the LLM can run `aws` commands to:
10+
11+
- Describe and list EC2 instances, security groups, VPCs, subnets
12+
- Query EKS clusters (`aws eks describe-cluster`, `aws eks list-clusters`)
13+
- Inspect load balancers (ALB/NLB via `aws elbv2`)
14+
- Check IAM roles and policies (`aws iam get-role`, `aws iam list-attached-role-policies`)
15+
- Query CloudWatch metrics and log groups
16+
- List S3 buckets and objects
17+
- Query RDS instances and snapshots
18+
- Inspect Route53 hosted zones and records
19+
- Run `aws sts get-caller-identity` to verify credentials
20+
21+
### Commands that are always blocked
22+
23+
Regardless of `ENABLE_AWS_TOOL`, the following are rejected at the validation layer:
24+
25+
| Blocked command | Reason |
26+
| -------------------------------------- | ---------------------------- |
27+
| `aws secretsmanager get-secret-value` | Secret retrieval |
28+
| `aws ssm get-parameter` | Secret retrieval |
29+
| `aws ssm get-parameters` | Secret retrieval |
30+
| `aws ssm get-parameters-by-path` | Secret retrieval |
31+
| `aws kms decrypt` | Credential/secret decryption |
32+
| `aws kms generate-data-key` | Credential/secret decryption |
33+
| `aws iam create-access-key` | Credential creation |
34+
| `aws sts assume-role` | Credential escalation |
35+
| Any compound command (`\|`, `&&`, `;`) | Shell injection prevention |
36+
| `aws configure` / `aws sso login` | Interactive mode |
37+
38+
---
39+
40+
## Enabling the AWS Tool
41+
42+
The tool is **disabled by default**. Set `ENABLE_AWS_TOOL=true` in the pod's environment:
43+
44+
```yaml
45+
# values.yaml
46+
env:
47+
ENABLE_AWS_TOOL: "true"
48+
```
49+
50+
No other configuration is required when the pod already has IRSA credentials — the tool inherits them from the process environment automatically.
51+
52+
---
53+
54+
## IAM Permissions
55+
56+
The IRSA role attached to the KubeAI service account must have the permissions you want the LLM to use. A minimal read-only policy example:
57+
58+
```json
59+
{
60+
"Version": "2012-10-17",
61+
"Statement": [
62+
{
63+
"Effect": "Allow",
64+
"Action": [
65+
"ec2:Describe*",
66+
"eks:DescribeCluster",
67+
"eks:ListClusters",
68+
"elasticloadbalancing:Describe*",
69+
"iam:GetRole",
70+
"iam:ListAttachedRolePolicies",
71+
"iam:ListRolePolicies",
72+
"iam:GetRolePolicy",
73+
"cloudwatch:GetMetricStatistics",
74+
"cloudwatch:ListMetrics",
75+
"logs:DescribeLogGroups",
76+
"logs:DescribeLogStreams",
77+
"s3:ListAllMyBuckets",
78+
"s3:ListBucket",
79+
"rds:DescribeDBInstances",
80+
"route53:ListHostedZones",
81+
"route53:ListResourceRecordSets",
82+
"sts:GetCallerIdentity"
83+
],
84+
"Resource": "*"
85+
}
86+
]
87+
}
88+
```
89+
90+
Attach this policy to the IRSA role used by the KubeAI service account. Follow the IRSA setup in [cross_cluster_access.md](cross_cluster_access.md) for the full role and trust policy setup.
91+
92+
---
93+
94+
## Cross-Account AWS Access
95+
96+
`aws sts assume-role` is blocked by the tool's validation layer, so cross-account access must be configured at the infrastructure level. There are two approaches depending on how many accounts you need.
97+
98+
---
99+
100+
### Option A — IRSA only (direct OIDC trust, no config file needed)
101+
102+
Account B registers Account A's OIDC provider directly in its own IAM. The pod then assumes Account B's role in a **single** `sts:AssumeRoleWithWebIdentity` call — no `~/.aws/config`, no role chaining. This is the approach described in the [AWS cross-account IRSA guide](https://docs.aws.amazon.com/eks/latest/userguide/cross-account-access.html).
103+
104+
Use this when KubeAI only needs to access Account B resources (the service account annotation points to Account B's role, so all commands run in Account B's context).
105+
106+
#### Step 1: Get Account A's OIDC issuer URL
107+
108+
Run this in **Account A**:
109+
110+
```bash
111+
OIDC_ISSUER=$(aws eks describe-cluster \
112+
--name cluster-a --region ap-southeast-1 \
113+
--query "cluster.identity.oidc.issuer" --output text)
114+
# e.g. https://oidc.eks.ap-southeast-1.amazonaws.com/id/EXAMPLED539D4633E53DE1B716D3041E
115+
```
116+
117+
#### Step 2: Register Account A's OIDC provider in Account B
118+
119+
Run this in **Account B**:
120+
121+
```bash
122+
# Get the OIDC thumbprint
123+
THUMBPRINT=$(openssl s_client -connect oidc.eks.ap-southeast-1.amazonaws.com:443 \
124+
-servername oidc.eks.ap-southeast-1.amazonaws.com 2>/dev/null \
125+
| openssl x509 -fingerprint -noout \
126+
| sed 's/SHA1 Fingerprint=//' | tr -d ':' | tr '[:upper:]' '[:lower:]')
127+
128+
aws iam create-open-id-connect-provider \
129+
--url $OIDC_ISSUER \
130+
--client-id-list sts.amazonaws.com \
131+
--thumbprint-list $THUMBPRINT
132+
# Output: arn:aws:iam::ACCOUNT_B_ID:oidc-provider/oidc.eks.ap-southeast-1.amazonaws.com/id/EXAMPLED539D4633E53DE1B716D3041E
133+
```
134+
135+
#### Step 3: Create a role in Account B that trusts Account A's OIDC provider
136+
137+
Run this in **Account B**, using the OIDC issuer path (without `https://`):
138+
139+
```bash
140+
OIDC_PROVIDER="oidc.eks.ap-southeast-1.amazonaws.com/id/EXAMPLED539D4633E53DE1B716D3041E"
141+
142+
cat > trust-policy.json <<EOF
143+
{
144+
"Version": "2012-10-17",
145+
"Statement": [
146+
{
147+
"Effect": "Allow",
148+
"Principal": {
149+
"Federated": "arn:aws:iam::ACCOUNT_B_ID:oidc-provider/${OIDC_PROVIDER}"
150+
},
151+
"Action": "sts:AssumeRoleWithWebIdentity",
152+
"Condition": {
153+
"StringEquals": {
154+
"${OIDC_PROVIDER}:sub": "system:serviceaccount:kubeai-chatbot:kubeai-chatbot",
155+
"${OIDC_PROVIDER}:aud": "sts.amazonaws.com"
156+
}
157+
}
158+
}
159+
]
160+
}
161+
EOF
162+
163+
aws iam create-role \
164+
--role-name kubeai-chatbot-cross-account \
165+
--assume-role-policy-document file://trust-policy.json
166+
167+
aws iam attach-role-policy \
168+
--role-name kubeai-chatbot-cross-account \
169+
--policy-arn arn:aws:iam::aws:policy/ReadOnlyAccess
170+
```
171+
172+
#### Step 4: Annotate the KubeAI service account with the Account B role
173+
174+
```yaml
175+
# values.yaml
176+
serviceAccount:
177+
annotations:
178+
eks.amazonaws.com/role-arn: "arn:aws:iam::ACCOUNT_B_ID:role/kubeai-chatbot-cross-account"
179+
```
180+
181+
All `aws` commands now run as the Account B role. No `~/.aws/config` needed.
182+
183+
---
184+
185+
### Option B — IRSA + aws config profiles (multiple accounts)
186+
187+
Use this when KubeAI needs access to **both** Account A and Account B. The IRSA role stays in Account A (default credentials). A named profile in `~/.aws/config` tells the AWS CLI how to assume the Account B role on demand — the SDK handles the `sts:AssumeRole` call transparently when `--profile account-b` is used.
188+
189+
#### Step 1: Create a cross-account role in Account B
190+
191+
```bash
192+
cat > cross-account-trust.json <<EOF
193+
{
194+
"Version": "2012-10-17",
195+
"Statement": [
196+
{
197+
"Effect": "Allow",
198+
"Principal": {
199+
"AWS": "arn:aws:iam::ACCOUNT_A_ID:role/kubeai-chatbot"
200+
},
201+
"Action": "sts:AssumeRole"
202+
}
203+
]
204+
}
205+
EOF
206+
207+
aws iam create-role \
208+
--role-name KubeAICrossAccountRole \
209+
--assume-role-policy-document file://cross-account-trust.json
210+
211+
aws iam attach-role-policy \
212+
--role-name KubeAICrossAccountRole \
213+
--policy-arn arn:aws:iam::aws:policy/ReadOnlyAccess
214+
```
215+
216+
#### Step 2: Allow the Account A IRSA role to assume the Account B role
217+
218+
Add `sts:AssumeRole` to the IRSA role policy in Account A:
219+
220+
```json
221+
{
222+
"Effect": "Allow",
223+
"Action": "sts:AssumeRole",
224+
"Resource": "arn:aws:iam::ACCOUNT_B_ID:role/KubeAICrossAccountRole"
225+
}
226+
```
227+
228+
#### Step 3: Mount `~/.aws/config` via ConfigMap
229+
230+
```yaml
231+
apiVersion: v1
232+
kind: ConfigMap
233+
metadata:
234+
name: aws-config
235+
namespace: kubeai-chatbot
236+
data:
237+
config: |
238+
[default]
239+
region = ap-southeast-1
240+
241+
[profile account-b]
242+
role_arn = arn:aws:iam::ACCOUNT_B_ID:role/KubeAICrossAccountRole
243+
credential_source = EcsContainer
244+
```
245+
246+
```yaml
247+
# values.yaml
248+
volumes:
249+
- name: aws-config
250+
configMap:
251+
name: aws-config
252+
253+
volumeMounts:
254+
- name: aws-config
255+
mountPath: /home/kubeai/.aws
256+
readOnly: true
257+
```
258+
259+
> `credential_source = EcsContainer` works on EKS because IRSA is compatible with the ECS container credential provider. Alternatively use `credential_source = Environment`.
260+
261+
Users instruct KubeAI to use Account B by specifying the profile in their message:
262+
263+
```bash
264+
Show me all EC2 instances in Account B using profile account-b
265+
```
266+
267+
The LLM appends `--profile account-b` to commands, e.g.:
268+
269+
```bash
270+
aws ec2 describe-instances --region ap-southeast-1 --profile account-b
271+
```
272+
273+
---
274+
275+
## Verify the Setup
276+
277+
Exec into the KubeAI pod and confirm the AWS CLI works:
278+
279+
```bash
280+
kubectl exec -it deployment/kubeai-chatbot -n kubeai-chatbot -- /bin/bash
281+
282+
# Confirm IRSA credentials are active
283+
aws sts get-caller-identity
284+
285+
# Test a read query
286+
aws eks list-clusters --region ap-southeast-1
287+
288+
# Test cross-account (if configured)
289+
aws sts get-caller-identity --profile account-b
290+
```
291+
292+
---
293+
294+
## Troubleshooting
295+
296+
**`ENABLE_AWS_TOOL` is set but the LLM does not use AWS commands**
297+
298+
- Confirm the env var value is exactly `"true"` (string, not boolean).
299+
- Restart the pod after changing env vars.
300+
301+
**`NoCredentialProviders` error**
302+
303+
- IRSA is not active. Check the service account annotation: `eks.amazonaws.com/role-arn`.
304+
- Verify the IRSA token is mounted: `ls /var/run/secrets/eks.amazonaws.com/serviceaccount/`.
305+
306+
**`AccessDenied` on a specific command**
307+
308+
- The IRSA role policy does not include the required action. Add it to the IAM policy attached to the IRSA role.
309+
310+
**`aws sts assume-role` rejected**
311+
312+
- This command is intentionally blocked by the tool's validation layer. Use Option A (IRSA pointing directly at Account B) or Option B (named profiles via `~/.aws/config`) instead — see [Cross-Account AWS Access](#cross-account-aws-access).
313+
314+
**Cross-account `AccessDenied`**
315+
316+
- **Option A**: Verify Account A's OIDC provider is registered in Account B IAM (`aws iam list-open-id-connect-providers` in Account B). Verify the Account B role's trust policy references `arn:aws:iam::ACCOUNT_B_ID:oidc-provider/...` (not Account A's ARN) with the correct `sub` condition (`system:serviceaccount:<namespace>:<service-account>`).
317+
- **Option B**: Verify the Account B role's trust policy allows `arn:aws:iam::ACCOUNT_A_ID:role/kubeai-chatbot` to assume it, and that the Account A IRSA role has `sts:AssumeRole` permission targeting the Account B role ARN.

‎docs/cross_cluster_access.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -444,6 +444,12 @@ kubectl exec -it deployment/kubeai-chatbot -n kubeai-chatbot -- aws eks describe
444444

445445
---
446446

447+
## AWS command-line tool
448+
449+
If you also want KubeAI to run `aws` commands (e.g. `aws eks describe-cluster`, `aws ec2 describe-instances`) in addition to `kubectl`, see [aws_cli_tool.md](aws_cli_tool.md). That guide covers enabling the tool, required IAM permissions, and cross-account AWS CLI access via named profiles.
450+
451+
---
452+
447453
## References
448454

449455
- [EKS IAM Roles for Service Accounts](https://docs.aws.amazon.com/eks/latestuserguide/iam-roles-for-service-accounts.html)

‎pkg/agent/agent_e2e_test.go‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -174,11 +174,7 @@ func TestAgentEndToEndToolExecution(t *testing.T) {
174174
t.Fatalf("run: %v", err)
175175
}
176176

177-
// Expect greeting and prompt inline (UI-driven startup)
178-
m1 := recvMsg(t, ctx, a.Output)
179-
if m1.Type != api.MessageTypeText || m1.Source != api.MessageSourceAgent {
180-
t.Fatalf("expected greeting text from agent, got type=%v source=%v", m1.Type, m1.Source)
181-
}
177+
// Expect prompt (UI-driven startup)
182178
m2 := recvMsg(t, ctx, a.Output)
183179
if m2.Type != api.MessageTypeUserInputRequest {
184180
t.Fatalf("expected user-input-request, got %v", m2.Type)
@@ -388,8 +384,7 @@ func TestAgentEndToEndAutomaticModifyDisabled(t *testing.T) {
388384
t.Fatalf("run: %v", err)
389385
}
390386

391-
// Skip greeting
392-
_ = recvMsg(t, ctx, a.Output)
387+
// Wait for user-input-request
393388
_ = recvMsg(t, ctx, a.Output)
394389

395390
// Send a query

0 commit comments

Comments
 (0)