#!/usr/bin/env bash

configure_logging() {
  log_folder="/var/log/aws-vpn-client"
  mkdir -p $log_folder

  # $script_type is set by OpenVPN
  if [[ $script_type == *up* ]]; then
    log_file=$log_folder/configure-dns-up.log
  elif [[ $script_type == *down* ]]; then
    log_file=$log_folder/configure-dns-down.log
  fi

  if [[ -f "$log_file" ]]; then
    rm -f "$log_file"
  fi
}

log() {
  echo "$(date) $1" >> "$log_file"
}

check_systemd_resolve_status() {
  output=$(resolvectl status 2>&1)
  exit_code=$?
  log "'resolvectl status' exit code: $exit_code output:"
  log "$output"
}

call_resolvectl() {
  output=$(resolvectl "$@" 2>&1)
  exit_code=$?
  log "resolvectl command exit code: $exit_code, output: $output"
  
  if [[ $exit_code -ne 0 ]]; then
    exit $exit_code
  fi 
}

# $foreign_option_<n> is set by OpenVPN.
# Currently we support DNS and DOMAIN types.
# Search domains are used in the order they appear in the OpenVPN config
# Example:
# "dhcp-option DNS 172.168.0.1"
# "dhcp-option DNS6 2001:4860:4860::8888"
# "dhcp-option DOMAIN amazon.com"
# "dhcp-option DOMAIN amazonaws.com"
# Domains will be set in the order 1) amazon.com, 2) amazonaws.com
get_dns_from_server() {
  log "Getting DNS servers from OpenVPN"
  for fopt in "${!foreign_option_@}"; do
    fopt_val="${!fopt}"
    log "$fopt from OpenVPN: $fopt_val"

    if [[ "${fopt_val}" == *dhcp-option\ DNS* ]]; then
      dns_server_count=$((dns_server_count + 1))
      # dhcp-option DNS 172.168.0.1 -> 172.168.0.1
      # dhcp-option DNS6 2001:4860:4860::8888 -> 2001:4860:4860::8888
      dns_server_ip="${fopt_val#dhcp-option DNS* }"
      dns_servers+=(${dns_server_ip})
    elif [[ "${fopt_val}" == *dhcp-option\ DOMAIN\ * ]]; then
      dns_domain_count=$((dns_domain_count + 1))
      # dhcp-option DOMAIN example.com -> example.com
      dns_domain="${fopt_val#dhcp-option DOMAIN }"
      dns_domains+=("${dns_domain}")
    elif [[ "${fopt_val}" == *dhcp-option\ DOMAIN-ROUTE* ]]; then
      dns_domain_count=$((dns_domain_count + 1))
      # dhcp-option DOMAIN-ROUTE . -> .
      dns_domain="${fopt_val#dhcp-option DOMAIN-ROUTE }"
      # Domain route '~.' needs to be added to prevent DNS leakage
      # https://github.com/systemd/systemd/issues/6076#issuecomment-387332572
      if [[ "${dns_domain}" == "." ]]; then
        dns_domains+=("~${dns_domain}")
      else
        dns_domains+=("${dns_domain}")
      fi 
    fi
  done
  log "Got DNS: ${dns_servers[*]}. DNS domains: ${dns_domains[*]}"
}

# $dev is set by OpenVPN
# "The actual name of the TUN/TAP device, including a unit number if it exists."
get_device_index() {
  log "Getting device index for $dev"
  device_info="$(ip link show dev "$dev")"
  exit_code=$?
  log "'ip link show dev "$dev"' exit code: $exit_code, output: $device_info"
  
  if [[ $exit_code -ne 0 ]]; then
    exit $exit_code
  fi

  device_index="${device_info%%:*}"
  log "Device index for $dev: $device_index"
}

configure_dns() {
  local -a dns_servers=() dns_domains=()
  local -i dns_server_count=0 dns_domain_count=0
  local device_index

  log "Configuring to use DNS servers from OpenVPN"
  get_dns_from_server

  if [[ $dns_server_count -eq 0 && $dns_domain_count -eq 0 ]]; then
    log "No DNS configured for the endpoint"
    return 0
  fi

  if [[ $dns_server_count -gt 0 ]]; then
    params=("dns" "$dev" "${dns_servers[@]}")
    log "Calling resolvectl with parameters: '${params[*]}'"
    # Example: resolvectl dns tun0 172.168.1.2 2001:4860:4860::8888
    call_resolvectl "${params[@]}"
  fi
  if [[ $dns_domain_count -gt 0 ]]; then
    params=("domain" "$dev" "${dns_domains[@]}")
    log "Calling resolvectl with parameters: '${params[*]}'"
    # Example: resolvectl domain tun0 amazon.com ~.
    call_resolvectl "${params[@]}"
  fi
  check_systemd_resolve_status
}

reset_dns() {
  log "Restoring DNS setting to the state before VPN connection"
  local device_index
  get_device_index
  log "Calling resolvectl revert $dev"
  call_resolvectl "revert" "$dev"
  check_systemd_resolve_status
}

main() {
  configure_logging
  
  # $script_type is set by OpenVPN
  log "Executing $script_type script with parameters '$*'"
  
  if [[ $script_type == *up* ]]; then
    configure_dns
  elif [[ $script_type == *down* ]]; then
    reset_dns
  fi
}

main "$@"
