#!/bin/bash
#
PATH="${PATH}:/usr/local/bin:/usr/bin:/usr/local/sbin:/usr/sbin"
#
if [ -f /usr/sbin/service ]; then
    SERVICE_COMMAND=/usr/sbin/service
else
    SERVICE_COMMAND=/sbin/service
fi

QUIET=0
CONF_FORMAT="yaml"
AUTH_SECTION_EXISTS=0

SCRIPT_PWD=$(cd `dirname $0` && pwd)
MONGO_CLIENT_BIN="${SCRIPT_PWD}/mongo"

if [ "${SCRIPT_PWD}" = "/usr/bin" ]; then
    TARBALL=0
else
    TARBALL=1
fi

usage() { 
    echo "Usage: `basename $0` [-h] [-q] [-c]"
    echo
    echo "  This script is used to enable authentication in Percona Server for MongoDB automatically."
    echo "  User 'dba' with role 'root' will be created with autogenerated password and the service will be restarted."
    echo "  -q - quiet mode. Do not ask any questions and enable auth automatically. For tarball installation it is not used."
    echo "  -c - provide path to config file. If -c is not set the default path will be used."
    echo
    exit 0
}
while getopts ":hqc:" arg; do
  case $arg in
    c) # Specify CONF value.
      CONF=${OPTARG}
      ;;
    q) # quiet mode.
      QUIET=1
      ;;
    h) # Display help.
      usage
      exit 0
      ;;
    *)
      ;;
  esac
done

[ -z "${CONF}" ] && CONF=/etc/mongod.conf

if [ "$EUID" -ne 0 ]
  then echo "Please run as root"
  exit
fi

parse_yaml() {
   local s='[[:space:]]*' w='[a-zA-Z0-9_]*' fs=$(echo @|tr @ '\034')
   sed -ne "s|^\($s\)\($w\)$s:$s\"\(.*\)\"$s\$|\1$fs\2$fs\3|p" \
        -e "s|^\($s\)\($w\)$s:$s\(.*\)$s\$|\1$fs\2$fs\3|p"  $1 |
   awk -F$fs '{
      indent = length($1)/2;
      vname[indent] = $2;
      for (i in vname) {if (i > indent) {delete vname[i]}}
      if (length($3) > 0) {
         vn=""; for (i=0; i<indent; i++) {vn=(vn)(vname[i])("_")}
         printf("%s%s=\"%s\"\n", vn, $2, $3);
      }
   }'
}

get_value_from_yaml() {
    section=$1
    param=$2
    array=$(parse_yaml $CONF)
    result=0
    while IFS=' ' read -ra VALUES; do
        for value in "${VALUES[@]}"; do
            if [[ $value =~ ([^,]+).*"="\"([^,]+)\" ]]; then
                name=${BASH_REMATCH[1]}
                val=${BASH_REMATCH[2]}
            fi
            if [[ $name =~ $section && $name =~ $param ]]; then
                result=$val
                break
            fi
        done
    done <<< "$array"
    echo $result
}

add_value_to_yaml() {
    section=$1
    param=$2
    value=$3
    if [[ $CONF_FORMAT == "yaml" ]]; then
        if [[ $AUTH_SECTION_EXISTS == 1 ]]; then
            sed -i "s/authorization: disabled/authorization: enabled/" $CONF
        else
            secutity_line=$(grep "security:" $CONF)
            if [[ $secutity_line =~ '#security:' ]]; then
                auth_line=$(grep "authorization:" $CONF)
                regex='#  authorization:'
                if [[ $auth_line =~ $regex ]]; then
                    sed -i "s/#security:/security:/" $CONF
                    sed -i "s/#  authorization:.*$/  authorization: enabled/" $CONF
                else
                    sed -i "s/#security:/security:\n  authorization: enabled/" $CONF
                fi
            elif [[ $secutity_line != '' ]]; then
                auth_line=$(grep "authorization:" $CONF)
                regex='#  authorization:'
                if [[ $auth_line =~ $regex ]]; then
                    sed -i "s/#  authorization:.*$/  authorization: enabled/" $CONF
                else
                    sed -i "s/security:/security:\n  authorization: enabled/" $CONF
                fi
            else
                delim=$( grep '## Enterprise-Only Options:' $CONF )
                if [[ $delim == '' ]]; then
                    echo "security:" >> $CONF
                    echo "  authorization: enabled" >> $CONF
                else
                    sed -i "s/## Enterprise-Only Options:/security:\n  authorization: enabled\n## Enterprise-Only Options:/" $CONF
                fi
            fi
        fi
    fi
}

add_user_to_mongo() {
    port="$(get_value_from_yaml net port)"
    user="dba"
    password="$(tr -dc 'a-zA-Z0-9' < /dev/urandom | head -c 32)"
    password="${password%\\n}"
    echo "db.createUser({user: \"$user\", pwd: \"$password\", roles: [ \"root\" ] });" | ${MONGO_CLIENT_BIN} -p $port localhost/admin
    if [ $? -eq 0 ];then
        echo -e "User has been created successfully!\nUser:${user}\nPassword:${password}"
    else
        echo "ERROR! User cannot be added!"
        rm -rf /tmp/mongodb_create.lock
        exit 1
    fi
}

if [ ! -f /tmp/mongodb_create.lock ]; then
    AUTH_ENABLED=0
    auth_res=$(get_value_from_yaml security auth)
    if [[ $auth_res == enabled  ]]; then
        AUTH_ENABLED=1
    elif [[ $auth_res == disabled  ]]; then
        AUTH_ENABLED=0
        AUTH_SECTION_EXISTS=1
    elif [[ `egrep '^auth=1' $CONF` ]]; then
        AUTH_ENABLED=1
        CONF_FORMAT="conf"
    elif [[ `egrep '^auth=0' $CONF` ]]; then
        AUTH_ENABLED=0
        AUTH_SECTION_EXISTS=1
    fi
    if [[ $AUTH_ENABLED == 0 ]]; then
        echo "We have detected authentication is not enabled."
        echo "Would you like help creating your first user?"
        if [[ $QUIET -eq 0 ]]; then
            echo "Please note that mongodb service could be restarted during this action"
            read -t 15 -p "Would you like to proceed?(Y/n)" setup_auth
        else
            setup_auth="y"
        fi
        if [ "$setup_auth" == "Y" -o "$setup_auth" == "y" ]; then
            touch /tmp/mongodb_create.lock
            started=$(pgrep mongod | wc -l)
            if [ $started == 0 ]; then
                if [ $TARBALL == 0 ]; then
                    $SERVICE_COMMAND mongod start
                    if [ "$?" != 0 ]; then
                        echo "Cannot start mongod service. Please check logs and start mongod manually."
                        rm -rf /tmp/mongodb_create.lock
                        exit 1
                    fi
                else
                    echo "Mongod service is stopped. Please start it and re-run script."
                    rm -rf /tmp/mongodb_create.lock
                    exit 1
                fi
            fi
            add_user_to_mongo
            add_value_to_yaml security authentication true
            if [ $TARBALL == 1 ]; then
                echo "It is needed to restart mongod service."
                echo "Once service is restarted authentication will be enabled"
                rm -rf /tmp/mongodb_create.lock
                exit 0
            else
                $SERVICE_COMMAND mongod stop
                if [ "$?" != 0 ]; then
                    PID=$(pgrep mongod)
                    kill -9 $PID
                fi
                rm -rf /tmp/mongodb_create.lock
                $SERVICE_COMMAND mongod start
            fi
        elif [ "$setup_auth" == "N" -o "$setup_auth" == "n" ]; then
            exit 0
        else
            echo "Incorrect input. You should press Y or N"
            exit 1
        fi
    else
        echo "Authentication is already enabled"
    fi
fi
