#!/bin/bash

set -e

###### Color Variables ######

RED='\033[0;31m'
GREEN='\033[1;32m'
NC='\033[0m'



###### General Variables ######

LOCATION_CONFIG="/etc/ngcp-usr-location/ngcp-usr-location.conf"
LOCATION_SCHEMA="/usr/share/kamailio/db_redis/kamailio/location"
DOMAIN=""
USERNAME=""
VERBOSE=false
BATCH=false
REDIS_LOC_PREFIX="1:location"
REDIS_CHUNK_SIZE="1000"


###### General Functions ######

#----------------------------------------------------------------------
# Helper function
usage() {
  echo "
Usage: $0 <option> [--verbose]

  options:
    --all [--batch]                 count the number of registered devices and subscribers in all the system (only the plain value with --batch)
    --domain <domain_name>          count the number of registered devices and subscribers per domain (ATTENTION: it could increase system load)
    --subscriber <username>         count the number of registered devices per subscriber (--username is allowed as well)
    --subscriber <username> <keys>  print the values of the specified keys for each device that belongs to the subscriber's username
    --remove <user@domain>          remove all the registrations from a subscriber
    --remove <user@domain> <ruid>   remove the specific registrations from a subscriber

    --verbose                       print the device's information.
                                    ATTENTION: it could increase system load if used together with --all or --domain options

"
}

#----------------------------------------------------------------------
# Ask for confirmation to proceed
confirmation() {
  while true ; do
    echo "ATTENTION: verbose option can increase the system load due to the high number of redis queries."
    echo "Are you sure you want to continue (yes/no)? "
    read -r choice
    case "${choice,,}" in
      yes) break;;
      no) echo "Exiting" && exit 0;;
      *) echo "Invalid choice. Enter yes or no, please";;
    esac
  done
}

#----------------------------------------------------------------------
# Check code injection in passed arguments
check_injection() {
  temporan=${1//;/}            #bash and mysql end of command
  temporan=${temporan//\\G/}   #mysql end of command
  temporan=${temporan//\\n/}   #new line
  temporan=${temporan//\\r/}   #carriage return
  echo "${temporan}"
}



###### REDIS Functions ######

#----------------------------------------------------------------------
# Query redis in non blocking chunks fashion (SSCAN command)
# result will be inside $REDIS_SSCAN_RESULT
redis_sscan() {

  # save old ifs and set it to newline (\n / 0x0a / 10)
  OLDIFS=$IFS
  IFS=$'\n'

  # reset cursor and output sscan result
  cursor="0"
  REDIS_SSCAN_RESULT=()

  # Iterate with sscan until cursor not equal 0
  while : ; do
    mapfile -t tmp_list < <(
      "${redis_command[@]}" sscan "$2" "$cursor" match "$1" count $REDIS_CHUNK_SIZE
    )
    cursor=${tmp_list[0]}

    # for each chunk we append new data to the final array except for the first element which contains the updated cursor
    REDIS_SSCAN_RESULT+=("${tmp_list[@]:1}")

    [[ $cursor != "0" ]] || break
  done

  #restore old ifs
  IFS=$OLDIFS
}

#----------------------------------------------------------------------
# Query redis in non blocking chunks fashion (SSCAN command)
# result will be inside $REDIS_COUNT_RESULT
redis_count() {

  # save old ifs and set it to newline (\n / 0x0a / 10)
  OLDIFS=$IFS
  IFS=$'\n'

  # reset cursor and output result
  cursor="0"
  REDIS_COUNT_RESULT="0"

  # Iterate with sscan until cursor not equal 0
  while : ; do
    mapfile -t tmp_list < <(
      "${redis_command[@]}" sscan "$2" "$cursor" match "$1" count $REDIS_CHUNK_SIZE
    )
    cursor=${tmp_list[0]}
    ((REDIS_COUNT_RESULT += ${#tmp_list[@]} - 1)) || true

    [[ $cursor != "0" ]] || break
  done

  #restore old ifs
  IFS=$OLDIFS
}

#----------------------------------------------------------------------
# Parse location schema to get columns definition
redis_location_schema() {
  if [ ! -s "${LOCATION_SCHEMA}" ] ; then
    echo "Warning: impossible to get redis db structure" >&2
    exit 1
  fi

  for element in $(head -1 "${LOCATION_SCHEMA}" | sed -e "s#/[a-z]\+,# #g") ; do
    loc_columns+=("${element}")
  done

  if [ "${#loc_columns[@]}" -eq 0 ] ; then
    echo "Warning: impossible to get redis db structure" >&2
    exit 1
  fi
}

#----------------------------------------------------------------------
# Parse username to get if domain is specified or not
redis_parse_username() {
  if [[ "${USERNAME}" =~ @ ]] ; then
    if  ${IGNORE_AUTH_REALM} ; then
      echo -e "Attention: kamailio.proxy.ignore_auth_realm is enabled, domain will be not considered in the query${NC}"
      search_key="${USERNAME/@*/}"
    else
      search_key="${USERNAME/@/:}"
    fi
  else
    if  ${IGNORE_AUTH_REALM} ; then
      search_key="${USERNAME}"
    else
      search_key="${USERNAME}:*"
    fi
  fi
}

#----------------------------------------------------------------------
# Parse and print device information
redis_print_device() {
  for device in "${device_list[@]}" ; do
    OLDIFS=$IFS
    IFS=$'\n'
    mapfile -t device_info < <(
      "${redis_command[@]}" hmget "${device}" "$@" \
        | awk '{ if ($0==""){print "NULL"} else {print $0} }'
    )
    IFS=$OLDIFS
    echo -e "********** Location row ${device} **********"
    index=0
    for info in "$@" ; do
      printf "%13s: %s\n" "${info}" "${device_info[$index]}"
      index=$((index+1))
    done
  done
}

#----------------------------------------------------------------------
# Count the number of registered devices and subscribers in all the system
# Print the device's information if verbose is activated
redis_action_all() {
  # If verbose is activated the list of devices has to be kept in memory
  if ${VERBOSE} ; then
    redis_sscan "$REDIS_LOC_PREFIX:entry*" "$REDIS_LOC_PREFIX:master"
    device_list=("${REDIS_SSCAN_RESULT[@]}")
    device_count="${#REDIS_SSCAN_RESULT[@]}"
  else
    redis_count "$REDIS_LOC_PREFIX:entry*" "$REDIS_LOC_PREFIX:master"
    device_count="${REDIS_COUNT_RESULT}"
  fi

  if [ "${device_count}" -eq 0 ] ; then
    if ${BATCH} ; then
      echo "0"
      exit 0
    else
      echo -e "${RED}No entries found${NC}" >&2
      exit 3
    fi
  fi

  if ${VERBOSE} ; then
    redis_print_device "${loc_columns[@]}"
  fi

  if ${BATCH} ; then
    echo "${device_count}"
    return
  else
    echo -e "${GREEN}There are ${device_count} registered devices${NC}"
  fi

  redis_count "$REDIS_LOC_PREFIX:usrdom*" "$REDIS_LOC_PREFIX::index::usrdom"
  user_count="${REDIS_COUNT_RESULT}"
  echo -e "${GREEN}There are ${user_count} users with at least one registered device${NC}"
}

#----------------------------------------------------------------------
# Count the number of registered devices and subscribers per domain
# Print the device's information if verbose is activated
redis_action_domain() {
  redis_sscan "$REDIS_LOC_PREFIX:usrdom::*:${DOMAIN}" "$REDIS_LOC_PREFIX::index::usrdom"
  subscriber_list=("${REDIS_SSCAN_RESULT[@]}")
  user_count="${#subscriber_list[@]}"
  if [ "${user_count}" -eq 0 ] ; then
    echo -e "${RED}No entries found for domain ${DOMAIN}${NC}" >&2
    exit 3
  fi

  for subscriber in "${subscriber_list[@]}" ; do

    mapfile -t device_list < <(
      "${redis_command[@]}" smembers "$subscriber" | awk '{print $1}'
    )
    device_count=$((device_count+"${#device_list[@]}"))

    if ${VERBOSE} && [ "${#device_list[@]}" -gt 0 ] ; then
      redis_print_device "${loc_columns[@]}"
    fi
  done

  echo -e "${GREEN}There are ${device_count} registered devices in domain ${DOMAIN}${NC}"
  echo -e "${GREEN}There are ${user_count} users with at least one registered device in domain ${DOMAIN}${NC}"
}

#----------------------------------------------------------------------
# Count the number of registered devices per subscriber
# Print the device's information if verbose is activated
redis_action_subscriber() {
  redis_sscan "$REDIS_LOC_PREFIX:usrdom::${search_key}" "$REDIS_LOC_PREFIX::index::usrdom"
  subscriber_list=("${REDIS_SSCAN_RESULT[@]}")
  if [ "${#subscriber_list[@]}" -eq 0 ] ; then
    echo -e "${RED}No entries found for subscriber ${USERNAME}${NC}" >&2
    exit 3
  fi
  for subscriber in "${subscriber_list[@]}" ; do
    mapfile -t device_list < <(
      "${redis_command[@]}" smembers "${subscriber}" | awk '{print $1}'
    )
    device_count=$((device_count+"${#device_list[@]}"))

    if [ "${#device_list[@]}" -gt 0 ] ; then
      if ${VERBOSE} ; then
        redis_print_device "${loc_columns[@]}"
        echo ""

      elif [ "${#KEYS[@]}" -gt 0 ] ; then
        redis_print_device "${KEYS[@]}"
      fi
    fi
    echo -e "${GREEN}There are ${device_count} registered devices in subscriber ${subscriber/$REDIS_LOC_PREFIX:usrdom::/}${NC}"
  done
}

#----------------------------------------------------------------------
# Remove all the registrations from a subscriber
redis_action_remove_all() {
  mapfile -t subscriber_list < <(
    "${redis_command[@]}" keys "$REDIS_LOC_PREFIX:usrdom::${search_key[*]}" \
      | awk '{print $1}'
  )
  if [ "${#subscriber_list[@]}" -eq 0 ] ; then
    echo -e "${RED}No entries found for subscriber ${USERNAME}${NC}" >&2
    exit 3
  fi
  for subscriber in "${subscriber_list[@]}" ; do
    error=false

    # First remove all the location entries associated to that subscriber
    mapfile -t device_list < <(
      "${redis_command[@]}" smembers "${subscriber}" | awk '{print $1}'
    )
    for device in "${device_list[@]}" ; do
      remove_device=$("${redis_command[@]}" del "${device}" | awk '{print $1}')
      if [ "${remove_device}" -eq 0 ] ; then
        echo -e "${RED}Warning: error deleting device with key ${device} on subscriber ${subscriber}${NC}" >&2
        error=true
      fi
    done

    # If there were no errors, remove the subscriber entry
    if ! ${error} ; then
      remove_subscriber=$("${redis_command[@]}" del "${subscriber}" | awk '{print $1}')
      if [ "${remove_subscriber}" -eq 0 ] ; then
        echo -e "${RED}Warning: error deleting subscriber with key ${subscriber}${NC}" >&2
        error=true
      fi
    fi
  done

  if ! ${error} ; then
    echo ""
    echo -e "${GREEN}Deleted all devices on subscriber ${USERNAME}${NC}"
  fi
}

#----------------------------------------------------------------------
# Remove specific registrations from a subscriber
redis_action_remove_records() {
  for device in "${KEYS[@]}" ; do
    # Check if it is the correct subscriber
    username_value=$("${redis_command[@]}" hget "$REDIS_LOC_PREFIX:entry::${device}" "username" | awk '{print $1}')
    if [ "${username_value}" != "${USERNAME/@*/}" ] ; then
      echo -e "${RED}Warning: device with key ${device} doesn't belong to subscriber ${USERNAME}${NC}" >&2
      continue
    fi

    # Remove the device from the subscriber entry list
    remove_device=$("${redis_command[@]}" del "$REDIS_LOC_PREFIX:entry::${device}" | awk '{print $1}')
    if [ "${remove_device}" -eq 0 ] ; then
      echo -e "${RED}Warning: error deleting device with ruid ${device} on subscriber ${USERNAME}${NC}" >&2
      exit 3
    fi

    # Remove the device entry
    remove_device=$("${redis_command[@]}" srem "$REDIS_LOC_PREFIX:usrdom::${search_key[*]}" "$REDIS_LOC_PREFIX:entry::${device}" | awk '{print $1}')
    if [ "${remove_device}" -eq 0 ] ; then
      echo -e "${RED}Warning: error deleting device with ruid ${device} on subscriber ${USERNAME}${NC}" >&2
      exit 3
    fi

    echo -e "${GREEN}Deleted device with ruid ${device} from subscriber ${USERNAME}${NC}"
  done

  # Check if the subscriber entry is empty
  device_count=$("${redis_command[@]}" scard "$REDIS_LOC_PREFIX:usrdom::${search_key[*]}" | awk '{print $1}')
  if [ "${device_count}" -eq 0 ] ; then
    echo ""
    echo -e "${GREEN}No more entries for subscriber ${USERNAME}${NC}"
  fi
}



###### Main Function ######

# load config
if ! [ -r "${LOCATION_CONFIG}" ] ; then
  echo "Warning: cannot read ${LOCATION_CONFIG}" >&2
  exit 1
fi
# shellcheck disable=SC1090
. ${LOCATION_CONFIG}

REDIS_CMD="$(which keydb-cli)"
redis_command=("${REDIS_CMD}" -h "${REDIS_IP}" -p "${REDIS_PORT}" -n "${REDIS_LOC_DB}")
redis_location_schema


device_count=0
user_count=0


# parse input argument
case "$1" in
    --all|-a)
      if [ "$2" == "--verbose" ] ; then
        VERBOSE=true
        confirmation
      fi
      if [ "$2" == "--batch" ] ; then
        BATCH=true
      fi
      redis_action_all
      ;;

    --domain|-d)
      DOMAIN="${2,,}"
      if [ -z "${DOMAIN}" ] ; then
        echo "Warning: domain option requires a domain name" >&2
        exit 1
      fi
      if ${IGNORE_AUTH_REALM} ; then
        echo -e "${RED}Impossible to get data because kamailio.proxy.ignore_auth_realm is enabled${NC}" >&2
        exit 1
      fi
      DOMAIN=$(check_injection "${DOMAIN}")
      if [ "$3" == "--verbose" ] ; then
        VERBOSE=true
        confirmation
      fi
      redis_action_domain
      ;;

    --subscriber|-s|--username|-u)
      USERNAME="${2,,}"
      if [ -z "${USERNAME}" ] ; then
        echo "Warning: subscriber option requires an username or username@domain" >&2
        exit 1
      fi
      USERNAME=$(check_injection "${USERNAME}")
      if [ "$3" == "--verbose" ] ; then
        VERBOSE=true
      elif [ "$3" != "" ] ; then
        shift
        shift
        KEYS=()
        for argument in "$@" ; do
          mapfile -t ADD_KEYS < <(check_injection "${argument}")
          KEYS+=("${ADD_KEYS[@]}")
        done
      fi
      redis_parse_username
      redis_action_subscriber
      ;;

    --remove|-r)
      USERNAME="${2,,}"
      if [ -z "${USERNAME}" ] ; then
        if ${IGNORE_AUTH_REALM} ; then
          echo "Warning: remove option requires an username" >&2
          exit 1
        else
          echo "Warning: remove option requires an username@domain" >&2
          exit 1
        fi
      fi
      if ! [[ "${USERNAME}" =~ @ ]] && ! ${IGNORE_AUTH_REALM} ; then
        echo "Warning: remove option requires an username@domain" >&2
        exit 1
      fi
      USERNAME=$(check_injection "${USERNAME}")
      if [ "$3" != "" ] ; then
        shift
        shift
        KEYS=()
        for argument in "$@" ; do
          mapfile -t ADD_KEYS < <(check_injection "${argument}")
          KEYS+=("${ADD_KEYS[@]}")
        done
        redis_parse_username
        redis_action_remove_records
      else
        redis_parse_username
        redis_action_remove_all
      fi
      ;;

    --help|-h)
      usage
      exit 0
      ;;

    *)
      echo "Warning: please verify the script usage" >&2
      usage
      exit 1
      ;;
  esac


# exit
exit 0
