Setup the necessary Athena components to analyze CloudTrail logs for incident response purposes

aws-setup-cloudtrail-investigation.sh

#!/bin/bash

# This script will create the necessary AWS Athena resources needed to conduct 
# an investigation using CloudTrail logs. Note that the S3 bucket defined as the
# Athena query results output location must already exist and will _not_ be 
# created by this script.

# Example: 
# ./aws-setup-cloudtrail-investigation.sh \
#   -c example-bucket-123/<PREFIX>/AWSLogs/<ACCOUNTID>/CloudTrail/<REGION> \
#   -w test-workgroup \
#   -d "this is a test" \
#   -a us-east-1 \
#   -r example-athena-output-bucket-456 \
#   -n cloudtrail_logs_test \
#   -s "2023/03/01"
#   -e "2023/03/08"

# Parse command-line arguments
while getopts ":c:w:d:a:r:n:s:e:" opt; do
  case ${opt} in
    c) CLOUDTRAIL_S3_PATH=$OPTARG;;
    w) ATHENA_WORKGROUP_NAME=$OPTARG;;
    d) ATHENA_WORKGROUP_DESCRIPTION=$OPTARG;;
    a) ATHENA_AWS_REGION=$OPTARG;;
    r) ATHENA_RESULTS_S3_PATH=$OPTARG;;
    n) ATHENA_TABLE_NAME=$OPTARG;;
    s) START_TIMESTAMP=$OPTARG;;
    e) END_TIMESTAMP=$OPTARG;;
    \?)
      echo "Invalid option: -$OPTARG" >&2
      exit 1
      ;;
    :)
      echo "Option -$OPTARG requires an argument." >&2
      exit 1
      ;;
  esac
done

# Check if all required arguments are provided
if [ -z "$CLOUDTRAIL_S3_PATH" ] || [ -z "$ATHENA_WORKGROUP_NAME" ] || [ -z "$ATHENA_RESULTS_S3_PATH" ] || [ -z "$ATHENA_WORKGROUP_DESCRIPTION" ] || [ -z "$START_TIMESTAMP" ] || [ -z "$ATHENA_TABLE_NAME" ] || [ -z "$ATHENA_AWS_REGION" ]; then
  echo "Usage: $0 -c <CLOUDTRAIL_S3_PATH> -w <ATHENA_WORKGROUP_NAME> -r <ATHENA_RESULTS_S3_PATH> -s <START_TIMESTAMP> [-e <END_TIMESTAMP>] -d <ATHENA_WORKGROUP_DESCRIPTION> -n <ATHENA_TABLE_NAME> -a <ATHENA_AWS_REGION>"
  exit 1
fi

if [ -z "$END_TIMESTAMP" ]; then
  END_TIMESTAMP="NOW"
fi

CREATE_TABLE_QUERY=$(cat <<EOF
  CREATE EXTERNAL TABLE $ATHENA_TABLE_NAME(
    eventVersion STRING,
    userIdentity STRUCT<
        type: STRING,
        principalId: STRING,
        arn: STRING,
        accountId: STRING,
        invokedBy: STRING,
        accessKeyId: STRING,
        userName: STRING,
        sessionContext: STRUCT<
            attributes: STRUCT<
                mfaAuthenticated: STRING,
                creationDate: STRING>,
            sessionIssuer: STRUCT<
                type: STRING,
                principalId: STRING,
                arn: STRING,
                accountId: STRING,
                userName: STRING>,
            ec2RoleDelivery:string,
            webIdFederationData:map<string,string>
        >
    >,
    eventTime STRING,
    eventSource STRING,
    eventName STRING,
    awsRegion STRING,
    sourceIpAddress STRING,
    userAgent STRING,
    errorCode STRING,
    errorMessage STRING,
    requestparameters STRING,
    responseelements STRING,
    additionaleventdata STRING,
    requestId STRING,
    eventId STRING,
    readOnly STRING,
    resources ARRAY<STRUCT<
        arn: STRING,
        accountId: STRING,
        type: STRING>>,
    eventType STRING,
    apiVersion STRING,
    recipientAccountId STRING,
    serviceEventDetails STRING,
    sharedEventID STRING,
    vpcendpointid STRING,
    tlsDetails struct<
        tlsVersion:string,
        cipherSuite:string,
        clientProvidedHostHeader:string>
  )
  PARTITIONED BY (
    \`timestamp\` string)
  ROW FORMAT SERDE 'org.apache.hive.hcatalog.data.JsonSerDe'
  STORED AS INPUTFORMAT 'com.amazon.emr.cloudtrail.CloudTrailInputFormat'
  OUTPUTFORMAT 'org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat'
  LOCATION
    's3://$CLOUDTRAIL_S3_PATH'
  TBLPROPERTIES (
    'projection.enabled'='true', 
    'projection.timestamp.format'='yyyy/MM/dd', 
    'projection.timestamp.interval'='1', 
    'projection.timestamp.interval.unit'='DAYS', 
    'projection.timestamp.range'='$START_TIMESTAMP,$END_TIMESTAMP', 
    'projection.timestamp.type'='date', 
    'storage.location.template'='s3://$CLOUDTRAIL_S3_PATH/\${timestamp}')
EOF
)

# Create Athena workgroup
echo "Creating Athena workgroup..."
aws athena create-work-group \
  --region $ATHENA_AWS_REGION \
  --name $ATHENA_WORKGROUP_NAME \
  --description "$ATHENA_WORKGROUP_DESCRIPTION" \
  --configuration 'ResultConfiguration={OutputLocation='"s3://$ATHENA_RESULTS_S3_PATH"'}' \
  --output json

if [ $? -eq 0 ]; then
    echo "Athena workgroup $ATHENA_WORKGROUP_NAME created!"
fi 

# Create partitioned Athena table
echo "Creating Athena table..."
aws athena start-query-execution \
  --region $ATHENA_AWS_REGION \
  --query-string "$CREATE_TABLE_QUERY" \
  --work-group $ATHENA_WORKGROUP_NAME 

if [ $? -eq 0 ]; then
    echo "Athena table created successfully!"
fi