#!/bin/bash
#
# Copyright (C) 2018 VMware, Inc. All rights reserved.
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA.
#

#
# SPDX-License-Identifier: GPL-2.0-only
#

#
# Configuration required for vmw_conn_notify
#

### BEGIN INIT INFO
# Provides:             vmw_conn_notifyd
# Required-Start:       $network $remote_fs $syslog $time
# Required-Stop:        $network $remote_fs $syslog $time
# Default-Start:        2 3 5
# Default-Stop:         0 1 6
# Short-Description:    VMware Network Connection Introspection Service
# Description:          Manages the service need to run Network Connecion Introspection
### END INIT INFO
shopt -s expand_aliases > /dev/null 2>&1
export _SYSTEMCTL_SKIP_REDIRECT=1

VMW_CONN_NOTIFY_NAME=vmw_conn_notify
VMW_CONN_NOTIFY_BIN=/usr/sbin/$VMW_CONN_NOTIFY_NAME
VMW_CONN_NOTIFY_CONFIG=/etc/vmw_conn_notify/vmw_conn_notify.cfg
LOCK_FILE=/var/lock/subsys/vmw_conn_notifyd
VMW_CONN_NOTIFY_STOP_TIME=2
VMW_CONN_NOTIFY_WATCHDOG_NAME=vmw_conn_notifyd_watchdog
VMW_CONN_NOTIFY_WATCHDOG=/usr/sbin/$VMW_CONN_NOTIFY_WATCHDOG_NAME

# GM_STATUS related variables
GM_SUCCESS=0
GM_FAIL=1

# logger variables
VERBOSE=1
SILENT=0
LOG_LEVEL=$VERBOSE
SUCCESS="info"
ERROR="err"
WARNING="warning"
INFO="info"

NFQUEUE_BYPASS="NFQUEUE --queue-num 0 --queue-bypass"
TCP_PACKET_FILTER_OP="-p tcp --tcp-flags FIN,SYN,RST,ACK,PSH SYN"
VNET_PACKET_STAMP="-m mark ! --mark 0x1/0x1"
VNET_CHAIN=vnetchain

if [ -x  '/sbin/iptables' ]; then
   IPTABLES="/sbin/iptables"
elif [ -x '/usr/sbin/iptables' ]; then
   IPTABLES="/usr/sbin/iptables"
else
   IPTABLES="iptables"
fi

if [ -x  '/sbin/ip6tables' ]; then
   IP6TABLES="/sbin/ip6tables"
elif [ -x '/usr/sbin/ip6tables' ]; then
   IP6TABLES="/usr/sbin/ip6tables"
else
   IP6TABLES="ip6tables"
fi

MODPROBE="/sbin/modprobe"

#####################################################################
# Definition:
# Function to log Success/Fail
#
# Arguments:
# Argument 1: Return code
# Argument 2: LOG_LEVEL
#####################################################################
eval_cmd() {
   local rc=$1
   local VERBOSE=1
   local SILENT=0
   local LOG_LEVEL=${2:-$VERBOSE}

   if [ "${LOG_LEVEL}" -eq $VERBOSE ]; then
     if [ "$rc" -eq "$GM_SUCCESS" ]; then
       echo '[ OK ]'
     else
       echo '[ FAILED ]'
     fi
   fi

   return "$rc"
}

#####################################################################
# Definition:
# Function to log messages
#
# Arguments:
# Argument 1: String to print
# Argument 2: LOG_LEVEL(SILENT(0) / VERBOSE(1)) (optional)
# Argument 3: Log Status (success|info, warning, failure|err)
#####################################################################
vmw_logger() {
   local VERBOSE=1
   local SILENT=0
   local STR=$1
   local LOG_LEVEL=${2:-$VERBOSE}
   local LOG_STATUS=${3:-"info"}

   if [ "${LOG_LEVEL}" -eq $VERBOSE ]; then
     echo "${STR}"
   fi

   case "$LOG_STATUS" in
   "err"|"failure") # failure
       LOG_STATUS="err"
       ;;
   "warn") # warning
       LOG_STATUS="warning"
       ;;
   "info"|"success") # success
       LOG_STATUS="info"
       ;;
   *)
       LOG_STATUS="info"
       ;;
   esac

   logger -p daemon.$LOG_STATUS -t "vmw_conn_notifyd[$$]" "$STR"

}

# Check for missing binaries
if [ ! -f $VMW_CONN_NOTIFY_BIN ] ; then
   vmw_logger "$VMW_CONN_NOTIFY_BIN not installed" $VERBOSE $ERROR
   if [ "$1" = "stop" ] ; then
      exit 0
   else
      exit 5
   fi
fi

#
# Enable this when we have a config file
# Check for missing config file
[ -e $VMW_CONN_NOTIFY_CONFIG ] && . $VMW_CONN_NOTIFY_CONFIG

# Determine OS type based on the functions library path
OS=""
if [ -f /etc/init.d/functions ] ; then
   OS="Redhat"
   #. /etc/init.d/functions

   MODPROBE_PARAM=""
elif [ -f /lib/lsb/init-functions ] && [ -f /etc/rc.status ] ; then
   OS="Suse"

   MODPROBE_PARAM="--allow-unsupported-modules"
elif [ -f /lib/lsb/init-functions ] ; then
   OS="Ubuntu"

   MODPROBE_PARAM=""
else
   OS="Unsupported OS, exiting."
   vmw_logger $OS $SILENT $ERROR
   exit 1
fi

# Execute the given command and bail-out on failure
exec_or_die() {
   status="0"
   cmd="$@"
   out=$($cmd 2>&1) || status="$?"
   if [ "$status" -ne 0 ]; then
      # Remove extra spaces if any because of indentation
      cmd=$(echo $cmd | /bin/sed 's/[[:space:]]/ /g')
      vmw_logger "Failure in executing \"$cmd\", error code: $status, \"$out\"" $SILENT $ERROR
      exit 4
   fi
}

# Execute the given command and log warning message on failure
exec_or_warn() {
   status="0"
   cmd="$@"
   out=$($cmd 2>&1) || status="$?"
   if [ "$status" -ne 0 ]; then
      # Remove extra spaces if any because of indentation
      cmd=$(echo $cmd | /bin/sed 's/[[:space:]]/ /g')
      vmw_logger "\"$cmd\" returns with code: $status, message: \"$out\"" $SILENT $ERROR
      return 4
   fi
}

# Load netfilter kernel modules
load_netfilter_modules() {
   exec_or_die ${MODPROBE} iptable_filter
   exec_or_die ${MODPROBE} ip6table_filter
   exec_or_die ${MODPROBE} xt_NFQUEUE
   exec_or_die ${MODPROBE} nf_conntrack_ipv4
   exec_or_die ${MODPROBE} nf_conntrack_ipv6
}

add_vnetchain_filter_rules() {
   exec_or_die "${IPTABLES} -N ${VNET_CHAIN}"
   exec_or_die "${IP6TABLES} -N ${VNET_CHAIN}"

   exec_or_die "${IPTABLES} -I INPUT ${VNET_PACKET_STAMP} \
                ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_die "${IPTABLES} -I OUTPUT ${VNET_PACKET_STAMP} \
                ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_die "${IPTABLES} -I ${VNET_CHAIN} -j ${NFQUEUE_BYPASS}"

   exec_or_die "${IP6TABLES} -I INPUT ${VNET_PACKET_STAMP} \
                ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_die "${IP6TABLES} -I OUTPUT ${VNET_PACKET_STAMP} \
                ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_die "${IP6TABLES} -I ${VNET_CHAIN} -j ${NFQUEUE_BYPASS}"
}

remove_vnetchain_filter_rules() {
   exec_or_warn "${IPTABLES} -D INPUT ${VNET_PACKET_STAMP} \
                 ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_warn "${IPTABLES} -D OUTPUT ${VNET_PACKET_STAMP} \
                 ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_warn "${IPTABLES} -D ${VNET_CHAIN} -j ${NFQUEUE_BYPASS}"

   exec_or_warn "${IP6TABLES} -D INPUT ${VNET_PACKET_STAMP} \
                 ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_warn "${IP6TABLES} -D OUTPUT ${VNET_PACKET_STAMP} \
                 ${TCP_PACKET_FILTER_OP} -j ${VNET_CHAIN}"
   exec_or_warn "${IP6TABLES} -D ${VNET_CHAIN} -j ${NFQUEUE_BYPASS}"

   exec_or_warn "${IPTABLES} -X ${VNET_CHAIN}"
   exec_or_warn "${IP6TABLES} -X ${VNET_CHAIN}"
}

if [ "`id -u`" -ne 0 ] ; then
   vmw_logger "User has insufficient privilege." $VERBOSE $ERROR
   exit 4
fi

#####################################################################
# Starts the watchdog service
#####################################################################
start_watchdog() {
   local retval="$GM_SUCCESS"

   # setsid creates a new session id for watchdog and
   # keeps it running even if shell session is closed
   # Watchdog is redundant when there is no systemd
   # But necessary when there is upstart or SysVinit
   setsid $VMW_CONN_NOTIFY_WATCHDOG >/dev/null 2>&1 < /dev/null &
   retval=$?
   return "$retval"
}

status_vmw_conn_notify() {

   # Check if there's already running instance of vmw_conn_notify, using lock_file
   # and using pidof command.. so at any time we will run only one instance
   #
   # This assumes that we use full pathname to start (/usr/sbin/vmw_conn_notify)
   local retval_vmw_conn_notify="$GM_SUCCESS"

   if [ -f $LOCK_FILE -a -n "`pidof $VMW_CONN_NOTIFY_BIN`" ] ; then
      retval_vmw_conn_notify="$GM_SUCCESS"
   else
      retval_vmw_conn_notify="$GM_FAIL"
   fi

   return "$retval_vmw_conn_notify"
}

#####################################################################
# Queries the watchdog service
#####################################################################
status_watchdog() {
   local retval="$GM_SUCCESS"
   local pids=$(pgrep -f $VMW_CONN_NOTIFY_WATCHDOG_NAME)

   if [ -n "$pids" ]; then
      retval="$GM_SUCCESS"
   else
      retval="$GM_FAIL"
   fi
   return "$retval"
}

#####################################################################
# Arguments:
# Argument 1: LOG_LEVEL(SILENT(0) / VERBOSE(1))
# Argument 2: NO_WATCHDOG (optional)
#####################################################################
start () {
   local retval_watchdog="$GM_SUCCESS"
   local retval_vmw_conn_notify="$GM_SUCCESS"
   local retval="$GM_SUCCESS"
   local VERBOSE=1
   local SILENT=0
   local LOG_LEVEL=${1:-$VERBOSE}

   status_vmw_conn_notify
   retval_vmw_conn_notify=$?
   if [ "$retval_vmw_conn_notify" -ne "$GM_SUCCESS" ]; then
      vmw_logger "Starting $VMW_CONN_NOTIFY_NAME service" $SILENT $INFO
      $VMW_CONN_NOTIFY_BIN
      retval_vmw_conn_notify=$?
   fi

   if [ "$retval_vmw_conn_notify" -ne "$GM_SUCCESS" ]; then
      retval_vmw_conn_notify="$GM_FAIL"
      vmw_logger "Unable to start $VMW_CONN_NOTIFY_NAME binary" $SILENT $ERROR
   else
      retval_vmw_conn_notify="$GM_SUCCESS"
      vmw_logger "$VMW_CONN_NOTIFY_NAME binary started successfully" $SILENT $SUCCESS
      ### Create the lock file ###
      mkdir -p /var/lock/subsys
      touch $LOCK_FILE
      load_netfilter_modules
      add_vnetchain_filter_rules
   fi

   status_watchdog
   retval_watchdog=$?
   if [ "$retval_watchdog" -ne "$GM_SUCCESS" ]; then
      vmw_logger "Starting watchdog for $VMW_CONN_NOTIFY_NAME" $SILENT $INFO
      start_watchdog
      sleep 1
      status_watchdog
      retval_watchdog=$?
   fi
   # Display results for watchdog
   if [ "$retval_watchdog" -ne "$GM_SUCCESS" ]; then
      retval_watchdog="$GM_FAIL"
      vmw_logger "Unable to start $VMW_CONN_NOTIFY_WATCHDOG_NAME" $SILENT $ERROR
   else
      retval_watchdog="$GM_SUCCESS"
      vmw_logger "$VMW_CONN_NOTIFY_WATCHDOG_NAME started successfully" $SILENT $SUCCESS
   fi

   retval=$((retval_vmw_conn_notify || retval_watchdog))
   if [ "$retval" -eq "$GM_SUCCESS" ]; then
     vmw_logger "$VMW_CONN_NOTIFY_NAME service is running" $LOG_LEVEL $SUCCESS
   else
     vmw_logger "$VMW_CONN_NOTIFY_NAME service is stopped" $LOG_LEVEL $ERROR
     remove_vnetchain_filter_rules
   fi
   eval_cmd "$retval" "$LOG_LEVEL"

   return "$retval"
}




#####################################################################
# Arguments:
# Argument 1: LOG_LEVEL(SILENT(0) / VERBOSE(1))
# Argument 2: NO_WATCHDOG (optional)
#####################################################################
stop() {
   local retval_watchdog="$GM_SUCCESS"
   local retval_vmw_conn_notify="$GM_SUCCESS"
   local retval="$GM_SUCCESS"
   local VERBOSE=1
   local SILENT=0
   local LOG_LEVEL=${1:-$VERBOSE}

   vmw_logger "Stopping $VMW_CONN_NOTIFY_NAME service" $SILENT $INFO

   status_watchdog
   retval_watchdog=$?
   if [ "$retval_watchdog" -eq "$GM_SUCCESS" ]; then
      vmw_logger "Stopping $VMW_CONN_NOTIFY_WATCHDOG_NAME for vmw_conn_notify" $SILENT $SUCCESS
      kill -SIGTERM `pgrep -f $VMW_CONN_NOTIFY_WATCHDOG_NAME`
      sleep 1
      status_watchdog
      retval_watchdog=$?
   fi
   if [ "$retval_watchdog" -ne "$GM_SUCCESS" ]; then
      retval_watchdog="$GM_SUCCESS"
      vmw_logger "$VMW_CONN_NOTIFY_WATCHDOG_NAME stopped successfully" $SILENT $SUCCESS
   else
      retval_watchdog="$GM_FAIL"
      vmw_logger "Failed to stop $VMW_CONN_NOTIFY_WATCHDOG_NAME" $SILENT $ERROR
   fi

   ### sleep for few seconds, wait for watchdog to stop
   sleep $VMW_CONN_NOTIFY_STOP_TIME

   status_vmw_conn_notify
   retval_vmw_conn_notify=$?
   if [ "$retval_vmw_conn_notify" -eq "$GM_SUCCESS" ]; then
      vmw_logger "Stopping $VMW_CONN_NOTIFY_NAME binary" $SILENT $INFO
      kill -SIGTERM `pidof $VMW_CONN_NOTIFY_NAME`
      sleep 1
      status_vmw_conn_notify
      retval_vmw_conn_notify=$?
   fi

   if [ "$retval_vmw_conn_notify" -ne "$GM_SUCCESS" ]; then
      retval_vmw_conn_notify="$GM_SUCCESS"
      rm -f $LOCK_FILE
      ### sleep for few seconds, wait for thin-agent to stop
      sleep $VMW_CONN_NOTIFY_STOP_TIME
      vmw_logger "$VMW_CONN_NOTIFY_NAME binary stopped successfully" $SILENT $SUCCESS
   else
      retval_vmw_conn_notify="$GM_FAIL"
      vmw_logger "Failed to stop $VMW_CONN_NOTIFY_NAME binary" $SILENT $ERROR
   fi

   retval=$((retval_vmw_conn_notify || retval_watchdog))
   if [ "$retval" -eq "$GM_SUCCESS" ]; then
     vmw_logger "$VMW_CONN_NOTIFY_NAME service is stopped" $LOG_LEVEL $SUCCESS
   else
     vmw_logger "$VMW_CONN_NOTIFY_NAME service is running" $LOG_LEVEL $ERROR
   fi
   remove_vnetchain_filter_rules

   eval_cmd "$retval" "$LOG_LEVEL"

   return "$retval"
}


#####################################################################
# Arguments:
# Argument 1: LOG_LEVEL(SILENT(0) / VERBOSE(1))
# Argument 2: NO_WATCHDOG (optional)
#####################################################################
status() {
   local retval_vmw_conn_notify="$GM_SUCCESS"
   local retval_watchdog="$GM_SUCCESS"
   local retval="$GM_SUCCESS"
   local VERBOSE=1
   local SILENT=0
   local LOG_LEVEL=${1:-$VERBOSE}
   local NO_WATCHDOG=${2:-0}

   status_vmw_conn_notify
   retval_vmw_conn_notify=$?
   if [ "$retval_vmw_conn_notify" -ne "$GM_SUCCESS" ]; then
      retval_vmw_conn_notify="$GM_FAIL"
      vmw_logger "$VMW_CONN_NOTIFY_NAME binary is stopped" $SILENT $INFO
   else
      retval_vmw_conn_notify="$GM_SUCCESS"
      vmw_logger "$VMW_CONN_NOTIFY_NAME binary is running" $SILENT $INFO
   fi

   if [ "${NO_WATCHDOG}" -eq 0 ]; then
      status_watchdog
      retval_watchdog=$?
      if [ "$retval_watchdog" -ne "$GM_SUCCESS" ]; then
         retval_watchdog="$GM_FAIL"
         vmw_logger "$VMW_CONN_NOTIFY_WATCHDOG_NAME is stopped" $SILENT $INFO
      else
         retval_watchdog="$GM_SUCCESS"
         vmw_logger "$VMW_CONN_NOTIFY_WATCHDOG_NAME is running" $SILENT $INFO
      fi
   fi

   retval=$((retval_vmw_conn_notify || retval_watchdog))
   if [ "$retval" -eq "$GM_SUCCESS" ]; then
      vmw_logger "$VMW_CONN_NOTIFY_NAME service is running" $LOG_LEVEL $INFO
   else
      vmw_logger "$VMW_CONN_NOTIFY_NAME service is stopped" $LOG_LEVEL $INFO
   fi

   return "$retval"
}

### main logic ###
case "$1" in
   # Argument : LOG_LEVEL(SILENT(0) / VERBOSE(1))
   start)
      start $2
      ;;
   # Argument : LOG_LEVEL(SILENT(0) / VERBOSE(1))
   stop)
      stop $2
      ;;
   # Argument : LOG_LEVEL(SILENT(0) / VERBOSE(1))
   status)
      status $2
      ;;
   # Argument : LOG_LEVEL(SILENT(0) / VERBOSE(1))
   restart|force-reload)
      stop $2
      # Sleep for few seconds for the MUX connection to get refreshed
      sleep 3
      start $2
      ;;
   *)
      vmw_logger "Usage: $0 {start|stop|restart|status|force-reload}" $VERBOSE $ERROR
      exit 1
esac
exit 0
