#!/bin/sh -eu

usage() {
  echo "Usage: $0 {start|stop|reload|restart|block|drop|deny|unblock|allow|flush} [address ...]"
  exit 1
}

start() {
  /usr/sbin/nft -f /etc/nftables.conf
  [ -f /var/lib/nftables/ip-blocklist.nft ] && /usr/sbin/nft -f /var/lib/nftables/ip-blocklist.nft || :
  [ -f /var/lib/nftables/ip6-blocklist.nft ] && /usr/sbin/nft -f /var/lib/nftables/ip6-blocklist.nft || :
}

stop() {
  /usr/sbin/nft list set inet chef-filter ip-blocklist > /var/lib/nftables/ip-blocklist.nft
  /usr/sbin/nft list set inet chef-filter ip6-blocklist > /var/lib/nftables/ip6-blocklist.nft
  /usr/sbin/nft delete table inet chef-filter
<% if node[:roles].include?("gateway") -%>
  /usr/sbin/nft delete table ip chef-nat
<% end -%>
}

reload() {
  stop
  start
}

block() {
  for address in "$@"
  do
    case "$address" in
      *.*) /usr/sbin/nft --check add element inet chef-filter ip-blocklist "{ $address }" && /usr/sbin/nft add element inet chef-filter ip-blocklist "{ $address }" ;;
      *:*) /usr/sbin/nft --check add element inet chef-filter ip6-blocklist "{ $address }" && /usr/sbin/nft add element inet chef-filter ip6-blocklist "{ $address }" ;;
    esac
  done
}

unblock() {
  for address in "$@"
  do
    case "$address" in
      *.*) /usr/sbin/nft --check delete element inet chef-filter ip-blocklist "{ $address }" && /usr/sbin/nft delete element inet chef-filter ip-blocklist "{ $address }" ;;
      *:*) /usr/sbin/nft --check delete element inet chef-filter ip6-blocklist "{ $address }" && /usr/sbin/nft delete element inet chef-filter ip6-blocklist "{ $address }" ;;
    esac
  done
}

flush() {
  /usr/sbin/nft --check flush set inet chef-filter ip-blocklist && /usr/sbin/nft flush set inet chef-filter ip-blocklist
  /usr/sbin/nft --check flush set inet chef-filter ip6-blocklist && /usr/sbin/nft flush set inet chef-filter ip6-blocklist
}

command="${1:-}"
if [ -n "$command" ]; then
  shift
fi

if [ -z "$command" ]; then
  usage
fi

case "$command" in
  start) start;;
  stop) stop;;
  reload|restart) reload;;
  block|drop|deny) block "$@";;
  unblock|allow) unblock "$@";;
  flush) flush;;
  *)
    usage
    ;;
esac

exit 0