#!/bin/bash
#
# Copyright (c) Microsoft Corporation.
#
# azcmagent_proxy, a script to set/remove proxy configuration for azcmagent
#
# Note #1:  this script must be run as root
# Note #2:  this script only works for systemd platforms for now

DAEMON_CONFIG_DIR=/lib/systemd/system
DAEMON_CONFIG_DIR2=/usr/lib/systemd/system
HIMDSD_SERVICE=himdsd.service
GCAD_SERVICE=gcad.service
EXTD_SERVICE=extd.service

DAEMON_GLOBAL_CONFIG_DIR=/lib/systemd/system.conf.d
PROXY_CONF=${DAEMON_GLOBAL_CONFIG_DIR}/proxy.conf
PROXY_ENV=https_proxy
AHA_PATH=/opt/azcmagent/bin/azcmagent

function print_usage
{
    echo "Usage:  azcmagent_proxy add <URL> - to add URL as the proxy"
    echo "        azcmagent_proxy remove - to delete configured proxy"
    exit 1
}

# One argument, the location of daemon service file
function remove_proxy_from_daemon ()
{
    if [ $# -ne 1 ]; then
        echo "remove_proxy() must be called with one argument"
        exit 1
    fi

    UNIT_FILE=${DAEMON_CONFIG_DIR}/$1
    if [ ! -f ${UNIT_FILE} ]; then
	UNIT_FILE=${DAEMON_CONFIG_DIR2}/$1
	if [ ! -f ${UNIT_FILE} ]; then
            echo "Warning:  unit file $1 does not exist"
            return 0
	fi
    fi

    grep -q "Environment=${PROXY_ENV}=" ${UNIT_FILE}
    if [ $? -ne 0 ]; then
        return 0
    fi

    sed -i "/Environment=${PROXY_ENV}=/d" ${UNIT_FILE}
}

function remove_global_proxy_from_daemon ()
{
    if [ $# -ne 1 ]; then
        echo "remove_proxy() must be called with one argument"
        exit 1
    fi

    if [ ! -f $1 ]; then
        echo "    No proxy previously configured"
        return 0
    fi

    echo "Removing proxy environment file:  $1"

    rm -f $1
}

# One argument, the location of daemon conf file
function add_global_proxy_for_daemon ()
{
    if [ $# -ne 1 ]; then
        echo "add_proxy() must be called with one argument"
        exit 1
    fi

    if [ -f $1 ]; then
        echo "Error:  Proxy already configured"
        exit 1
    fi

    echo "Setting proxy environment variable to file:  $1"

    echo "# Proxy setting for all systemd daemons, created by /opt/azcmagent/bin/azcmagent_proxy" > $1
    echo "" >> $1
    echo "[Manager]" >> $1
    echo "" >> $1
    echo "DefaultEnvironment=\"${PROXY_ENV}=${url}\"" >>$1
}

function remove_proxy_from_aha ()
{
    if [ ! -f ${AHA_PATH} ]; then
        echo "Error:  file ${AHA_PATH} does not exist"
        exit 1
    fi

    echo "Removing proxy environment variable from file:  ${AHA_PATH}"

    grep -q "export ${PROXY_ENV}=" ${AHA_PATH}
    if [ $? -ne 0 ]; then
        echo "    No proxy previously configured"
        return 0
    fi

    sed -i "/export ${PROXY_ENV}=/d" ${AHA_PATH}
}

function add_proxy_to_aha ()
{
    echo "Adding proxy environment variable to file:  ${AHA_PATH}"

    sed -i "/ Environment Variables below ======/aexport ${PROXY_ENV}=${url}" ${AHA_PATH}
}

# Make sure we are running as root
if [ $EUID != 0 ]; then
   echo "$0 must be invoked as root!"
   exit 1
fi

if [ $# -ne 1 -a $# -ne 2 ]; then
   print_usage
fi

if [ $1 == "add" ]; then
   command=1
   url=$2
   if [ -z "${url}" ]; then
       print_usage
   fi
elif [ $1 == "remove" ]; then
   command=2
else
   print_usage
fi

# For backward compatibility, we always check to see if proxy settings should be
# removed from individual daemon settings
remove_proxy_from_daemon ${HIMDSD_SERVICE}
remove_proxy_from_daemon ${GCAD_SERVICE}
remove_proxy_from_daemon ${EXTD_SERVICE}

remove_global_proxy_from_daemon ${PROXY_CONF}
remove_proxy_from_aha

# On RHEL systems, must also explicitly remove the environment variable
systemctl unset-environment ${PROXY_ENV}

if [ ${command} -eq 1 ]; then
    add_global_proxy_for_daemon ${PROXY_CONF}
    add_proxy_to_aha
fi

systemctl daemon-reexec
systemctl restart himdsd
systemctl restart gcad
systemctl restart extd
