#!/bin/bash
set -euo pipefail

LOG_PREFIX="$(pwd) $0"
log() {
    logline="[$LOG_PREFIX] $1\n"
    printf "$logline" >&1
}
log_error() {
    logline="[$LOG_PREFIX] $1\n"
    printf "$logline" >&1
    printf "$logline" >&2
}

if [ -n "${HAPROXY_SYSLOGD-}" ]; then
    SYSLOGD_SERVICE="/marathon-lb/service/syslogd"
    mkdir -p $SYSLOGD_SERVICE
    cp /marathon-lb/syslogd/run "$SYSLOGD_SERVICE/"
fi

# Custom syslog socket for marathon-lb.py logging
SYSLOG_SOCKET=${SYSLOG_SOCKET:-/dev/null}

LB_SERVICE="/marathon-lb/service/lb"
mkdir -p $LB_SERVICE

HAPROXY_SERVICE="/marathon-lb/service/haproxy"
mkdir -p $HAPROXY_SERVICE/env

if [ -n "${PORTS-}" ]; then
  log "$PORTS > $HAPROXY_SERVICE/env/PORTS"
else
  log_error "Define $PORTS with a comma-separated list of ports to which HAProxy binds"
  exit 1
fi

# Find the --ssl-certs arg if one was provided,
# get the certs and remove them and the arg from the list
# of positional parameters so we don't duplicate them
# further down when we pass $@ to marathon_lb.py
declare -i ssl_certs_pos=0
for ((i=1; i<=$#; i++)); do
  if [ "${!i}" = '--ssl-certs' ]; then
    ssl_certs_pos=$(($i+1))
    break
  fi
done
if [ $ssl_certs_pos -gt 0 ]; then
  SSL_CERTS=${!ssl_certs_pos}
  set -- "${@:1:$(($ssl_certs_pos-2))}" "${@:$(($ssl_certs_pos+1))}"
  [ -n "${HAPROXY_SSL_CERT-}" ] && SSL_CERTS+=",/etc/ssl/cert.pem"
else
  SSL_CERTS="/etc/ssl/cert.pem"
fi

if [ -n "${HAPROXY_SSL_CERT-}" ]; then
  # if provided via environment variable, use it.
  echo -e "$HAPROXY_SSL_CERT" > /etc/ssl/cert.pem

  # if additional certs were provided as $HAPROXY_SSL_CERT0 .. 100
  for i in {0..100}; do
    certenv="HAPROXY_SSL_CERT$i"
    if [ -n "${!certenv-}" ]; then
      certfile="/etc/ssl/cert$i.pem"
      echo -e "${!certenv}" > $certfile
      SSL_CERTS+=",$certfile"
    fi
  done
elif [ $ssl_certs_pos -eq 0 ]; then  # if --ssl-certs wasn't passed as arg to this script
  # if no environment variable or command line argument is provided,
  # create self-signed ssl certificate
  openssl genrsa -out /tmp/server-key.pem 2048
  openssl req -new -key /tmp/server-key.pem -out /tmp/server-csr.pem -subj /CN=*/
  openssl x509 -req -in /tmp/server-csr.pem -out /tmp/server-cert.pem -signkey /tmp/server-key.pem -days 3650
  cat /tmp/server-cert.pem /tmp/server-key.pem > /etc/ssl/cert.pem
  rm /tmp/server-*.pem
fi

if [ -n "${MESOS_SANDBOX-}" ] && [ -d "$MESOS_SANDBOX/templates" ]; then
  mkdir -p templates
  cp -v "$MESOS_SANDBOX/templates/"* templates/
fi

if [ -n "${HAPROXY_SYSCTL_PARAMS-}" ]; then
  log "setting sysctl params to: ${HAPROXY_SYSCTL_PARAMS}"
  if [ -n "${HAPROXY_SYSCTL_NONSTRICT-}" ]; then
    # ignore errors
    sysctl -w $HAPROXY_SYSCTL_PARAMS || true
  else
    sysctl -w $HAPROXY_SYSCTL_PARAMS
  fi
fi

MODE=$1; shift
case "$MODE" in
  poll)
    POLL_INTERVAL="${POLL_INTERVAL:-60}"
    ARGS=""
    ;;
  sse)
    ARGS="--sse"
    ;;
  *)
    log_error "Unknown mode $MODE. Synopsis: $0 poll|sse [marathon_lb.py args]"
    exit 1
    ;;
esac

if [ -n "${VAULT_TOKEN-}" ] && [ -n "${VAULT_HOST-}" ] && [ -n "${VAULT_PORT-}" ] && [ -n "${VAULT_PATH-}" ]; then
  MARATHON_LB_PASSWORD=$(curl -k -L -H "X-Vault-Token:$VAULT_TOKEN" "$VAULT_URL" -s| python -m json.tool | python -c 'import json,sys;obj=json.load(sys.stdin);print obj["data"]["pass"]')
  MARATHON_LB_USER=$(curl -k -L -H "X-Vault-Token:$VAULT_TOKEN" "$VAULT_URL" -s | python -m json.tool | python -c 'import json,sys;obj=json.load(sys.stdin);print obj["data"]["user"]')
  CREDENTIALS="$MARATHON_LB_USER:$MARATHON_LB_PASSWORD"
  ARGS="$ARGS --auth-credentials $CREDENTIALS"
fi

for arg in "$@"; do
    escaped=$(printf %q "$arg")
    ARGS="$ARGS $escaped"
done

cat > $LB_SERVICE/run << EOF
#!/bin/sh
exec 2>&1
cd /marathon-lb
exec /marathon-lb/marathon_lb.py \
    --syslog-socket $SYSLOG_SOCKET \
    --haproxy-config /marathon-lb/haproxy.cfg \
    --ssl-certs "${SSL_CERTS}" \
    --command "sv reload ${HAPROXY_SERVICE}" \
    $ARGS
EOF
chmod 755 $LB_SERVICE/run

log "Created $LB_SERVICE/run with contents:"
LB_RUN=$(cat $LB_SERVICE/run)
log "$LB_RUN"

if [ "${MODE}" == "poll" ]; then

cat > $LB_SERVICE/finish << EOF
#!/bin/sh
sleep ${POLL_INTERVAL}

EOF
chmod 755 $LB_SERVICE/finish

fi

runsvdir -P /marathon-lb/service &
trap "kill -s 1 $!" TERM INT
wait
