# The plugin configuration file
###############################
PLUGIN_CONF_FILE="traffic-accounting.conf"

# Preinit return value for success
PLUGIN_RET_VAL=0

# Define some global variables
VERBOSE=0

# Check sanity of eg. environment
traffic_accounting_helper_sanity_check()
{
  # Check whether chains exists
  if ! ip4tables -nL ACCOUNTING_INPUT_CHAIN >/dev/null 2>&1; then
    echo "** ERROR: ACCOUNTING_INPUT_CHAIN does not exist! **" >&2
    return 1
  fi

  if ! ip4tables -nL ACCOUNTING_OUTPUT_CHAIN >/dev/null 2>&1; then
    echo "** ERROR: ACCOUNTING_OUTPUT_CHAIN does not exist! **" >&2
    return 1
  fi

  # Check if chains inserted in the main chains
#  if ! ip4tables -nL INPUT |grep -q '^ACCOUNTING_INPUT_CHAIN '; then
#    echo "** ERROR: ACCOUNTING_INPUT_CHAIN is not inserted in the INPUT chain! **" >&2
#    return 1
#  fi

#  if ! ip4tables -nL OUTPUT |grep -q '^ACCOUNTING_OUTPUT_CHAIN '; then
#    echo "** ERROR: ACCOUNTING_OUTPUT_CHAIN is not inserted in the OUTPUT chain! **" >&2
#    return 1
#  fi

  return 0
}


traffic_accounting_helper_do_work()
{
  local RETVAL=0

  # Touch the log file (just in case it doesn't exist yet):
  touch /var/log/traffic-accounting.log

  # Truncate file
  printf "" >/tmp/traffic-accounting.new

  # Process the input chain
  if [ "$VERBOSE" = "1" ]; then
    echo "${INDENT}Traffic Accounting Hosts:"
    echo "${INDENT}-------------------------"
  fi

  # Also include default unicast route addresses, (0.0.0.0/0 and ::/0)
  DEFAULT_ADDR="0.0.0.0/0"
  if [ "$IPV6_SUPPORT" = "1" ]; then
    DEFAULT_ADDR="$DEFAULT_ADDR ::/0"
  fi

  IFS=' ,'
  for host in $TRAFFIC_ACCOUNTING_HOSTS $DEFAULT_ADDR; do
    old_entry="$(grep "^$host " /var/log/traffic-accounting.log)"
    old_ip="$(echo "$old_entry" |cut -s -d' ' -f2)"
    old_in_value="$(echo "$old_entry" |cut -s -d' ' -f3)"
    old_out_value="$(echo "$old_entry" |cut -s -d' ' -f4)"

    # If value is non-existant make it zero
    if [ -z "$old_in_value" ]; then
      old_in_value=0
    fi

    # If value is non-existant make it zero
    if [ -z "$old_out_value" ]; then
      old_out_value=0
    fi

    # Get host_ip, if it fails, skip rule
    # Parse/get hostname. Try to use host cache if applicable
    # NOTE: get_dynamic_host_cached returns hostname in $host_ip
    if ! get_dynamic_host_cached $host || [ -z "$host_ip" ]; then
      echo "** WARNING: Skipping rule for \"$host\"! **" >&2
      RETVAL=1
      continue
    fi

    IFS=' ,'
    for mon_host_ip in $host_ip; do
      echo "${INDENT}Monitoring host \"$host\" with IP: $mon_host_ip"

      if [ "$VERBOSE" = "1" ]; then
        printf "${INDENT}old_ip=$old_ip host_ip=$mon_host_ip "
      fi

      # Process input chain
      OLDFOUND=0
      if [ -n "$old_ip" ]; then
        get_numeric_ip_version "$mon_host_ip"
        case $? in
        4)
          LCOUNT=0
          IFS=$EOL
          for LINE in `ip4tables -xnvL ACCOUNTING_INPUT_CHAIN |sed -e "1,2d"`; do
            ipt_ip="$(echo "$LINE" |awk '{ print $8 }')"

            LCOUNT=$((LCOUNT + 1))
            if [ "$ipt_ip" = "$old_ip" ]; then
              ip4tables -R ACCOUNTING_INPUT_CHAIN $LCOUNT -s $mon_host_ip -j RETURN
              if [ "$VERBOSE" = "1" ]; then
                printf "in_action=update "
              fi
              OLDFOUND=1
              ipt_in_value="$(echo "$LINE" |awk '{ print $2 }')"

              break
            fi
          done
          ;;
        6)
          if [ "$IPV6_SUPPORT" = "1" ]; then
            LCOUNT=0
            IFS=$EOL
            for LINE in `ip6tables -xnvL ACCOUNTING_INPUT_CHAIN |sed -e "1,2d"`; do
              ipt_ip="$(echo "$LINE" |awk '{ print $7 }')"

              LCOUNT=$((LCOUNT + 1))
              if [ "$ipt_ip" = "$old_ip" ]; then
                ip6tables -R ACCOUNTING_INPUT_CHAIN $LCOUNT -s $mon_host_ip -j RETURN
                if [ "$VERBOSE" = "1" ]; then
                  printf "in_action=update "
                fi
                OLDFOUND=1
                ipt_in_value="$(echo "$LINE" |awk '{ print $2 }')"

                break
              fi
            done
          fi
          ;;
        esac
      fi

      if [ $OLDFOUND -eq 0 ]; then
        if [ "$VERBOSE" = "1" ]; then
          printf "in_action=add "
        fi

        if [ "$mon_host_ip" = "0.0.0.0/0" -o "$mon_host_ip" = "::/0" ]; then
          iptables -A ACCOUNTING_INPUT_CHAIN -s $mon_host_ip -j RETURN
        else
          iptables -I ACCOUNTING_INPUT_CHAIN 1 -s $mon_host_ip -j RETURN
        fi

        # Preset values to zero as none exist yet
        ipt_in_value=0
      fi

      # Process output chain
      OLDFOUND=0
      if [ -n "$old_ip" ]; then
        get_numeric_ip_version "$mon_host_ip"
        case $? in
        4)
          LCOUNT=0
          IFS=$EOL
          for LINE in `ip4tables -xnvL ACCOUNTING_OUTPUT_CHAIN |sed -e "1,2d"`; do
            ipt_ip="$(echo "$LINE" |awk '{ print $9 }')"

            LCOUNT=$((LCOUNT + 1))
            if [ "$ipt_ip" = "$old_ip" ]; then
              ip4tables -R ACCOUNTING_OUTPUT_CHAIN $LCOUNT -d $mon_host_ip -j RETURN
              if [ "$VERBOSE" = "1" ]; then
                printf "out_action=update "
              fi
              OLDFOUND=1
              ipt_out_value="$(echo "$LINE" |awk '{ print $2 }')"

              break
            fi
          done
          ;;
        6)
          if [ "$IPV6_SUPPORT" = "1" ]; then
            LCOUNT=0
            IFS=$EOL
            for LINE in `ip6tables -xnvL ACCOUNTING_OUTPUT_CHAIN |sed -e "1,2d"`; do
              ipt_ip="$(echo "$LINE" |awk '{ print $8 }')"

              LCOUNT=$((LCOUNT + 1))
              if [ "$ipt_ip" = "$old_ip" ]; then
                ip6tables -R ACCOUNTING_OUTPUT_CHAIN $LCOUNT -d $mon_host_ip -j RETURN
                if [ "$VERBOSE" = "1" ]; then
                  printf "out_action=update "
                fi
                OLDFOUND=1
                ipt_out_value="$(echo "$LINE" |awk '{ print $2 }')"

                break
              fi
            done
          fi
          ;;
        esac
      fi

      if [ $OLDFOUND -eq 0 ]; then
        if [ "$VERBOSE" = "1" ]; then
          printf "out_action=add "
        fi

        if [ "$mon_host_ip" = "0.0.0.0/0" -o "$mon_host_ip" = "::/0" ]; then
          iptables -A ACCOUNTING_OUTPUT_CHAIN -d $mon_host_ip -j RETURN
        else
          iptables -I ACCOUNTING_OUTPUT_CHAIN 1 -d $mon_host_ip -j RETURN
        fi

        # Preset values to zero as none exist yet
        ipt_out_value=0
      fi

      # Calculate new in value
      new_in_value=$((old_in_value + ipt_in_value))

      # Calculate new out value
      new_out_value=$((old_out_value + ipt_out_value))
      if [ "$VERBOSE" = "1" ]; then
        printf "old_in_val=$old_in_value ipt_in_val=$ipt_in_value new_in_val=$new_in_value old_out_val=$old_out_value ipt_out_val=$ipt_out_value new_out_val=$new_out_value"
      fi

      # Create entry in accounting file
      echo "$host $mon_host_ip $new_in_value $new_out_value" >>/tmp/traffic-accounting.new

      if [ "$VERBOSE" = "1" ]; then
        printf "\n\n"
      fi
    done
  done

  # FIXME: Don't use old-file
  if [ -f /var/log/traffic-accounting.log ]; then
    if [ -f /var/log/traffic-accounting.log.old ]; then
      rm -f /var/log/traffic-accounting.log.old
    fi

    mv /var/log/traffic-accounting.log /var/log/traffic-accounting.log.old
  fi
  mv /tmp/traffic-accounting.new /var/log/traffic-accounting.log

  return $RETVAL
}


############
# Mainline #
############

if [ "$1" = "-v" -o "$1" = "--verbose" ]; then
  VERBOSE=1
fi

# Check where to find the config file
CONF_FILE=""
if [ -n "$PLUGIN_CONF_PATH" ]; then
  CONF_FILE="$PLUGIN_CONF_PATH/$PLUGIN_CONF_FILE"
fi

# Check if the config file exists
if [ ! -f "$CONF_FILE" ]; then
  echo "** ERROR: Config file \"$CONF_FILE\" not found! **" >&2
  PLUGIN_RET_VAL=1
else
  # Source the plugin config file
  . "$CONF_FILE"

  # Only proceed if environment ok
  if ! traffic_accounting_helper_sanity_check; then
    PLUGIN_RET_VAL=1
  else
    # Create actual rules
    if ! traffic_accounting_helper_do_work; then
      PLUGIN_RET_VAL=1
    fi
  fi
fi
