Add script to check for security update of given ami

check-update-security.sh

@@ -0,0 +1,199 @@
+#!/usr/bin/env bash
+
+set -eio pipefail
+
+usage() {
+    echo "Usage:"
+    echo "  $0 AMI_PLATFORM"
+    echo "Example:"
+    echo "  $0 al2_arm"
+    echo "AMI_PLATFORM Must be one of: al, al2, al2_arm, al2023, al2023_arm"
+}
+
+error_msg() {
+    local msg="$1"
+    echo "ERROR: $msg"
+}
+
+#paths to get the ami ids from ssm params
+AL1_PATH="/aws/service/ecs/optimized-ami/amazon-linux/recommended"
+AL2_PATH="/aws/service/ecs/optimized-ami/amazon-linux-2/recommended"
+AL2_ARM_PATH="/aws/service/ecs/optimized-ami/amazon-linux-2/arm64/recommended"
+AL2023_PATH="/aws/service/ecs/optimized-ami/amazon-linux-2023/recommended"
+AL2023_ARM_PATH="/aws/service/ecs/optimized-ami/amazon-linux-2023/arm64/recommended"
+
+#Indicates that an update exists
+UPDATE_EXISTS_CODE="100"
+#Indicates that a wait operation failed
+WAIT_FAIL_CODE="255"
+#Indicates success
+SUCCESS_CODE="0"
+
+#in case of failure, terminate instance
+failure_cleanup() {
+    terminate_out=$(aws ec2 terminate-instances --instance-ids $instance_id)
+}
+
+check_wait_response() {
+    local wait_response=$1
+    if [ "$wait_response" -eq "$WAIT_FAIL_CODE" ]; then
+        error_msg "Failed to launch instance, instance timeout"
+        exit 1
+    fi
+}
+
+platform=$1
+if [ -z "$platform" ]; then
+    error_msg "Must specify an AMI platform"
+    usage
+    exit 1
+fi
+if [ -z "$IAM_INSTANCE_PROFILE_ARN" ]; then
+    error_msg "IAM_INSTANCE_PROFILE_ARN environment variable must exist"
+    exit 1
+fi
+
+#get ecs-optimized ami's PATH in ssm parameters
+platform=$1
+skip_user_data=0
+instance_type="c5.large"
+case "$platform" in
+"al1")
+    ami_path=$AL1_PATH
+    instance_type="t3.small"
+    ;;
+"al2")
+    ami_path=$AL2_PATH
+    ;;
+"al2_arm")
+    ami_path=$AL2_ARM_PATH
+    instance_type="c6g.medium"
+    ;;
+"al2023")
+    ami_path=$AL2023_PATH
+    skip_user_data=1
+    ;;
+"al2023_arm")
+    ami_path=$AL2023_ARM_PATH
+    skip_user_data=1
+    instance_type="c6g.medium"
+    ;;
+*)
+    error_msg "Incorrect platform selection"
+    usage
+    exit 1
+    ;;
+esac
+
+#Query ssm to get latest ecs optimized ami
+ami_id=$(aws ssm get-parameters --names $ami_path --region us-west-2 | jq -r '.Parameters[0].Value' | jq -r '.image_id')
+
+user_data=$(mktemp user_data.txt)
+echo "#cloud-config" >>user_data.txt
+echo "repo_upgrade: none" >>user_data.txt
+
+#Launch ec2 instance with given ami and SSM access for command execution
+#Also get instance id
+if [ $skip_user_data -eq 0 ]; then
+    #modify user data to ignore automatic updates by al and al2
+    instance_id=$(aws ec2 run-instances \
+        --image-id $ami_id \
+        --instance-type $instance_type \
+        --iam-instance-profile Arn=$IAM_INSTANCE_PROFILE_ARN \
+        --user-data file://user_data.txt |
+        jq -r '.Instances[0].InstanceId')
+    command_params='commands=["yum check-update --security --sec-severity=critical  --exclude=nvidia*,docker*,cuda*,containerd* -q"]'
+else
+    #in the case that we
+    instance_id=$(aws ec2 run-instances \
+        --image-id $ami_id \
+        --instance-type $instance_type \
+        --iam-instance-profile Arn=$IAM_INSTANCE_PROFILE_ARN |
+        jq -r '.Instances[0].InstanceId')
+    distribution_release_al2023=$(curl -s https://al2023-repos-us-west-2-de612dc2.s3.dualstack.us-west-2.amazonaws.com/core/releasemd.xml |
+        xmllint --xpath "string(//root/releases/release[last()]/@version)" -)
+    command_params='commands=["sudo dnf check-update --releasever='$distribution_release_al2023' --security --sec-severity=Critical --exclude=nvidia*,docker*,cuda*,containerd* -q"]'
+fi
+
+rm "$user_data"
+
+#Wait for instance status to reach ok, fail at timeout code
+aws ec2 wait instance-running --instance-ids $instance_id
+check_wait_response $(echo $?)
+
+#Instance has been launched, terminate in case of an error
+trap 'failure_cleanup' ERR
+
+#Assert that ssm agent is running before moving forward
+ssm_agent_status() {
+    aws ssm describe-instance-information \
+        --instance-information-filter-list key=InstanceIds,valueSet=$instance_id \
+        --query 'InstanceInformationList[0].PingStatus' --output text
+}
+max_retries=10
+success=0
+for ((r = 0; r < max_retries; r++)); do
+    if [ "$(ssm_agent_status)" = "Online" ]; then
+        success=1
+        break
+    fi
+    sleep 10
+done
+if [ $success -ne 1 ]; then
+    echo "SSM Agent connection timed out"
+    failure_cleanup
+    exit 1
+fi
+
+#Send command
+cmd_id=$(aws ssm send-command \
+    --document-name 'AWS-RunShellScript' \
+    --parameters "$command_params" \
+    --targets Key=instanceids,Values=$instance_id \
+    --comment "run security check" |
+    jq -r '.Command.CommandId')
+
+#Wait for command to be executed
+command_status() {
+    aws ssm get-command-invocation \
+        --command-id $cmd_id \
+        --instance-id $instance_id \
+        --query 'Status' \
+        --output text
+}
+max_retries=20
+success=0
+for ((r = 0; r < max_retries; r++)); do
+    cmd_status=$(command_status)
+    if [ "$cmd_status" = "Failed" ] || [ "$cmd_status" = "Success" ]; then
+        success=1
+        break
+    fi
+    sleep 5
+done
+if [ $success -ne 1 ]; then
+    echo "Command execution timed out"
+    failure_cleanup
+    exit 1
+fi
+
+#Get command output
+cmd_response_code=$(aws ssm get-command-invocation \
+    --command-id $cmd_id \
+    --instance-id $instance_id |
+    jq -r '.ResponseCode')
+
+#Delete the instance
+terminate_out=$(aws ec2 terminate-instances --instance-ids $instance_id)
+
+#Return whether update is necessary
+if [ "$cmd_response_code" -eq "$UPDATE_EXISTS_CODE" ]; then
+    echo "true"
+elif [ "$cmd_response_code" -ne "$SUCCESS_CODE" ]; then
+    #if update doesn't exist and there was a fail code, something went wrong
+    echo "Unknown issue with the command execution"
+else
+    echo "false"
+fi
+
+exit 0