#!/bin/bash

get_vpc_cidrs()
{
    cidrs=$(curl -s $metadata/network/interfaces/macs/$mac_address/vpc-ipv4-cidr-blocks)
    echo "$cidrs"
}

run_dnsmasq()
{
    all_domains="$(grep ^search $resolv_conf | cut -d' ' -f2- )"
    pid=$(ps -ef | grep dnsmasq | grep synth-domain | awk '{print $2}')
    if [ $pid ]; then
        sudo kill $pid
    fi
    for d in $all_domains; do
        for c in $(get_vpc_cidrs); do
            syn_domains="$syn_domains --synth-domain=$d,$c,ip- "
        done
    done
    runmasq="sudo dnsmasq --listen-address=127.0.0.1 $syn_domains "
    eval "$runmasq"
    echo "started dnsmasq : $runmasq"
}

rewrite_setup_dns()
{
  tmpfile=$(mktemp /tmp/setupdnsXXXXXX)
  cat > "$tmpfile" << 'EOF'
  #!/bin/bash
  #
  # Set up DNS for EMR master/slave instance in VPC.
  # This script also set up DNS in us-east-1 for non-VPC to handle ec2 instances,
  # whose host name begin with domU, with invalid dns domain name (TT0055043598).
  #
  set -e
  set -x

  alias curl="curl --connect-timeout 2 -q -f --retry-delay 2 --retry 5"

  resolv_conf="/etc/resolv.conf"
  dhclient_conf="/etc/dhcp/dhclient.conf"
  localhost="127.0.0.1"
  metadata="http://169.254.169.254/latest/meta-data"

  restart_network="false"
  in_vpc="false"

  mac_address="$(curl $metadata/mac/ | tr '[:upper:]' '[:lower:]')"
  region="$(curl http://169.254.169.254/latest/dynamic/instance-identity/document | jq -r .region)"

  # wait for the network to come up before proceeding
  if [ -e /usr/bin/nm-online ]; then
      /usr/bin/nm-online
  fi

  get_default_domain()
  {
      if [ "$region" = "us-east-1" ]; then
          echo 'ec2.internal'
      else
          echo "$region.compute.internal"
      fi
  }

  get_first_nameserver_from_resolv_conf()
  {
      awk '$1 ~ /^nameserver/ { print $2 }' "$resolv_conf"
  }

  check_vpc()
  {
      if "$(curl $metadata/network/interfaces/macs/$mac_address/)" | grep -q vpc; then
          in_vpc="true"
      fi
  }

  get_vpc_cidrs()
  {
      cidrs=$(curl $metadata/network/interfaces/macs/$mac_address/vpc-ipv4-cidr-blocks)
      echo "$cidrs"
  }

  append_line_to_dhclient_conf()
  {
      echo "$1" | tee -a "$dhclient_conf"
  }

  prepend_domain()
  {
      #sample line : prepend domain-name "ec2.internal  ";
      if grep -Eq "^prepend domain-name \"$1[:space:]+\";$" "$dhclient_conf"; then
          return
      else
          append_line_to_dhclient_conf "prepend domain-name \"$1 \";"
          restart_network="true"
      fi
  }

  prepend_domain_server()
  {
      #sample line : prepend domain-name-servers 127.0.0.1;
      if grep -Eq "^prepend domain-name-servers $1;$" "$dhclient_conf"; then
          return
      fi
      append_line_to_dhclient_conf "prepend domain-name-servers $1;"
      restart_network="true"
  }

  run_dnsmasq()
  {
      all_domains="$(grep ^search $resolv_conf | cut -d' ' -f2- )"
      pid=$(ps -ef | grep dnsmasq | grep synth-domain | awk '{print $2}')
      if [ $pid ]; then
          kill $pid
      fi
      for d in $all_domains; do
          for c in $(get_vpc_cidrs); do
              syn_domains="$syn_domains --synth-domain=$d,$c,ip- "
          done
      done
      runmasq="dnsmasq --listen-address=127.0.0.1 $syn_domains "
      eval "$runmasq"
      echo "started dnsmasq : $runmasq"
  }

  get_host_name()
  {
      echo "$(hostname -f)"
  }

  show_dns_status()
  {
      type="$1"
      echo "------------ $type $resolv_conf ------------"
      cat "$resolv_conf"
      echo "------------ $type $dhclient_conf ------------"
      cat "$dhclient_conf"
      hostname="$(get_host_name)"
      status="$?"
      "'hostname -f' returns : $hostname"
      return $status
  }

  restart_network_if_needed()
  {
      if "$restart_network"; then
          echo "Updating DNS settings."
          service network restart
          restart_network="false"
      fi
  }

  main()
  {
      show_dns_status "BeforeSetup"

      old_domain="$(grep search $resolv_conf | cut -d' ' -f2-)"
      default_domain="$(get_default_domain)"

      check_vpc

      if [ "$in_vpc" = "false" ]; then
          # NON-VPC
          if [ "$region" = "us-east-1" ]; then
              if [[ "$old_domain" == "${default_domain}"* ]]; then
                  echo "$default_domain is already used in us-east-1."
              else
                  echo "Making sure $default_domain is used in us-east-1."
                  prepend_domain $default_domain
              fi
          else
              echo "Not in VPC, do nothing and exit."
          fi
      else
          # VPC
          first_nameserver="$(get_first_nameserver_from_resolv_conf)"
          resolving_host_name="$(get_host_name)"
          if [ "$1" = "rundnsmasq" -o -z "$resolving_host_name" ]; then
              echo "Run dnsmasq"
              run_dnsmasq
              if [ "$first_nameserver" != "$localhost" ]; then
                  prepend_domain_server "$localhost"
              fi
          else
              echo "Resolving hostname(${resolving_host_name}) successfully, do nothing and exit."
          fi
      fi

      restart_network_if_needed
      return show_dns_status "AfterSetup"
  }

  main "$@"
  exit "$?"
EOF

  sudo mv $tmpfile /usr/bin/setup-dns
}

if [ ! -f /tmp/dns_flag ]; then
  resolv_conf="/etc/resolv.conf"
  metadata="http://169.254.169.254/latest/meta-data"
  mac_address=`curl -s $metadata/mac`
  run_dnsmasq
  rewrite_setup_dns
  touch /tmp/dns_flag
  pid="$(/bin/ps axwwo pid,cmd | awk '$12 ~ /aws157.instancecontroller.Main/ { print $1 }')"
  sudo kill "$pid"
fi