From a6043157977bc7b6bd964d38255bba76d107b8f9 Mon Sep 17 00:00:00 2001
From: Boxuan Hu <boxuan@amazon.com>
Date: Fri, 7 Feb 2025 02:07:54 +0000
Subject: [PATCH] Replace User creds with Role creds for CWL agent

---
 bin/run-integ-tests.sh | 21 +++++++++++++++++++++
 bin/start-agent.sh     |  1 +
 tox.ini                |  1 +
 3 files changed, 23 insertions(+)

diff --git a/bin/run-integ-tests.sh b/bin/run-integ-tests.sh
index f8d5d98..51fb3de 100755
--- a/bin/run-integ-tests.sh
+++ b/bin/run-integ-tests.sh
@@ -18,6 +18,27 @@ status_code=0
 # Configure and start the agent
 ###################################
 
+# Check if IAM user credentials exist
+if [ -z "$AWS_ACCESS_KEY_ID" ] || [ -z "$AWS_SECRET_ACCESS_KEY" ]; then
+    echo "No IAM user credentials found, assuming we are running on CodeBuild pipeline, falling back to IAM role..."
+    
+    # Store the AWS STS assume-role output and extract credentials
+    CREDS=$(aws sts assume-role \
+        --role-arn $Code_Build_Execution_Role_ARN \
+        --role-session-name "session-$(uuidgen)" \
+        --query 'Credentials.[AccessKeyId,SecretAccessKey,SessionToken]' \
+        --output text \
+        --duration-seconds 3600)
+
+    # Parse the output into separate variables
+    read AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN <<< $CREDS
+
+    # Export the variables
+    export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
+else
+    echo "Using provided IAM user credentials..."
+fi
+
 $rootdir/bin/start-agent.sh
 
 ###################################
diff --git a/bin/start-agent.sh b/bin/start-agent.sh
index 0fb09a8..e4b0c7d 100755
--- a/bin/start-agent.sh
+++ b/bin/start-agent.sh
@@ -22,6 +22,7 @@ cd $rootdir/tests/integ/agent
 echo "[AmazonCloudWatchAgent]
 aws_access_key_id = $AWS_ACCESS_KEY_ID
 aws_secret_access_key = $AWS_SECRET_ACCESS_KEY
+aws_session_token = $AWS_SESSION_TOKEN
 " > ./.aws/credentials
 
 echo "[profile AmazonCloudWatchAgent]
diff --git a/tox.ini b/tox.ini
index fe51c57..a8a60fc 100644
--- a/tox.ini
+++ b/tox.ini
@@ -23,6 +23,7 @@ passenv =
     AWS_REGION
     AWS_ACCESS_KEY_ID
     AWS_SECRET_ACCESS_KEY
+    AWS_SESSION_TOKEN
 
 [testenv:flake8]
 basepython = python3.7