#!/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=$(systemd-resolve --status 2>&1)
  exit_code=$?
  log "'systemd-resolve --status' exit code: $exit_code output:"
  log "$output"
}

call_systemd_resolved() {
  log "Calling busctl with parameters '$*'"
  output=$(busctl call "org.freedesktop.resolve1" "/org/freedesktop/resolve1" "org.freedesktop.resolve1.Manager" "$@" 2>&1)
  exit_code=$?
  log "busctl 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 DOMAIN amazon.com"
# "dhcp-option DOMAIN amazonaws.com"
# Domains will be set in the order 1) amazon.com, 2) amazonaws.com
# Client VPN service doesn't support IPv6.
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
      dns_server_ip="${fopt_val#dhcp-option DNS }"
      # Append: 2 4 172 168 0 1
      dns_servers+=(2 4 ${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 }"
      # https://www.freedesktop.org/software/systemd/man/org.freedesktop.resolve1.html
      # "It takes a network interface index and an array of domains, each with a boolean parameter indicating 
      # whether the specified domain shall be used as a search domain (false), or just as a routing domain (true)"
      dns_domains+=("${dns_domain}" false)
    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
      dns_domains+=("${dns_domain}" true)
    fi
  done
}

# $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

  get_device_index
  if [[ $dns_server_count -gt 0 ]]; then
    params=("$device_index" "$dns_server_count" "${dns_servers[@]}")
    log "Calling SetLinkDNS with parameters: '${params[*]}'"
    # Example: SetLinkDNS 'ia(iay)' 125 2 2 4 172 31 29 109 2 4 172 31 4 164
    call_systemd_resolved SetLinkDNS 'ia(iay)' "${params[@]}"
  fi
  if [[ $dns_domain_count -gt 0 ]]; then
    params=("$device_index" "$dns_domain_count" "${dns_domains[@]}")
    log "Calling SetLinkDomains with parameters: '${params[*]}'"
    # Example: SetLinkDomains 'ia(sb)' amazon.com false example.com false
    call_systemd_resolved SetLinkDomains 'ia(sb)' "${params[@]}"
  fi
  check_systemd_resolve_status
}

reset_dns() {
  log "Restoring DNS setting to the state before VPN connection"
  local device_index
  get_device_index
  call_systemd_resolved RevertLink i "$device_index"
  check_systemd_resolve_status
}

# Script args example: tun0 1500 1552 10.0.1.2 255.255.255.224 init
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 "$@"
