#!/bin/bash

DIR=$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )
ACTION=$1

# Functions
function abort() {
    echo "$1" >&2 && exit 1
}

function cn_config() {
    # /opt/cn-config/cn-config.ldif is managed by kubernetes configmap
    configsrc=/opt/cn-config/cn-config.ldif
    if [ ! -f $configsrc ]; then
        abort "$configsrc does not exist. Exiting."
    fi
    # always re-install config from latest repp
    rm -rf ${LDAP_CONF_DIR}/*}
    echo -e "slapadd -v -F ${LDAP_CONF_DIR} -bcn=config -l ${configsrc}"
    slapadd -v -F ${LDAP_CONF_DIR} -bcn=config -l ${configsrc}
    # test data load success
    if [ "$?" -ne 0 ]; then
        # remove restored data since config faulty, forcing reload on next boo
        rm -rfv ${LDAP_MDB_DIR}/*.mdb
        abort "Config data load did not succeed. Exiting."
  fi
}

function mdb_load() {
    echo -e "Assessing data state for ldap role: ${LDAP_ROLE}"
    file_data=$(ls -t ${LDAP_BAK_DIR}/mdb/data.mdb* 2>/dev/null | head -n 1)

    if [[ ( ! -f "$file_data" ) ]] ; then
        abort "No mdb snapshot found in ${LDAP_BAK_DIR}/mdb. Exiting."
    fi

    # Get latest data.mdb in backup
    bakData=$(ls -t ${LDAP_BAK_DIR}/mdb/data.mdb.[[:digit:]]* | head -n 1)
    echo "Latest mdb file in backup: $bakData"
    # Restore data and config if LDAP_FORCE_RESTORE is true or newer backup exists or size is smaller then usual.
    echo "Comparing timestamps of $bakData and ${LDAP_MDB_DIR}/data.mdb"
    if [[ ${LDAP_FORCE_RESTORE} = "true" ]] || [[ ${bakData} -nt ${LDAP_MDB_DIR}/data.mdb ]] \
        || du -sh ${LDAP_MDB_DIR}/data.mdb | grep -q 'M';
    then
        if [[ ${LDAP_FORCE_RESTORE} = "true" ]]; then
          echo "Force restore from ${bakData}"
        else
          echo -e "Backup data is newer than ${LDAP_MDB_DIR}/data.mdb.\nLoading most recent data and config..."
        fi
        find ${LDAP_MDB_DIR} -type f -name '*.mdb' -delete
        stat --format "Restoring snapshot '%n' modified on %y" ${bakData}
        cp -v ${bakData} ${LDAP_MDB_DIR}/data.mdb
        if [[ ${LDAP_ROLE} = "master" ]]; then
            bakDataAccessLog=$(ls -t ${LDAP_BAK_DIR}/mdb/accesslog/accesslog.mdb.[[:digit:]]* | head -n 1)
            # restore snapshot of latest accesslog mdb iff it corresponds to most recent data snapshot,
            # ie, it is tagged with the same timestamp, eg, <filename>.mdb.<timestamp>
            if [[ ${bakData##*.} = ${bakDataAccessLog##*.} ]]; then
                echo -e "Loading accesslog corresponding to most recent data snapshot"
                stat --format "restoring snapshot of accesslog '%n' modified on %y" ${bakDataAccessLog}
                cp -v ${bakDataAccessLog} ${LDAP_MDB_DIR}/accesslog/data.mdb
            fi
        fi
    else
        echo "${LDAP_MDB_DIR}/data.mdb is most recent data available."
    fi
}

function mdb_dump() {

    if [ ! $(which /usr/bin/mdb_copy) ]; then
        abort "mdb_copy not installed. Exiting."
    fi

    timestamp=$(date "+%Y%m%d%H%M%S")
    tmpdir=/tmp/$timestamp
    mkdir -p $tmpdir
    # data mdb
    # Clean up old data, just in case
    rm -rf ${LDAP_BAK_DIR}/mdb/data.mdb ${LDAP_BAK_DIR}/mdb/accesslog/data.mdb
    echo "Performing snapshot using mdbcopy"
    if ! /usr/bin/mdb_copy ${LDAP_MDB_DIR} ${LDAP_BAK_DIR}/mdb; then
        abort "mdb_copy mdb failed. Existing."
    else
        mv ${LDAP_BAK_DIR}/mdb/data.mdb ${LDAP_BAK_DIR}/mdb/data.mdb.$timestamp
        ls -lt ${LDAP_BAK_DIR}/mdb/data.mdb.${timestamp}
    fi

    # accesslog mdb
    if [[ ${LDAP_ROLE} = "master" ]]; then
        echo "Performing accesslog snapshot using mdbcopy"
        if ! /usr/bin/mdb_copy ${LDAP_MDB_DIR}/accesslog ${LDAP_BAK_DIR}/mdb/accesslog; then
            abort "mdb_copy accesslog failed. Existing."
        else
            mv ${LDAP_BAK_DIR}/mdb/accesslog/data.mdb \
                    ${LDAP_BAK_DIR}/mdb/accesslog/accesslog.mdb.${timestamp}
            ls -lt  ${LDAP_BAK_DIR}/mdb/accesslog/accesslog.mdb.${timestamp}
        fi
    fi

    # config
    echo "Performing config dump using slapcat"
    slapcat -F /etc/ldap/slapd.d -b cn=config > ${tmpdir}/config-${LDAP_ROLE}.ldif.$timestamp
    # ensure slapcat issued no errors before copying config to backups
    if [ "$?" -ne 0 ]; then
        abort "Config dump did not succeed. Exiting."
    fi
    cp -v ${tmpdir}/config-${LDAP_ROLE}.ldif.$timestamp ${LDAP_BAK_DIR}/config
    # clean up
    rm -Rf ${tmpdir}
}

function ldif_dump() {

    echo -e "Performing ldif dump for ${LDAP_ROLE} database"

    # exit if backup directories not present
    if [ ! -d ${LDAP_BAK_DIR}/ldif ]; then
        abort "ldif dump directory not present: ${LDAP_BAK_DIR}/ldif"
    fi
    # wait for slapd process to exit before backing up
    j=0
    while pgrep -x slapd; do
        sleep 3
        j=$(($j+1))
        if [ $j -gt 40 ]; then
            abort "slapd shutdown is hung so backup cannot continue. Exiting."
        fi
    done
    # dump and zip
    timestamp=$(date "+%Y%m%d%H%M%S")
    slapcat -v -b ${LDAP_BASE} > ${LDAP_BAK_DIR}/ldif/${LDAP_ROLE}-${timestamp}.db-ldif
    # ensure slapcat issued no errors
    if [ "$?" -ne 0 ]; then
        abort "Database dump did not succeed. Exiting."
    fi
    gzip -v ${LDAP_BAK_DIR}/ldif/${LDAP_ROLE}-${timestamp}.db-ldif
}

function prune_dumpdir() {

    # Clean up staled backup files (no timestamp suffix)
    rm -rf ${LDAP_BAK_DIR}/mdb/data.mdb ${LDAP_BAK_DIR}/mdb/accesslog/accesslog.mdb

    files=($(ls -t ${LDAP_BAK_DIR}/mdb/data.mdb*))
    file_count=${#files[@]}
    trim_num=$(expr ${file_count} - ${MDB_BAK_HISTORY})

    if [ $trim_num -gt 0 ]; then
        for (( i=$MDB_BAK_HISTORY; i<=$(( $file_count -1 )); i++ ))
        do
            #echo -e "deleting file: ${files[$i]}\n"
            rm -v ${files[$i]}
        done
    else
        echo "No files to remove in ${LDAP_BAK_DIR}/mdb"
    fi

    files=($(ls -t ${LDAP_BAK_DIR}/mdb/accesslog/accesslog.mdb*))
    file_count=${#files[@]}
    trim_num=$(expr ${file_count} - ${MDB_BAK_HISTORY})

    if [ $trim_num -gt 0 ]; then
        for (( i=$MDB_BAK_HISTORY; i<=$(( $file_count -1 )); i++ ))
        do
            #echo -e "deleting file: ${files[$i]}\n"
            rm -v ${files[$i]}
        done
    else
        echo "No files to remove in ${LDAP_BAK_DIR}/mdb/accesslog"
    fi

    files=($(ls -t ${LDAP_BAK_DIR}/config/*.ldif*))
    file_count=${#files[@]}
    trim_num=$(expr ${file_count} - ${CONF_BAK_HISTORY})

    if [ $trim_num -gt 0 ]; then
        for (( i=$CONF_BAK_HISTORY; i<=$(( $file_count -1 )); i++ ))
        do
            #echo -e "deleting file: ${files[$i]}\n"
            rm -v ${files[$i]}
        done
    else
        echo "No files to remove in ${LDAP_BAK_DIR}/config"
    fi
}

function krenew() {
  echo "renewing kerberos tgt every $renew_secs seconds"
  # sleep $renew_secs before first tkt renewal, since the initContainer obtains initial tkt
  sleep $renew_secs
  while true
  do
    kinit_svc
    sleep $renew_secs
  done
}

function kinit_svc() {

    ccfile=${KRB5CCNAME#*:}

    # Make the principal be the first component of HOSTNAME followed by
    # 'stanford.edu'.

    # Example 1. If HOSTNAME is 'ldap-test-smaster.cluster.local' then the
    # principal will be 'ldap-test-smaster.stanford.edu'.

    # Example 2. If HOSTNAME is 'ldap-test-sh.stanford.edu' then the
    # principal will be 'ldap-test-sh.stanford.edu'.
    principal=${HOSTNAME%%.*}.${REALM}
    kinit -k -t $KRB5_KTNAME -c ${ccfile} ldap/${principal}
    klist
}

# MAIN
if [ -e /usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 ]
then
  export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4"
fi

# ensure data dump dirs present
mkdir -p ${LDAP_BAK_DIR}/mdb ${LDAP_BAK_DIR}/mdb/accesslog ${LDAP_BAK_DIR}/config ${LDAP_MDB_DIR}/accesslog

case "${ACTION}" in
  dump)
    # Only do snaphost on one ldap server
    pod_last_number=$(hostname -s | rev | cut -d "-" -f 1)
    while true
    do
      sleep ${MDB_BAK_FREQ}
      # Don't do dump if no-backup file exists
      if [ -f ${LDAP_BAK_DIR}/no-backup ]; then
        echo "${LDAP_BAK_DIR}/no-backup exists. Skip backup"
      elif [ "$pod_last_number" = "0" ]; then
        echo "making mdb_copy every ${MDB_BAK_FREQ} seconds"
        echo "MDB backup stared."
        mdb_dump && prune_dumpdir
        echo "MDB backup ended."
        next_backup=$(date -d "+$MDB_BAK_FREQ seconds")
        echo "Next backup will start at $next_backup."
      fi
    done
    ;;

  ldifdump)
    ldif_dump
    ;;

  kinit)
    kinit_svc
    ;;

  krenew)
    shift
    renew_secs=$1
    krenew
    ;;

  cn_config)
    cn_config
    ;;

  datapop)
    cn_config && mdb_load
    ;;

  ldap)
    # Default slapd url
    endpoints="ldap:///"
    if [[ "$ENABLE_SIMPLE" = "yes" ]]; then
      endpoints="ldap:/// ldaps:///"

      # Simple bind needs saslauthd.
      # Put the process PID in /var/run/saslauthd.
      KRB5_KTNAME=/etc/krb5.keytab /usr/sbin/saslauthd -O /etc/saslauthd.conf \
        -a kerberos5 -c -m /var/run/saslauthd -n 5
    fi
    if ! /usr/sbin/slapd -h "${endpoints}" -F /etc/ldap/slapd.d -d ${LDAP_LOG_LEVEL} > /dev/stdout 2>&1; then
        echo "slapd failed to start. Wipe out local copy, forcing reload on next start."
        rm -rfv ${LDAP_MDB_DIR}/*.mdb ${LDAP_MDB_DIR}/accesslog/*.mdb
    fi
    ;;

  stop)
    if [[ -f /var/run/saslauthd/saslauthd.pid ]]; then
      # Stop saslauthd
      kill -9 `cat /var/run/saslauthd/saslauthd.pid`
    fi

    pid=$(cat /var/run/slapd.pid 2>/dev/null || pidof slapd)
    retry=TERM/20/KILL/forever
    if [ -n $pid ]; then
      start-stop-daemon --stop --retry ${retry} --pid ${pid} 2>&1
    else
      start-stop-daemon --stop --retry ${retry} --exec /usr/sbin/slapd 2>&1
    fi
    ;;

  *)
    abort "Usage: $0 (ldap|dump|ldifdump|krenew|kinit|cn_config)"
    ;;
esac
