# -*-Shell-script-*-

#
# Copyright (C) 2018 velostrata.com
# All Rights Reserved.
#

GCP_META_FALLBACK_HOST="169.254.169.254"
GCP_META_ROOT="computeMetadata/v1"
GCP_META_ZONE="instance/zone"
GCP_META_PROJECT_ATTRIBUTES="project/attributes"
DNS_DATA_NOT_FOUND="###gcpdns-not-found###"

DHCLIENT_UP_HOOKS=/etc/dhcp/dhclient-up-hooks
NM_CONFD=/etc/NetworkManager/conf.d
NETCONFIG=/etc/sysconfig/network/config
SYSTEMD_RESOLVED=/etc/systemd/resolved.conf
RESOLV_CONF=/etc/resolv.conf

BACKUP_DIR=/etc/migrate-for-gce/backup/gcpdns

log_err() {
	echo -e "[`date`] MigrateForGCE: $*" >> /var/run/migrate-for-gce-gcpdns.log
}

get_meta_with_retries() {
	local key header attempts rc
	key=$1
	header=$2
	rc=1

	attempts=$3
	[ -z "${attempts}" ] && attempts=2

	log_err "QUERY ${key}"
	while [ "${rc}" -ne 0 ] && [ "${attempts}" -gt 0 ] ; do
		meta=$(curl -s -f -H "${header}" ${key})
		if [ "${?}" -eq 0 ] ; then
			echo ${meta}
			return
		fi

		let attempts--
		if [ "${attempts}" -gt 0 ]; then
			sleep 1
		fi
	done
	echo "${DNS_DATA_NOT_FOUND}"
}

get_gcp_meta() {
	get_meta_with_retries "http://${GCP_META_FALLBACK_HOST}/${GCP_META_ROOT}/$1" "Metadata-flavor: Google"
}

get_gcp_project_meta() {
	get_gcp_meta "${GCP_META_PROJECT_ATTRIBUTES}/$1"
}

get_gcp_region() {
	local current_zone
	current_zone=$(get_gcp_meta "${GCP_META_ZONE}")
	if [ ${current_zone} != "${DNS_DATA_NOT_FOUND}" ] ; then
		echo $current_zone | cut -d "/" -f 4 | cut -d "-" -f 1-2
	fi
}

get_dns_data() {
	local region key regiondata
	region=$1
	key=$2
	regiondata=${DNS_DATA_NOT_FOUND}

	# try region first
	if [ ! -z "${region}" ] ; then
		regiondata=$(get_gcp_project_meta "${region}_${key}")
	fi

	# fallback to default if not specified
	if [ "${regiondata}" != "${DNS_DATA_NOT_FOUND}" ] ; then
		echo ${regiondata}
	else
		get_gcp_project_meta "default_${key}"
    fi
}

gcp_dns_servers() {
	get_dns_data "$(get_gcp_region)" "dns-servers"
}

gcp_dns_suffixes() {
	get_dns_data "$(get_gcp_region)" "dns-domain-suffixes"
}

backup_file() {
	local src backup_name
	src=$1
	backup_name=$2

	mkdir -p ${BACKUP_DIR}
	if [ ! -f ${BACKUP_DIR}/${backup_name} ] ; then
		cp ${src} ${BACKUP_DIR}/${backup_name}
	fi
}

restore_original_file() {
	local backup_name restore_path
	backup_name=$1
	restore_path=$2

	if [ -f ${BACKUP_DIR}/${backup_name} ] ; then
		cat ${BACKUP_DIR}/${backup_name} > ${restore_path}
	fi
}

execute_resolvconf() {
	local to_exec nic
	to_exec=$1
	nic=$(ls /sys/class/net | grep -v lo | head -n 1)

	sed -i '/#/d' ${to_exec}
	sed -i '/^\s*$/d' ${to_exec}
	if [ -s ${to_exec} ] ; then
		cat ${to_exec} | resolvconf -a $nic
		resolvconf -u
	fi
}

configure_resolvconf() {
	local dns_servers dns_suffixes
	dns_servers="$1"
	dns_suffixes="$2"

	backup_file ${RESOLV_CONF} resolvconf_exec

	to_exec=`mktemp`
	generate_resolv_conf "${to_exec}" "${dns_servers}" "${dns_suffixes}"
	execute_resolvconf ${to_exec}
	rm -f ${to_exec}
}

generate_resolv_conf() {
	local out_file dns_servers dns_suffixes temp_file
	conf_file=$1
	dns_servers="$2"
	dns_suffixes="$3"

	temp_file=`mktemp`

	if [ -z "${dns_suffixes}" ] ; then
		echo "search google.internal" >> ${temp_file}
	elif [ "${dns_suffixes}" == "${DNS_DATA_NOT_FOUND}" ] ; then
		cat ${conf_file} | grep search >> ${temp_file}
	else
		echo search $(echo ${dns_suffixes} | tr "," " ") >> ${temp_file}
	fi

	if [ -z "${dns_servers}" ] ; then
		echo "nameserver ${GCP_META_FALLBACK_HOST}" >> ${temp_file}
	elif [ "${dns_servers}" == "${DNS_DATA_NOT_FOUND}" ] ; then
		cat ${conf_file} | grep nameserver >> ${temp_file}
	else
		for ns in $(echo ${dns_servers} | tr "," " ") ; do
			echo nameserver ${ns} >> ${temp_file}
		done
	fi

	echo "# auto-generated from GCP project meta-data" > ${conf_file}
	cat ${temp_file} >> ${conf_file}
	rm -f ${temp_file}
}


get_network_config_path() {
	if [ -f "/etc/sysconfig/network-scripts/ifcfg-lo" ]; then
		echo "/etc/sysconfig/network-scripts"
	else
		echo "/etc/sysconfig/network"
	fi
}

update_netconfig() {
	local dns_servers dns_suffixes temp_file
	dns_servers="$1"
	dns_suffixes="$2"

	backup_file ${NETCONFIG} netconfig

	temp_file=`mktemp`
	cat ${NETCONFIG} > ${temp_file}

	if [ -z "${dns_suffixes}" ] ; then
		sed -i 's/NETCONFIG_DNS_POLICY=.*/NETCONFIG_DNS_POLICY="auto"/g' ${temp_file}
		sed -i "s/NETCONFIG_DNS_STATIC_SEARCHLIST=.*/NETCONFIG_DNS_STATIC_SEARCHLIST=\"\"/g" ${temp_file}
	elif [ "${dns_suffixes}" != "${DNS_DATA_NOT_FOUND}" ] ; then
		sed -i 's/NETCONFIG_DNS_POLICY=.*/NETCONFIG_DNS_POLICY="auto"/g' ${temp_file}
		sed -i "s/NETCONFIG_DNS_STATIC_SEARCHLIST=.*/NETCONFIG_DNS_STATIC_SEARCHLIST=\"$(echo ${dns_suffixes} | tr "," " ")\"/g" ${temp_file}
	fi

	if [ -z "${dns_servers}" ] ; then
		sed -i 's/NETCONFIG_DNS_POLICY=.*/NETCONFIG_DNS_POLICY="auto"/g' ${temp_file}
		sed -i "s/NETCONFIG_DNS_STATIC_SERVERS=.*/NETCONFIG_DNS_STATIC_SERVERS=\"\"/g" ${temp_file}
	elif [ "${dns_servers}" != "${DNS_DATA_NOT_FOUND}" ] ; then
		sed -i 's/NETCONFIG_DNS_POLICY=.*/NETCONFIG_DNS_POLICY="auto"/g' ${temp_file}
		sed -i "s/NETCONFIG_DNS_STATIC_SERVERS=.*/NETCONFIG_DNS_STATIC_SERVERS=\"$(echo ${dns_servers} | tr "," " ")\"/g" ${temp_file}
	fi

	cat ${temp_file} > ${NETCONFIG}
	rm -f ${temp_file}
	
	netconfig update -f
}

update_setting() {
	local work_file key value
	work_file=$1
	key=$2
	value="$3"

	if ! grep -q "^[#]*${key}=" ${work_file} ; then
		echo "#${key}=" >> ${work_file}
	fi

	if [ -z "${value}" ] ; then
		sed -i "s/^[#]*${key}=.*/#${key}=/g" ${work_file}
	else
		sed -i "s/^[#]*${key}=.*/${key}=${value}/g" ${work_file}
	fi
}


update_resolved() {
	local dns_servers dns_suffixes temp_file
	dns_servers="$1"
	dns_suffixes="$2"

	backup_file ${SYSTEMD_RESOLVED} systemd_resolved

	temp_file=`mktemp`
	cat ${SYSTEMD_RESOLVED} > ${temp_file}

	if ! grep -q "[Resolve]" ${temp_file} ; then
		echo "[Resolve]" >> ${temp_file}
	fi

	if [ -z "${dns_suffixes}" ] ; then
		update_setting ${temp_file} "Domains" ""
	elif [ "${dns_suffixes}" != "${DNS_DATA_NOT_FOUND}" ] ; then
		update_setting ${temp_file} "Domains" "$(echo ${dns_suffixes} | tr "," " ")"
	fi

	if [ -z "${dns_servers}" ] ; then
		update_setting ${temp_file} "DNS" ""
	elif [ "${dns_servers}" != "${DNS_DATA_NOT_FOUND}" ] ; then
		update_setting ${temp_file} "DNS" "$(echo ${dns_servers} | tr "," " ")"
	fi

	cat ${temp_file} > ${SYSTEMD_RESOLVED}
	rm -f ${temp_file}

	systemctl try-restart systemd-resolved.service
}

enable_resolv_conf_generation()
{
	local ifcfg
	ifcfg=$1

	rm -f ${NM_CONFD}/20-nodns.conf
	sed -i 's/PEERDNS=.*/PEERDNS="yes"/g' ${ifcfg}
	sed -i -e '/#MIGRATE-FOR-GCE-GCPDNS/,/#END-MIGRATE-FOR-GCE-GCPDNS/c\' ${DHCLIENT_UP_HOOKS}
}

disable_resolv_conf_generation()
{
	local ifcfg
	ifcfg=$1

	printf "[main]\ndns=none\n" > ${NM_CONFD}/20-nodns.conf
	if [ ! -z "${ifcfg}" ] && [ -f "${ifcfg}" ] ; then
		if grep -q "PEERDNS" ${ifcfg} ; then
			sed -i 's/PEERDNS=.*/PEERDNS="no"/g'  ${ifcfg}
		else
			echo "PEERDNS=no" >> ${ifcfg}
		fi
	fi

	if [ ! -e ${DHCLIENT_UP_HOOKS} ] || [ ! -s ${DHCLIENT_UP_HOOKS} ] ; then
		echo "" > ${DHCLIENT_UP_HOOKS}
	fi
	if ! grep -q "MIGRATE-FOR-GCE-GCPDNS" ${DHCLIENT_UP_HOOKS} ; then
		sed -i '1i #MIGRATE-FOR-GCE-GCPDNS\nmake_resolve_conf() { true; }\n#END-MIGRATE-FOR-GCE-GCPDNS\n' ${DHCLIENT_UP_HOOKS}
	else
		sed -i -e '/#MIGRATE-FOR-GCE-GCPDNS/,/#END-MIGRATE-FOR-GCE-GCPDNS/c\#MIGRATE-FOR-GCE-GCPDNS\nmake_resolve_conf() { true; }\n#END-MIGRATE-FOR-GCE-GCPDNS' ${DHCLIENT_UP_HOOKS}
	fi
	chmod +x ${DHCLIENT_UP_HOOKS}
}


wait_for_meta() {
	local dummy
	dummy=$(get_meta_with_retries "http://${GCP_META_FALLBACK_HOST}/${GCP_META_ROOT}" "Metadata-flavor: Google" 10)
}

update_dns_config() {
	local dns_servers dns_suffixes nic ifcfg
	nic=$1

	wait_for_meta

	dns_servers="$(gcp_dns_servers)"
	dns_suffixes="$(gcp_dns_suffixes)"

	if [ "${dns_servers}" == "${DNS_DATA_NOT_FOUND}" ] && [ "${dns_suffixes}" == "${DNS_DATA_NOT_FOUND}" ]; then
		return
	fi

	if [ -z "${nic}" ] ; then
		nic=`ls /sys/class/net/ | grep -v lo | awk '{print $1}'`
	fi
	ifcfg=`get_network_config_path`/ifcfg-${nic}

	if [ -z "${dns_servers}" ] && [ -z "${dns_suffixes}" ]; then
		enable_resolv_conf_generation ${ifcfg}
	else
		disable_resolv_conf_generation ${ifcfg}
	fi

	managed=0
	if [ -e ${NETCONFIG} ] ; then
		managed=1
		update_netconfig "${dns_servers}" "${dns_suffixes}"
	elif [ -L ${RESOLV_CONF} ] ; then
		if [ -e ${SYSTEMD_RESOLVED} ] ; then
			managed=1
			update_resolved "${dns_servers}" "${dns_suffixes}"
		fi
		if command -v resolvconf >/dev/null ; then
			managed=1
			configure_resolvconf "${dns_servers}" "${dns_suffixes}"
		fi
	fi

	if [ ${managed} -eq 0 ] ; then
		backup_file ${RESOLV_CONF} resolv_conf
		generate_resolv_conf ${RESOLV_CONF} "${dns_servers}" "${dns_suffixes}"
	fi
}

restore_dns_config() {
	rm -f ${NM_CONFD}/20-nodns.conf
	sed -i -e '/#MIGRATE-FOR-GCE-GCPDNS/,/#END-MIGRATE-FOR-GCE-GCPDNS/c\' ${DHCLIENT_UP_HOOKS}

	restore_original_file resolv_conf ${RESOLV_CONF}
	restore_original_file netconfig ${NETCONFIG}
	restore_original_file systemd_resolved ${SYSTEMD_RESOLVED}

	if [ -f ${BACKUP_DIR}/resolvconf_exec ] ; then
		execute_resolvconf ${BACKUP_DIR}/resolvconf_exec
	fi

	systemctl try-restart systemd-resolved.service

	rm -rf ${BACKUP_DIR}.executed
	mv -f ${BACKUP_DIR} ${BACKUP_DIR}.executed
}