#!/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
HIMDSD_SERVICE=${DAEMON_CONFIG_DIR}/himdsd.service
GCD_SERVICE=${DAEMON_CONFIG_DIR}/gcd.service
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

    if [ ! -f $1 ]; then
        echo "Error:  file $1 does not exist"
        exit 1
    fi

    echo "Removing proxy environment variable from file:  $1"

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

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

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

    if [ ! -f $1 ]; then
        echo "Error:  file $1 does not exist"
        exit 1
    fi

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

    sed -i "/\[Service\]/aEnvironment=${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

remove_proxy_from_daemon ${HIMDSD_SERVICE}
remove_proxy_from_daemon ${GCD_SERVICE}
remove_proxy_from_aha

if [ ${command} -eq 1 ]; then
   add_proxy_to_daemon ${HIMDSD_SERVICE}
   add_proxy_to_daemon ${GCD_SERVICE}   
   add_proxy_to_aha
fi

systemctl daemon-reload
systemctl restart himdsd
systemctl restart gcd
