#!/bin/bash
set -e

source ${SCRIPT_DIR}/logging
source ${SCRIPT_DIR}/kasm_api

kasm_id=$1
container_name=$2
container_ns_id=$3
container_pid=$4
script_data_encoded=$5
notification_pipe=$6
container_setup_progress=0

wg_log_file=/var/run/kasm-sidecar/$container_ns_id/wg.log

notify_failure() {
  set +e

  # update the workspace status
  set_container_status "starting" "$container_setup_progress" "Connection to egress gateway failed. Contact an administrator if the issue persists."
  
  # log the error
  if [ -f "$wg_log_file" ]; then
    wg_log=$(cat "$wg_log_file")
    error "Wireguard log: $wg_log"
  fi

  # show the error notification
  if [ "$show_vpn_status" == "1" ]; then
    show_container_notification "error" "VPN Connection Failed" "VPN Configuration Failed. An error has occurred starting the VPN. Refer to the logs for more information."
  fi

  # mark the container as failed with a delay to ensure the workspace error is visible to the user since
  # the container relies on REST API calls to update the status
  sleep 10
  nsenter --mount=/proc/$container_pid/ns/mnt --net=/proc/$container_pid/ns/net -- sh -c "echo 'error' > /dockerstartup/.egress_status"
  exit 0
}

if [ -z "$container_name" ] || [ -z "$container_ns_id" ] || [ -z "$container_pid" ] || [ -z "$script_data_encoded" ]; then
  error "Missing parameters. Got $@. Exiting."
  exit 1
fi

trap notify_failure ERR

container_setup_progress=5
set_container_status "starting" "$container_setup_progress" "Connecting to egress gateway..."

# tweak to wg
sed -i '/cmd sysctl -q/d' $(which wg-quick)

# extract script data
script_data=$(echo $script_data_encoded | base64 -d)
egress_provider=$(echo $script_data | jq -r '.egress_provider')
egress_gateway=$(echo $script_data | jq -r '.egress_gateway')
egress_country=$(echo $script_data | jq -r '.egress_country')
egress_city=$(echo $script_data | jq -r '.egress_city')
show_vpn_status=$(echo $script_data | jq -r '.show_vpn_status')
show_ip_status=$(echo $script_data | jq -r '.show_ip_status')

if [ "$show_vpn_status" == "1" ]; then
  show_container_notification "offline" "Connecting to VPN" "Please wait while the VPN connection is being established..."
fi

# prepare the config file
config_file=/var/run/kasm-sidecar/$container_ns_id/wg.conf
info "Preparing config file" 
echo $script_data | jq -r '.config' > $config_file
chmod 600 $config_file

sed -i '/^DNS = /d' $config_file
info "Config file \"$config_file\" ready"

# run the VPN connection
info "Connecting to egress gateway \"$egress_gateway\" ($wg_log_file)"
nsenter --net=/proc/$container_pid/ns/net wg-quick up $config_file > >(tee -a $wg_log_file &> /dev/null) 2>&1

# verify a network device was created
nsenter --net=/proc/$container_pid/ns/net wg show > /dev/null
ret=$?
if [ $ret -ne 0 ]; then
  error "Connection failed."
  notify_failure
fi

wg_log=$(cat "$wg_log_file")
info "Connection established successfully: $wg_log"

# mark the container as ready
container_setup_progress=10
set_container_status "starting" "$container_setup_progress" "Connection to egress gateway established."
nsenter --mount=/proc/$container_pid/ns/mnt --net=/proc/$container_pid/ns/net -- sh -c "echo 'ready' > /dockerstartup/.egress_status"

# show the VPN status
if [ "$show_vpn_status" == "1" ]; then
  show_container_notification "info" "VPN Connected!" "Provider: $egress_provider\nCountry: $egress_country\nCity: $egress_city\nGateway: $egress_gateway"
fi

# show the connection details
if [ "$show_ip_status" == "1" ]; then
  attempt=1
  attempt_delay=5
  max_attempts=10

  while [ $attempt -le $max_attempts ]
  do  
    ip_status=$(nsenter --mount=/proc/$container_pid/ns/mnt --net=/proc/$container_pid/ns/net -- bash -c "curl -q --max-time 5 -s https://ipleak.net/json/ || true")
    
    if [ -n "$ip_status" ]; then
      info "Successfully retrieved public IP information"
      break
    fi
    
    error "Failed to retrieve public IP information (${attempt} attempts)"
    attempt=$((attempt + 1))
    sleep $attempt_delay
  done

  if [ -z "$ip_status" ]; then
    show_container_notification "error" "Public IP Status" "Failed to retrieve public IP information. The connection may not be fully established."
  else  
    show_container_notification "info" "Public IP Information" "$ip_status"
  fi
fi