#!/bin/bash

set -euo pipefail

# tp - Twitch proxy script to make network requests in the Twitch network when connected to the Amazon network (WPA2)

# The host you SSH into to make your proxy tunnel
BASTION_HOST="${BASTION_HOST:-us-west-2.prod.bastion.live-video.a2z.com}"
# The SSH control file that we send shutdown signals on (also helps with pkill to kill outstanding tunnels)
CONTROL_FILE="${HOME}/.tp.ssh.sock"
# The local port we connect to the SOCKS proxy on
LOCAL_PROXY_PORT="${LOCAL_PROXY_PORT-1080}"
# The local port we connect to the HTTP proxy on
LOCAL_HTTP_PROXY_PORT="${LOCAL_HTTP_PROXY_PORT-12345}"
# The URL to test for twitch connectivity
TEST_URL="${TEST_URL-https://code.justin.tv}"
# mwinit command. --aea works on the Amazon network
MWINIT="${MWINIT-mwinit --aea}"
# midway cookie path
MIDWAY_COOKIE="${MIDWAY_COOKIE-${HOME}/.midway/cookie}"
# SSH config file
SSH_CONFIG="${SSH_CONFIG-${HOME}/.ssh/config}"

function auth() {
  if [[ "${1-}" == "help" ]]; then
    echo "Authenticate for the proxy"
    exit 0
  fi

  auth_midway

  # Only add key if it's not already loaded in the agent
  if ssh-add -l | grep -q "$(ssh-keygen -lf ~/.ssh/id_*sa | awk '{print $2}')"; then
    echo "Key loaded in SSH agent. Skipping ssh-add"
  else
    ssh-add -q
  fi
}


function auth_midway() {
  if [[ ! -f "${MIDWAY_COOKIE}" ]]; then
    $MWINIT
    return
  fi

  if check_midway_auth; then
    echo "Authenticated to midway. Skipping mwinit"
    return
  fi

  $MWINIT
}

# Last successful authenticated session status HTTP body
AUTHENTICATED_MIDWAY_SESSION_HTTP_BODY=""

function check_midway_auth() {
  # Check if the midway cookie is valid. This assumes the cookie and SSH certificate have the same expiration time.
  local session_status=$(curl --silent --write-out "HTTPSTATUS:%{http_code}" -q -b "${MIDWAY_COOKIE}" --cookie-jar "${MIDWAY_COOKIE}" -L https://midway-auth.amazon.com/api/session-status);
  local http_status=$(echo "${session_status}" | tr -d '\n' | sed -e 's/.*HTTPSTATUS://')
  local http_body=$(echo "${session_status}" | sed -e 's/HTTPSTATUS\:.*//g')
  if (( "200" != "${http_status}")); then
    echo "Invalid midway session-status status code: ${http_status}" >&2
    return 1
  fi

  local authenticated=$(echo "$http_body" | tr '{},' '\n\n\n' | sed -n 's/"authenticated":\([^,]*\)/\1/p')
  if [[ "false" = "$authenticated" ]]; then
    return 1
  fi

  # Store the body in case extracting other fields is needed
  AUTHENTICATED_MIDWAY_SESSION_HTTP_BODY="$http_body"

  return 0
}

function midway_auth_expires_at() {
  if [[ -z "$AUTHENTICATED_MIDWAY_SESSION_HTTP_BODY" ]]; then
    return
  fi

  local http_body="$AUTHENTICATED_MIDWAY_SESSION_HTTP_BODY"
  local expires_at=$(echo "$http_body" | tr '{},' '\n\n\n' | sed -n 's/"expires_at":\([^,]*\)/\1/p')
  # This should work cross-platform
  echo "$(perl -e 'print scalar localtime $ARGV[0]' $expires_at)"
}

function root__enable() {
  if [[ "${1-}" == "help" ]]; then
    echo "Enable proxy. Authenticates if needed. Restarts if already enabled"
    exit 0
  fi

  if ssh_status 2>/dev/null; then
    echo "Proxy SSH already enabled"
    root__disable
  fi

  auth

  insert_bastion_ssh_config

  # Use timeout because the ssh command hangs when not on the Amazon network

  # -f Requests ssh to go to background just before command execution
  # -T disable pseudo-terminal allocation
  # -N do not execute remote command
  # -D Specifies a local dynamic application-level port forwarding
  timeout 5 ssh -fTN \
    -o LogLevel=fatal \
    -o ServerAliveInterval=5 \
    -o ServerAliveCountMax=3 \
    -o ConnectionAttempts=1 \
    -o ConnectTimeout=3 -D 127.0.0.1:$LOCAL_PROXY_PORT ${BASTION_HOST} \
    -o ControlPath="${CONTROL_FILE}" \
    -o ControlPersist=no \
    || (echo "Failed to connect to proxy. Are you on the Amazon network? Is teleport-bastion disabled?" && exit 1)
  echo "Proxy connected"

  # Setup http proxy over socks proxy
  nohup hpts -s "127.0.0.1:${LOCAL_PROXY_PORT}" -p "${LOCAL_HTTP_PROXY_PORT}" &> /dev/null &

  echo "Running a sample to test twitch network access"
  check_proxy_connection >/dev/null

  echo "Network test success!"
}

function root__disable() {
  if [[ "${1-}" == "help" ]]; then
    echo "Disable proxy"
    exit 0
  fi

  remove_bastion_ssh_config

  # This should be enough, but sometimes SSH sessions linger. So we kill them with pkill below
  ssh -S ${CONTROL_FILE} -TO exit hostname || true
  rm -f ${CONTROL_FILE}
  pkill -f "ssh.*${BASTION_HOST}.*${CONTROL_FILE}" || true
  pkill -f "hpts -s 127.0.0.1:${LOCAL_PROXY_PORT} -p ${LOCAL_HTTP_PROXY_PORT}"
  echo "Disabled proxy"
}

function insert_bastion_ssh_config() {
  # Upsert bastion config into SSH config file
  local ssh_dir="$(dirname $SSH_CONFIG)"
  if [[ ! -d "$ssh_dir" ]]; then
    echo "Missing SSH directory $ssh_dir" >&2
    return 1
  fi

  if grep -q "# BEGIN video bastion ssh config" "${SSH_CONFIG}"; then
    # Already exists. Do nothing
    return 0
  fi

  echo "# BEGIN video bastion ssh config
# Each AZ has a bastion host, and they share the same name. Disable key checking so they don't conflict.

Host ${BASTION_HOST}
  CertificateFile ~/.ssh/id_rsa-cert.pub
  StrictHostKeyChecking no
  UserKnownHostsFile /dev/null
  ControlMaster auto
  ControlPersist 5
  ProxyCommand none

# Use the bastion hosts for all ssh access, except for accessing the bastions
# and services that are available on both networks.
Host !${BASTION_HOST} *.justin.tv *.live-video.a2z.com
  ProxyJump ${BASTION_HOST}

# END video bastion ssh config" >> "$SSH_CONFIG"
}

function remove_bastion_ssh_config() {
  if [[ ! -f "$SSH_CONFIG" ]]; then
    return
  fi

  # Remove bastion config from SSH config file
  #
  #  -i modifies the file in place
  #  -0 slurps the entire file
  #  -p loop around the program
  #  -e perl program
  #
  #   m - Treat string as multiple lines. That is, change "^" and "$" from matching the start or end of line only at the left and right ends of the string to matching them anywhere within the string.
  #   s - Treat string as single line. That is, change "." to match any character whatsoever, even a newline, which normally it would not match.
  perl -i -0pe 's/(\n)?# BEGIN video bastion ssh config.*# END video bastion ssh config//sm' "$SSH_CONFIG"
}


function ssh_status() {
  ssh -S ${CONTROL_FILE} -TO check hostname
}

function root__status() {
  if [[ "${1-}" == "help" ]]; then
    echo "Check status to proxy and midway"
    exit 0
  fi

  local exitcode=0

  local proxy_status="disconnected"
  if check_proxy_connection >/dev/null; then
    proxy_status="connected"
  else
    exitcode=1
  fi
  echo "SOCKS Proxy status: ${proxy_status}"

  local http_proxy_status="disconnected"
  if check_http_proxy_connection >/dev/null; then
    http_proxy_status="connected"
  else
    exitcode=1
  fi
  echo "HTTP Proxy status: ${http_proxy_status}"

  local midway_status="unauthenticated"
  if check_midway_auth; then
    midway_status="authenticated (Expires $(midway_auth_expires_at))"
  else
    exitcode=1
  fi
  echo "Midway status: ${midway_status}"

  return $exitcode
}

function check_proxy_connection() {
  root__proxy curl --silent --show-error -q --connect-timeout 1 --fail "${TEST_URL}"
}

function check_http_proxy_connection() {
  root__http_proxy curl  --silent --show-error -q --connect-timeout 1 --fail "${TEST_URL}"
}

function root__curl() {
  if [[ "${1-}" == "help" ]]; then
    echo "Executes curl over the twitch network. Shortcut for 'tp proxy curl'"
    exit 0
  fi

  root__proxy curl "$@"
}

function root__proxy() {
  if [[ "${1-}" == "help" ]]; then
    echo "Execute a command using SSH as a SOCKS proxy. The command must accept env variables: ALL_PROXY or HTTP_PROXY"
    exit 0
  fi

  # ALL_PROXY:  socks5h means to resolve the hostname through the proxy; this is important for resolving the code.justin.tv on
  #             the Twitch network instead of using the corp proxy. curl uses this variable.
  #
  # HTTP_PROXY: For go if the HTTP transport is using http.ProxyFromEnvironment.
  ALL_PROXY="socks5h://localhost:$LOCAL_PROXY_PORT" HTTP_PROXY="socks5://localhost:$LOCAL_PROXY_PORT" "$@"
}

function root__http_proxy() {
  if [[ "${1-}" == "help" ]]; then
    echo "Execute a command using http proxy. The command must accept env variables: HTTP_PROXY / HTTPS_PROXY / http_proxy / https_proxy"
    exit 0
  fi

  HTTP_PROXY="http://localhost:${LOCAL_HTTP_PROXY_PORT}" \
    HTTPS_PROXY="http://localhost:${LOCAL_HTTP_PROXY_PORT}" \
    http_proxy="http://localhost:${LOCAL_HTTP_PROXY_PORT}" \
    https_proxy="http://localhost:${LOCAL_HTTP_PROXY_PORT}" \
    "$@"
}

function root__ssh() {
  if [[ "${1-}" == "help" ]]; then
    echo "Execute SSH command on bastion"
    exit 0
  fi

  ssh -o ControlMaster=no ${BASTION_HOST} "$@"
}

# Process a help flag for every subcommand
function process_help() {
  prefix=$1
  shift
  if [[ "${1-}" == "help" ]]; then
    echo "Prints each command and help"
    exit 0
  fi

  compgen -A function | grep $prefix | sort | while read -r line; do
    echo
    echo "  ${line:${#prefix}}"
    echo "    $($line help)"
  done
}

# This is the core to our __ expansion magic
function process_build() {
  prefix=$1
  shift
  T="$prefix${1-}"
  if declare -F "$T" >/dev/null ; then
    func="$prefix${1}"
    shift; # pop $1
    "$func" "$@"    # invoke our named function w/ all remaining arguments
  else
    echo "Undefined subcommand ${1-}"
    process_help $prefix
    exit 1
  fi
}

process_build root__ "$@"
