# 1️⃣ Create a virtual environment (recommended)
python3 -m venv .venv
source .venv/bin/activate   # on Windows: .venv\Scripts\activate

# 2️⃣ Install core dependencies
pip install --upgrade pip
pip install boto3==1.34.0   # AWS SDK for Python
pip install python-dotenv==1.0.0   # load .env files
# Optional: if you plan to use a real LLM (e.g., OpenAI)
# pip install openai==1.30.0

# 1️⃣ Initialize a new Node.js project
mkdir agentcorruption-defender && cd $_
npm init -y   # or yarn init -y

# 2️⃣ Install AWS SDK v3 (modular) and dotenv
npm install @aws-sdk/client-s3 @aws-sdk/credential-providers dotenv
# Optional: TypeScript support
npm install --save-dev typescript ts-node @types/node

# 3️⃣ (TS only) Create a basic tsconfig.json
npx tsc --init --rootDir src --outDir dist --esModuleInterop --resolveJsonModule --lib es2020,dom --module commonjs

# file: agent.py
"""
Defensive LLM‑to‑AWS agent – mitigates AgentCorruption prompt‑injection.
"""

import os
import json
import logging
from typing import Literal

import boto3
from botocore.exceptions import BotoCoreError, ClientError
from dotenv import load_dotenv

# ----------------------------------------------------------------------
# Configuration & Logging
# ----------------------------------------------------------------------
load_dotenv()  # pulls AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION, etc.
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s - %(message)s",
)
logger = logging.getLogger("agentcorruption")

# ----------------------------------------------------------------------
# Constants – adjust to your least‑privilege policy
# ----------------------------------------------------------------------
ALLOWED_ACTIONS = {"list_s3_buckets"}          # whitelist of safe operations
ROLE_TO_ASSUME = os.getenv(
    "AGENT_ASSUME_ROLE_ARN",
    "arn:aws:iam::123456789012:role/AgentCorruptionDefenderRole",
)
SESSION_NAME = "AgentCorruptionSession"

# ----------------------------------------------------------------------
# Helper: Assume a least‑privilege role via STS
# ----------------------------------------------------------------------
def assume_role(role_arn: str, session_name: str) -> dict:
    """
    Returns temporary credentials dict.
    Raises RuntimeError on failure.
    """
    sts = boto3.client("sts")
    try:
        resp = sts.assume_role(
            RoleArn=role_arn,
            RoleSessionName=session_name,
            # Optional: ExternalID for added security if your trust policy requires it
            # ExternalId=os.getenv("EXTERNAL_ID"),
        )
        creds = resp["Credentials"]
        logger.info("Assumed role %s successfully", role_arn)
        return {
            "aws_access_key_id": creds["AccessKeyId"],
            "aws_secret_access_key": creds["SecretAccessKey"],
            "aws_session_token": creds["SessionToken"],
            "region_name": os.getenv("AWS_DEFAULT_REGION", "us-east-1"),
        }
    except (BotoCoreError, ClientError) as e:
        logger.exception("Failed to assume role %s", role_arn)
        raise RuntimeError(f"STS assume_role error: {e}") from e

# ----------------------------------------------------------------------
# Mock LLM – replace with real provider call
# ----------------------------------------------------------------------
def get_llm_response(user_prompt: str) -> str:
    """
    In a real implementation, call your LLM (OpenAI, Anthropic, etc.)
    and return the raw text.
    For demo purposes we simply echo the prompt.
    """
    # TODO: insert actual LLM SDK call here
    return user_prompt.strip()

# ----------------------------------------------------------------------
# Core agent logic
# ----------------------------------------------------------------------
def safe_agent(user_prompt: str) -> dict:
    """
    1. Validate prompt → allowed action.
    2. Assume least‑privilege role.
    3. Execute the action via boto3.
    4. Return structured result or error.
    """
    # ---- Step 1: Prompt validation -------------------------------------------------
    # We only allow a very small set of deterministic commands.
    # In practice you could parse the LLM output into a JSON schema.
    normalized = user_prompt.lower().strip()
    if normalized not in ALLOWED_ACTIONS:
        msg = f"Prompt '{user_prompt}' not in allowed actions {ALLOWED_ACTIONS}"
        logger.warning(msg)
        return {"status": "error", "reason": "prompt_not_allowed", "message": msg}

    action = normalized  # e.g., "list_s3_buckets"

    # ---- Step 2: Assume role -------------------------------------------------------
    try:
        creds = assume_role(ROLE_TO_ASSUME, SESSION_NAME)
    except RuntimeError as e:
        return {"status": "error", "reason": "sts_failure", "message": str(e)}

    # ---- Step 3: Execute AWS call --------------------------------------------------
    try:
        if action == "list_s3_buckets":
            s3 = boto3.client("s3", **creds)
            response = s3.list_buckets()
            bucket_names = [b["Name"] for b in response.get("Buckets", [])]
            logger.info("Listed %d S3 buckets", len(bucket_names))
            return {"status": "success", "action": action, "buckets": bucket_names}
        else:
            # Should never happen because of the whitelist, but keep for safety.
            raise ValueError(f"Unsupported action: {action}")
    except (BotoCoreError, ClientError) as e:
        logger.exception("AWS call failed for action %s", action)
        return {"status": "error", "reason": "aws_failure", "message": str(e)}
    except Exception as e:  # catch‑all for unexpected bugs
        logger.exception("Unexpected error")
        return {"status": "error", "reason": "unexpected", "message": str(e)}


# ----------------------------------------------------------------------
# Simple CLI for local testing
# ----------------------------------------------------------------------
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python agent.py \"<user prompt>\"")
        sys.exit(1)

    prompt = sys.argv[1]
    result = safe_agent(prompt)
    print(json.dumps(result, indent=2))

# Ensure .env contains AWS credentials for the *initial* IAM user/role
# (the one that can sts:AssumeRole the defender role)
export AWS_ACCESS_KEY_ID=AKIA...
export AWS_SECRET_ACCESS_KEY=...
export AWS_DEFAULT_REGION=us-east-1
export AGENT_ASSUME_ROLE_ARN=arn:aws:iam::123456789012:role/AgentCorruptionDefenderRole

python agent.py "list_s3_buckets"
# Expected output: JSON with status: success and bucket list
python agent.py "drop all tables"
# Expected output: JSON with status: error, reason: prompt_not_allowed

// file: src/agent.ts
/**
 * Defensive LLM‑to‑AWS agent (Node.js) – mitigates AgentCorruption.
 * Uses AWS SDK v3 (modular) and dotenv for config.
 */

import { config } from "dotenv";
import { STSClient, AssumeRoleCommand, AssumeRoleCommandOutput } from "@aws-sdk/client-sts";
import { S3Client, ListBucketsCommand, ListBucketsCommandOutput } from "@aws-sdk/client-s3";
import { Logger } from "./logger"; // simple wrapper (see below)

config(); // loads .env into process.env

// ----------------------------------------------------------------------
// Logger (tiny wrapper around console)
// ----------------------------------------------------------------------
export class Logger {
  static info(msg: string) {
    console.info(`[INFO] ${new Date().toISOString()} - ${msg}`);
  }
  static warn(msg: string) {
    console.warn(`[WARN] ${new Date().toISOString()} - ${msg}`);
  }
  static error(msg: string, err?: unknown) {
    console.error(`[ERROR] ${new Date().toISOString()} - ${msg}`, err);
  }
}

// ----------------------------------------------------------------------
// Constants – adjust to your least‑privilege policy
// ----------------------------------------------------------------------
const ALLOWED_ACTIONS = new Set(["list_s3_buckets"]);
const ROLE_TO_ASSUME =
  process.env.AGENT_ASSUME_ROLE_ARN ||
  "arn:aws:iam::123456789012:role/AgentCorruptionDefenderRole";
const SESSION_NAME = "AgentCorruptionSession";

// ----------------------------------------------------------------------
// Helper: Assume role via STS
// ----------------------------------------------------------------------
async function assumeRole(roleArn: string): Promise<AssumeRoleCommandOutput> {
  const sts = new STSClient({});
  const params = {
    RoleArn: roleArn,
    RoleSessionName: SESSION_NAME,
    // ExternalId: process.env.EXTERNAL_ID, // if needed
  };
  try {
    const data = await sts.send(new AssumeRoleCommand(params));
    Logger.info(`Assumed role ${roleArn}`);
    return data;
  } catch (err) {
    Logger.error(`Failed to assume role ${roleArn}`, err);
    throw new Error(`STS assumeRole failed: ${err instanceof Error ? err.message : String(err)}`);
  }
}

// ----------------------------------------------------------------------
// Mock LLM – replace with real provider call
// ----------------------------------------------------------------------
function mockLLM(prompt: string): string {
  // In production, call your LLM SDK here and return the raw text.
  return prompt.trim();
}

// ----------------------------------------------------------------------
// Core agent logic
// ----------------------------------------------------------------------
export async function safeAgent(userPrompt: string): Promise<any> {
  // 1️⃣ Validate prompt
  const normalized = userPrompt.toLowerCase().trim();
  if (!ALLOWED_ACTIONS.has(normalized)) {
    const msg = `Prompt "${userPrompt}" not in allowed actions ${Array.from(
      ALLOWED_ACTIONS
    ).join(", ")}`;
    Logger.warn(msg);
    return { status: "error", reason: "prompt_not_allowed", message: msg };
  }
  const action = normalized as "list_s3_buckets";

  // 2️⃣ Assume least‑privilege role
  let creds: any;
  try {
    const assumeResp = await assumeRole(ROLE_TO_ASSUME);
    creds = {
      accessKeyId: assumeResp.Credentials?.AccessKeyId,
      secretAccessKey: assumeResp.Credentials?.SecretAccessKey,
      sessionToken: assumeResp.Credentials?.SessionToken,
      region: process.env.AWS_DEFAULT_REGION ?? "us-east-1",
    };
  } catch (err) {
    return { status: "error", reason: "sts_failure", message: String(err) };
  }

  // 3️⃣ Execute AWS call
  try {
    if (action === "list_s3_buckets") {
      const s3 = new S3Client(creds);
      const data = await s3.send(new ListBucketsCommand({}));
      const bucketNames = data.Buckets?.map((b) => b.Name ?? "") ?? [];
      Logger.info(`Listed ${bucketNames.length} S3 buckets`);
      return { status: "success", action, buckets: bucketNames };
    } else {
      // Should never happen because of the whitelist
      throw new Error(`Unsupported action: ${action}`);
    }
  } catch (err) {
    Logger.error(`AWS call failed for action ${action}`, err);
    return { status: "error", reason: "aws_failure", message: String(err) };
  }
}

// ----------------------------------------------------------------------
// Simple CLI for local testing (ts-node)
// ----------------------------------------------------------------------
if (require.main === module) {
  const arg = process.argv[2];
  if (!arg) {
    console.error('Usage: ts-node src/agent.ts "<user prompt>"');
    process.exit(1);
  }
  safeAgent(arg)
    .then((res) => console.log(JSON.stringify(res, null, 2)))
    .catch((e) => {
      console.error("Unexpected error:", e);
      process.exit(1);
    });
}

# 1️⃣ Create .env (never commit this file)
cat > .env <<EOF
AWS_ACCESS_KEY_ID=AKIA...
AWS_SECRET_ACCESS_KEY=...
AWS_DEFAULT_REGION=us-east-1
AGENT_ASSUME_ROLE_ARN=arn:aws:iam::123456789012:role/AgentCorruptionDefenderRole
EOF

# 2️⃣ Install dependencies (if not done already)
npm install

# 3️⃣ Run with ts-node (or compile first)
npx ts-node src/agent.ts "list_s3_buckets"
# Expected: JSON with status: success and bucket list
npx ts-node src/agent.ts "delete everything"
# Expected: JSON with status: error, reason: prompt_not_allowed

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": ["sts:AssumeRole"],
      "Resource": "arn:aws:iam::123456789012:role/AgentCorruptionDefenderRole"
    },
    {
      "Effect": "Allow",
      "Action": ["s3:ListBucket", "s3:GetObject"],
      "Resource": [
        "arn:aws:s3:::my-public-bucket",
        "arn:aws:s3:::my-public-bucket/*"
      ]
    }
  ]
}