#!/bin/bash -e

# This file was copied from
# https://code.amazon.com/packages/SinglePassEC2CloudFormationScript/blobs/mainline/--/configuration/usr/singlepass/prod.sh

set -o pipefail

while getopts :o:h:b:r:l:p:c:x:adf OPTION
do
    case ${OPTION}
    in
    a) ANYWHERE=YES;;
    d) DIRECT=YES;;
    f) STRICT_FAIL=YES;;
    p) PROFILE_ARGUMENT=${OPTARG};;
    l) LDAP_ARGUMENT=${OPTARG};;
    h) HOSTCLASS_ARGUMENT=${OPTARG};;
    b) BUCKET_NAME_ARGUMENT=${OPTARG};;
    r) ROLE_ARN_ARGUMENT=${OPTARG};;
    c) CREDENTIALS_FILE_ARGUMENT=${OPTARG};;
    o) OVERRIDE_REGION=${OPTARG};;
    x) USE_PROXY=${OPTARG};;
    ?) echo "SPIE Bootstrapping Script - Argument '${OPTARG}' unknown. Aborting script."; exit 1;;
    esac
done

get_fallback_region_for_partition () {
    local partition="$1"
    case "${partition}"
    in
    aws)
        echo "us-east-1";;
    aws-cn)
        echo "cn-north-1";;
    aws-iso)
        echo "us-iso-east-1";;
    aws-iso-b)
        echo "us-isob-east-1";;
    aws-us-gov)
        echo "us-gov-west-1";;
    *)
        return 1;;
    esac
}

get_sts_endpoint () {
    local partition="$1" region="$2"

    local sts_domain
    case "${partition}"
    in
    aws)
        sts_domain="amazonaws.com";;
    aws-cn)
        sts_domain="amazonaws.com.cn";;
    aws-iso)
        sts_domain="c2s.ic.gov";;
    aws-iso-b)
        sts_domain="sc2s.sgov.gov";;
    aws-us-gov)
        sts_domain="amazonaws.com";;
    *)
        return 1;;
    esac

    if [ -z "${region}" ]; then
        region=$(get_fallback_region_for_partition "${partition}")
    fi
    echo "https://sts.${region}.${sts_domain}/"
}

ERROR_MESSAGES=""
if [ -n "$ANYWHERE" ]; then
    if [ -z "$HOSTCLASS_ARGUMENT" ]; then
        ERROR_MESSAGES+="\nSet for SinglePassAnywhere, but Hostclass not provided. Please use argument -h [hostclass name]"
    fi
    if [ -z "$BUCKET_NAME_ARGUMENT" ]; then
        ERROR_MESSAGES+="\nSet for SinglePassAnywhere, but Bucket Name not provided. Please use argument -b [bucket name]"
    fi
    if [ -z "$DIRECT" -a -z "$ROLE_ARN_ARGUMENT" ]; then
        ERROR_MESSAGES+="\nSet for SinglePassAnywhere, but neither Direct flag nor Role ARN provided. Please set to Direct (-d) or provide a Role ARN (-r [role arn])"
    fi
    case ${LDAP_ARGUMENT}
    in
    BJS)
        PARTITION=aws-cn;
        REGION=cn-north-1;;
    PDT)
        PARTITION=aws-us-gov;
        REGION=us-gov-west-1;;
    DCA)
        PARTITION=aws-iso;
        REGION=us-iso-east-1;;
    LCK)
        PARTITION=aws-iso-b;
        REGION=us-isob-east-1;;
    IAD)
        PARTITION=aws;;
    *)
        ERROR_MESSAGES+="\nSet for SinglePassAnywhere, but LDAP source is not provided or not valid. Please use argument -l [IAD/BJS/PDT/DCA/LCK] to specify which LDAP source should be used for credentials";;
    esac
    ACCOUNT_ID=IRRELEVANT
    if [ -n "$OVERRIDE_REGION" ] ; then
        REGION=$OVERRIDE_REGION
    fi
else
    if [ -n "$PROFILE_ARGUMENT" ]; then
        ERROR_MESSAGES+="\nAWS CLI Profile argument (-p) cannot be used without setting Anywhere (-a). If running on an EC2 Instance, please remove AWS CLI Profile (-p) argument. If running on a non-EC2 Instance, please set to Anywhere (-a)"
    fi

    if [ -n "$CREDENTIALS_FILE_ARGUMENT" ]; then
        ERROR_MESSAGES+="\nAWS CLI Credentials File argument (-c) cannot be used without setting Anywhere (-a). If running on an EC2 Instance, please remove AWS CLI Credentials File (-c) argument. If running on a non-EC2 Instance, please set to Anywhere (-a)"
    fi

    METADATA_SERVICE='http://169.254.169.254/latest'
    IDENTITY_DOCUMENT="${METADATA_SERVICE}/dynamic/instance-identity/document"
    PARTITION_REQ="${METADATA_SERVICE}/meta-data/services/partition"

    TOKEN=$(curl --silent --fail --connect-timeout 5 --max-time 5 --retry 3 -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 600") || $(echo "")
    if [ -z "$TOKEN" ]; then
        METADATA=$(curl --silent --fail --connect-timeout 5 --max-time 5 --retry 3 $IDENTITY_DOCUMENT)
        PARTITION=$(curl --silent --fail --connect-timeout 5 --max-time 5 --retry 3 $PARTITION_REQ)
    else
        METADATA=$(curl --silent --fail --connect-timeout 5 --max-time 5 --retry 3 $IDENTITY_DOCUMENT -H "X-aws-ec2-metadata-token: $TOKEN")
        PARTITION=$(curl --silent --fail --connect-timeout 5 --max-time 5 --retry 3 $PARTITION_REQ -H "X-aws-ec2-metadata-token: $TOKEN")
    fi

    METADATA_LINES=$(echo $METADATA | /usr/bin/tr -d '{[]" }')
    ACCOUNT_ID=$(echo $METADATA_LINES | grep -oP "(?<=accountId:)(.+?)(?=([,]|$)+)")
    if [ -z "$ACCOUNT_ID" ]; then
        ERROR_MESSAGES+="\nFailed to get ACCOUNT ID from Metadata Service"
    fi
    REGION=$(echo $METADATA_LINES | grep -oP "(?<=region:)(.+?)(?=([,]|$)+)")
    if [ -z "$REGION" ]; then
        ERROR_MESSAGES+="\nFailed to get REGION from Metadata Service"
    fi
    if [ -z "$PARTITION" ]; then
        ERROR_MESSAGES+="\nFailed to get PARTITION from Metadata Service"
    fi
fi

if ! STS_ENDPOINT=$(get_sts_endpoint "${PARTITION}" "${REGION}"); then
    ERROR_MESSAGES+="\nCalled get_sts_endpoint with unknown partition '${PARTITION}'";
fi

if [ -n "$ERROR_MESSAGES" ]; then
    printf "$ERROR_MESSAGES\n"
    exit 1
fi

if [ -n "$CREDENTIALS_FILE_ARGUMENT" ]; then
    export AWS_SHARED_CREDENTIALS_FILE=$CREDENTIALS_FILE_ARGUMENT
fi

if [ -n "$BUCKET_NAME_ARGUMENT" ]; then
    BUCKET=$BUCKET_NAME_ARGUMENT
else
    BUCKET=singlepass-$ACCOUNT_ID-client-data
fi

if [ -n "$REGION" ]; then
    REGION_CLI_ARGUMENT="--region $REGION"
else
    REGION_CLI_ARGUMENT=""
fi

if [ -n "$PROFILE_ARGUMENT" ]; then
    PROFILE_CLI_ARGUMENT="--profile $PROFILE_ARGUMENT"
else
    PROFILE_CLI_ARGUMENT=""
fi

STS_ENDPOINT_CLI_ARGUMENT="--endpoint-url ${STS_ENDPOINT}"

if [ -z "$DIRECT" ]; then
    if [ -n "$ROLE_ARN_ARGUMENT" ]; then
        ROLE_ARN=$ROLE_ARN_ARGUMENT
    else
        ROLE_ARN=arn:$PARTITION:iam::$ACCOUNT_ID:role/singlepass-sync-reader
    fi

    ROLE_ARN_CLI_ARGUMENT="--role-arn $ROLE_ARN"
    ROLE_SESSION_NAME_CLI_ARGUMENT="--role-session-name spie-bootstrap-r1"

    if [ "$PARTITION" = "aws" -a -z "$ANYWHERE" ]; then
        EXTERNAL_ID_CLI_ARGUMENT="--external-id eb0054ab-e9aa-4605-9a81-23c40a20d920"
    else
        EXTERNAL_ID_CLI_ARGUMENT=""
    fi

    if ! ROLE_RESPONSE=$(/usr/bin/aws sts assume-role $ROLE_ARN_CLI_ARGUMENT $ROLE_SESSION_NAME_CLI_ARGUMENT $REGION_CLI_ARGUMENT $STS_ENDPOINT_CLI_ARGUMENT $EXTERNAL_ID_CLI_ARGUMENT --query 'Credentials.[AccessKeyId, SecretAccessKey, SessionToken]' | /usr/bin/tr -d '[]" ,' | /usr/bin/tr '\n' ' '); then
        FALLBACK_REGION_CLI_ARGUMENT="--region $(get_fallback_region_for_partition "${PARTITION}")"
        FALLBACK_STS_ENDPOINT_CLI_ARGUMENT="--endpoint-url $(get_sts_endpoint "${PARTITION}")"
        ROLE_RESPONSE=$(/usr/bin/aws sts assume-role $ROLE_ARN_CLI_ARGUMENT $ROLE_SESSION_NAME_CLI_ARGUMENT $FALLBACK_REGION_CLI_ARGUMENT $FALLBACK_STS_ENDPOINT_CLI_ARGUMENT $EXTERNAL_ID_CLI_ARGUMENT --query 'Credentials.[AccessKeyId, SecretAccessKey, SessionToken]' | /usr/bin/tr -d '[]" ,' | /usr/bin/tr '\n' ' ');
    fi
    
    export AWS_ACCESS_KEY_ID=$(echo $ROLE_RESPONSE | /bin/cut --delimiter=' ' -f1)
    export AWS_SECRET_ACCESS_KEY=$(echo $ROLE_RESPONSE | /bin/cut --delimiter=' ' -f2)
    export AWS_SESSION_TOKEN=$(echo $ROLE_RESPONSE | /bin/cut --delimiter=' ' -f3)
fi

UPDATER_RPM=singlepass-ec2-client-self-updater-bootstrap.noarch.rpm
UPDATER_SRC=s3://$BUCKET/client/$UPDATER_RPM
UPDATER_DST=/usr/singlepass/$UPDATER_RPM

if RESULT=$(/usr/bin/aws s3 cp $UPDATER_SRC $UPDATER_DST $REGION_CLI_ARGUMENT $PROFILE_CLI_ARGUMENT); then
    echo $RESULT
else
    /usr/bin/aws s3 cp $UPDATER_SRC $UPDATER_DST $PROFILE_CLI_ARGUMENT
fi

/usr/bin/yum -y localinstall $UPDATER_DST

if [ -z "$DIRECT" ]; then
    unset AWS_ACCESS_KEY_ID
    unset AWS_SECRET_ACCESS_KEY
    unset AWS_SESSION_TOKEN
fi

if [ -n "$CREDENTIALS_FILE_ARGUMENT" ]; then
    unset AWS_SHARED_CREDENTIALS_FILE
fi

if [ -n "$HOSTCLASS_ARGUMENT" ]; then
    echo $HOSTCLASS_ARGUMENT > /etc/hostclass
fi

mkdir -p /usr/singlepass/etc
echo -n "" > /usr/singlepass/etc/override.properties

if [ -n "$USE_PROXY" ]; then
      echo "proxy_info = $USE_PROXY" >> /usr/singlepass/etc/override.properties
      echo "Selected USE_PROXY"
      echo "Proxy URL : $USE_PROXY"
fi
if [ -n "$STRICT_FAIL" ]; then
      echo "strict_fail = 1" >> /usr/singlepass/etc/override.properties
fi
if [ -n "$BUCKET_NAME_ARGUMENT" ]; then
    echo "s3_bucket = $BUCKET_NAME_ARGUMENT" >> /usr/singlepass/etc/override.properties
fi
if [ -n "$ROLE_ARN_ARGUMENT" ]; then
    echo "role_arn = $ROLE_ARN_ARGUMENT" >> /usr/singlepass/etc/override.properties
fi
if [ -n "$DIRECT" ]; then
    echo "use_sts = 0" >> /usr/singlepass/etc/override.properties
    echo "use_sts = 0" >> /usr/singlepass/updater/etc/config.properties
fi
if [ -n "$ANYWHERE" ]; then
    echo "anywhere = 1" >> /usr/singlepass/etc/override.properties
    echo "partition = $PARTITION" >> /usr/singlepass/etc/override.properties
    if [ -n "$REGION" ]; then
        echo "region = $REGION" >> /usr/singlepass/etc/override.properties
    fi
    if [ -n "$PROFILE_ARGUMENT" ]; then
        echo "aws_cli_profile = $PROFILE_ARGUMENT" >> /usr/singlepass/etc/override.properties
    fi
    if [ -n "$CREDENTIALS_FILE_ARGUMENT" ]; then
        echo "aws_cli_credentials_file = $CREDENTIALS_FILE_ARGUMENT" >> /usr/singlepass/etc/override.properties
    fi
fi

# Run the SPIE Client Updater to install the SPIE Client
/usr/singlepass/updater/bin/auth-sync-ec2-self-update.pl

# temporary workaround for AL1 issue in DCA/LCK: https://tt.amazon.com/0555321104
# can be removed once we get rid of all AL1 instances
if [ $PARTITION = aws-iso ] || [ $PARTITION = aws-iso-b ]; then
    echo -e "AWS_CA_BUNDLE=$AWS_CA_BUNDLE\n$(cat /etc/cron.d/singlepass_ec2_client_cron)" > /etc/cron.d/singlepass_ec2_client_cron
fi

# Ensure a successful run of the SPIE Client
/usr/singlepass/bin/auth-sync-ec2.pl
