#!/usr/bin/env bash

#================================================================
# DESCRIPTION
#   Script retrieves files from S3 via SQS at maximum warp.
#================================================================
# AUTHOR     Anastas Dancha (@anapsix)
# COPYRIGHT  Copyright (c) Anastas Dancha (@anapsix)
# LICENCE    GNU General Public License
#================================================================
# HISTORY
#   2017/01/24 : anapsix : initial release
#================================================================

help() {
cat <<EOH
Usage: $0 [-h] [-s] -q sqs_queue_name [-m {0..10}] [-t {0..100}]
  -h  display this help
  -s  silent
  -q  sqs queue name
  -m  max_messages to retrieve from SQS between 1 and 10, defaults to 1
  -t  number of threads to start, defaults to "auto"
EOH
}

OPTS=$(getopt -o shq:m:t: -n 'parse-options' -- "$@")
if [ $? != 0 ]; then echo "Failed parsing options." >&2; help; exit 1; fi

eval set -- "$OPTS"
while true; do
  case "$1" in
    -s ) export SILENT=1; shift ;;
    -h ) help; exit 0; shift ;;
    -q ) export SQS_QUEUE_NAME="$2"; shift; shift ;;
    -m ) MAX_MESSAGES="$2"; shift; shift ;;
    -t ) THREADS="$2"; shift; shift ;;
    -- ) shift; break ;;
    * ) break ;;
  esac
done

if [ ${#SQS_QUEUE_NAME[0]} -eq 0 ]; then
  echo >&2 "-s sqs_queue_name is a required argument.."
  help
  exit 1
fi

####

set -e

export MAX_MESSAGES=${MAX_MESSAGES:-1}
export THREADS=${THREADS:-auto}
export MAX_THREADS=50
run_id=$(head -16 /dev/urandom | md5sum | cut -c1-17)

export AWS_DEFAULT_REGION=${AWS_REGION:-us-west-2}
export LOG_DIR="/mnt/elb_logs"

SCRIPT_BASENAME=$(basename $0)
SCRIPT_LOG_DIR="/tmp/${SCRIPT_BASENAME%%.*}"
RUN_LOG_DIR="${SCRIPT_LOG_DIR}/${run_id}"
mkdir -p ${SCRIPT_LOG_DIR} ${RUN_LOG_DIR}

verbose() {
  if [ ${SILENT:-0} -eq 0 ]; then
    echo "${run_id}${thread_id:+/$thread_id}: $@" | tee -a ${SCRIPT_LOG_DIR}/run_${run_id}.log >&2
  fi
}

sqs_delete_message() {
  verbose "Deleting retrieved message id ${message_id}"
  aws sqs delete-message --queue-url ${SQS_QUEUE_URL} --receipt-handle ${1} >${RUN_LOG_DIR}/thread_${thread_id}.log 2>&1
}

s3_get_log() {
  aws s3 cp s3://$1 $2 >${RUN_LOG_DIR}/thread_${thread_id}.log 2>&1
}

main() {
  thread_id=$(head -16 /dev/urandom | md5sum | cut -c1-9)
  TMP_DIR="${LOG_DIR}_tmp/${thread_id}"
  THREAD_LOG="${RUN_LOG_DIR}/thread_${thread_id}.log"
  mkdir ${TMP_DIR}

  messages="$(aws sqs receive-message --max-number-of-messages ${MAX_MESSAGES} --queue-url ${SQS_QUEUE_URL} --output text)"

  IFS=$'\n'

  verbose "Downloading ELB logs to temp dir (${TMP_DIR}/).."
  for message in $messages; do
    message_id=$(echo "$message" | cut -f4)
    receipt_handle=$(echo "$message" | cut -f5)
    bucket_name=$(echo $message | cut -f2 | jq '.Records[].s3.bucket.name' | sed 's/^"//g;s/"$//g')
    logfile=$(echo $message | cut -f2 | jq '.Records[].s3.object.key' | sed 's/^"//g;s/"$//g')

    sqs_delete_message $receipt_handle >${THREAD_LOG} 2>&1
    s3_get_log ${bucket_name}/${logfile} ${TMP_DIR} >${THREAD_LOG} 2>&1
  done

  verbose "Moving ELB logs from temp dir to Splunk input dir (${LOG_DIR}/).."
  chown -v splunk:root ${TMP_DIR}/* >${THREAD_LOG} 2>&1
  mv -v ${TMP_DIR}/* ${LOG_DIR}/ >${THREAD_LOG} 2>&1
  rm -rv ${TMP_DIR} >${THREAD_LOG} 2>&1
}

verbose "Run ${run_id} started.."
verbose "Discovering SQS queue url for ${SQS_QUEUE_NAME}"
export SQS_QUEUE_URL="$(aws sqs get-queue-url --queue-name ${SQS_QUEUE_NAME} --output text)"

if [ "${THREADS}" == "auto" ]; then
  verbose "THREADS autodicovery mode selected.."

  if [ $MAX_MESSAGES -ne 10 ]; then
    verbose "Ignoring passed max_messages, defaulting to maximum of 10"
    export MAX_MESSAGES=10
  fi

  export SQS_NUMBER_OF_MESSAGES=$(aws sqs get-queue-attributes --queue-url ${SQS_QUEUE_URL} --attribute-names ApproximateNumberOfMessages --output text | cut -f2)
  verbose "$SQS_NUMBER_OF_MESSAGES messages detected.."
  export THREADS=$[$[$SQS_NUMBER_OF_MESSAGES/10]+1]
  verbose "Need $THREADS threads to be retriever all of $SQS_NUMBER_OF_MESSAGES messages.."
  if [ $THREADS -gt ${MAX_THREADS} ]; then
    verbose "Cowardly refusing to run more than ${MAX_THREADS} threads, defaulting to ${MAX_THREADS}.."
    export THREADS=${MAX_THREADS}
  fi

fi

verbose "Run is proceeding with ${THREADS} threads.."
for ((i=1;i<=${THREADS};i++)); do
  main &
done
wait
verbose "Run ${run_id} finished.."