#compdef nft
# ------------------------------------------------------------------------------
# Description
# -----------
#
#  Completion script for nft 0.9.0 (https://www.netfilter.org/projects/nftables/index.html).
#
# ------------------------------------------------------------------------------
# Authors
# -------
#
#  * Markus Richter ( https://github.com/mqus , <mqus@disroot.org>)
#
# ------------------------------------------------------------------------------

_nft(){
  local -a rules states prev args families options descriptors
  local state="start" line nextstate cmd_obj cmd_subcmd cmd_fam cmd_tab cmd_chain #curcontext="$curcontext"

  options=(
    '(-)'{-h,--help}'[show help]'
    '(-)'{-v,--version}'[print version information]'
    "(-i --interactive)"{-i,--interactive}'[read input from interactive CLI]: :->end'
    "(-f --file)"{-f,--file}'[read input from <filename>]:nftables rule file:_files'
    '(-c --check -n --numeric -N)'{-c,--check}"[check command's validity without actually applying the changes]"
    '(-j --json)'{-j,--json}'[format output in json]'
    '(-c --check -N)*'{-n,--numeric}'[can be specified up to 3 times, Shows 1:network addresses(default behaviour), 2:internet services (port numbers) and 3:protocols, user IDs, and group IDs numerically]'
    '(-s --stateless)'{-s,--stateless}'[omit stateful information of ruleset]'
    '(-N -n --numeric -c --check)'-N'[translate IP addresses to names]'
    '(-a --handle)'{-a,--handle}'[output rule handle]'
    '(-e --echo)'{-e,--echo}'[echo what has been added, inserted or replaced]'
    {-I,--includepath}'[add specified directory to the paths searched for include files]:include directory [/usr/share]:include directory:_directories'
  )

  # start a state machine. The state is modified by _arguments if the
  # current argument (descriptors) cannot be completed. Each state has to define is successive state and the
  # 'descriptors' for _arguments, which essentially tells _arguments how to complete
  local _i=0
  while true;do
    (( _i+=1 ))
    #Guard for endless loops
    [[ $_i -gt 100 ]] && return 1

    descriptors=()
    nextstate="end"
    case $state in
      (start)
        ##if line is empty (at the start) or ends with semicolon, autocomplete subcommands,
        # else if we are after a space,complete a semicolon (end the current nft command) and start anew
        if [[ $line[1] = "" || $line[1] =~ ';$' ]] ; then
          descriptors=( ":: :_nft_subcommands" )
          nextstate="category"
        else
          if [[ $words =~ ' $' ]]; then
            descriptors=(':: :(\;)')
          else
            descriptors=(':argument: ')
          fi
          nextstate="start"
        fi
        ;;
      (category)
        case $line[1] in
          (add | list | flush | delete | create | rename | insert | replace | reset)
            descriptors=( ":: :_nft_${line[1]}" )
            nextstate=$line[1]
            ;;
          (monitor)
            descriptors=( ":: : _nft_mon_filter"  )
            nextstate="mon1"
            ;;
          (export)
            descriptors=( ":: :(ruleset)" ":: :_nft_out_format" )
            nextstate="preend"
            ;;
          (describe)
            descriptors=( ":expression: ")
            nextstate="start" #x restart
            ;;
          (*)
            return 1;
            ;;
        esac

        #descriptors=( "(ruleset)" )
        #nextstate="end"
        ;;
      (mon1)
        case $line[1] in
          (new | destroy)
            #      descriptors=( ":: :_nft_mon_keywords"  ":: :_nft_out_format")
            descriptors=( ":: : _nft_mon_keywords")
            nextstate="mon1"
            ;;
          (tables | chains | sets | rules | elements | ruleset)
            descriptors=( ":: : _nft_out_format")
            nextstate="preend"
            ;;
        esac
        ;;
      #all completions for create and insert match with the completions of add
      (create | insert)
        state="add"
        ;|
      #all completions for reset and flush match with the completions of list
      (reset | flush)
        state="list"
        ;|
      #(add(^table)/create(^table)/delete/flush(^ruleset)/insert/list(^ruleset)/rename/replace)[family]table
      (reset | delete | insert | rename | replace | add | create | flush | list)
        if [[ $state = "add" && $line[1] = "table" ]]; then
          descriptors=( ":: :_nft_families" ":table name:")
          nextstate="start" #x restart
        elif [[ $state = "list" && ( $line[1] = "ruleset" || $line[1] = "tables" ) ]]; then
          descriptors=( ":: :_nft_families")
          nextstate="start" #x restart
        elif [[ $state = "delete" && $line[1] = "table" ]]; then
          descriptors=(": : _nft_table all-handle")
          nextstate="tcomplete-delete-table"
        else
          cmd_obj=$line[1]
          cmd_subcmd=$state
          descriptors=(": : _nft_table all")
          nextstate="tcomplete"
        fi
        ;;
      (tcomplete-delete-table)
        # if only a family was completed, complete the table name.
        case $line[1] in
          (arp | bridge | inet | ip | ip6 | netdev)
            descriptors=(": : _nft_table ${line[1]}-handle")
            cmd_fam=$line[1]
            ;;
          # if 'handle' was completed, complete the handle number.
          (handle)
            descriptors=(": : _nft_table_handle_all " )
            ;;
          # else, complete nothing and go to the next state. default family is 'ip'
          (*)
            descriptors=()
            cmd_fam="ip"
            ;;
        esac
        nextstate="delete-table"
        ;;
      (tcomplete)
        # if only a family was completed, complete the table name.
        case $line[1] in
          (arp | bridge | inet | ip | ip6 | netdev)
            descriptors=(": : _nft_table ${line[1]}")
            cmd_fam=$line[1]
            ;;
          # else, complete nothing and go to the next state. default family is 'ip'
          (*)
            descriptors=()
            cmd_fam="ip"
            ;;
        esac
        nextstate="$cmd_subcmd-$cmd_obj"
        ;;
      (list-table)
        descriptors=(":: :(\;)")
        nextstate="start"
        ;;
      (delete-table)
        #if family AND handle were input, complete handle number for given family.
        if [[ $line[1] == "handle" ]]; then
          descriptors=(":table handle: _nft_table_handle $cmd_fam" )
        else
          descriptors=()
        fi
        nextstate="start"
        ;;
      (delete-chain | delete-set | delete-quota | delete-counter | delete-ct\\ helper)
      cmd_tab=$line[1]
      descriptors=(": : _nft_object $cmd_fam $cmd_tab $cmd_obj true")
      nextstate="delete-obj-handle"
      ;;
      (delete-obj-handle)
        if [[ $line[1] == "handle" ]]; then
          descriptors=(": : _nft_object_handle $cmd_fam $cmd_tab $cmd_obj")
        else
          descriptors=(": :(\;)")
        fi
        nextstate="start"
        ;;
      (add-chain)
        descriptors=(":chain name:")
        nextstate="start"
        ;;
      (rename-chain)
        cmd_tab=$line[1]
        descriptors=(": : _nft_object $cmd_fam $cmd_tab chain false")
        nextstate="add-chain"
        ;;
      (replace-rule | delete-rule)
        cmd_tab=$line[1]
        descriptors=(": : _nft_object $cmd_fam $cmd_tab chain false")
        nextstate="repdel-rule"
        ;;
      (repdel-rule)
        cmd_chain=$line[1]
        descriptors=(": :(handle)" ": : _nft_rule_handle $cmd_fam $cmd_tab ${line[1]}")
        if [[ $cmd_subcmd = "replace" ]];then
          nextstate="rule-stmt"
        else
          nextstate="start"
        fi
        ;;
      (add-rule)
        cmd_tab=$line[1]
        descriptors=(": : _nft_object $cmd_fam $cmd_tab chain false")
        nextstate="add-rule-2"
        ;;
      (add-rule-2)
        cmd_chain=$line[1]
        descriptors=(": :(handle index position)")
        nextstate="add-rule-3"
        ;;
      (add-rule-3)
        case $line[1] in
          (index | position)
            descriptors=(":${line[1]}:")
            ;;
          (handle)
            descriptors=(": : _nft_rule_handle $cmd_fam $cmd_tab $cmd_chain")
            ;;
          (*)
            descriptors=()
            ;;
        esac
        nextstate="rule-stmt"
        ;;
      (rule-stmt)
        #TODO
        #    _nft_rule $cmd_fam $cmd_tab $cmd_chain\
          #  && return 0;
        descriptors=(":expression: ")
        nextstate="start"
        ;;
      (list-set | list-map | delete-map | list-chain | list-flowtable | delete-flowtable | list-ct\\ helper | list-counter | list-quota | list-meter)
      cmd_tab=$line[1]
      descriptors=(": : _nft_object $cmd_fam $cmd_tab $cmd_obj false")
      nextstate="start"
      ;;
      (add-element | delete-element)
        cmd_tab=$line[1]
        descriptors=(": : _nft_object $cmd_fam $cmd_tab '\(map\|set\)' false")
        nextstate="start"
        ;;
      #TODO:
      #(add-set | add-map)
      #(add-flowtable)
      #("add-ct\ helper")
      #(add-counter)
      #(add-quota)

      (*)
        return 1;
        ;;
    esac
    _arguments -C -s \
               "${options[@]}" \
               "${descriptors[@]}" \
               "*:: :->$nextstate" \
      && return 0;

  done
} # end _nft

_nft_subcommands(){
  local commands=(
    'add:add a table, chain, rule, set, map, or object'
    'list:list a ruleset, table, chain, set, map, or object'
    'flush:flush (delete everything from) a ruleset, table, chain, set, or map'
    'export:print the ruleset in a machine readable format (json or xml)'
    'delete:delete a table, chain, rule, set, element, map, or object'
    'create:similar to add but returns an error for existing chain'
    'rename:rename the specified chain'
    'insert:similar to the add command, but the rule is prepended to the beginning of the chain or before the rule at the given position'
    'replace:similar to the add command, but replaces the specified rule'
    'reset:list-and-reset stateful object'
    'monitor:listen to Netlink events'
    'describe:show information about the type of an expression and its data type'
  )
  _describe -t commands 'nft subcommand' commands "$@"
}

_nft_mon_filter(){
  local monitor_filters=(
    'new:show only events of created objects'
    'destroy:show only events of deleted objects'
  )
  _describe -t monitor_filters 'nft monitor' monitor_filters -J monitor_filters "$@"
  _nft_mon_keywords
}

_nft_mon_keywords(){
  local monitor_keywords=(
    'tables:show table events'
    'chains:show chain events'
    'sets:show set events'
    'rules:show rule events'
    'elements:show only events of element objects'
    'ruleset:show ruleset events, such as table, chain, rule, set, counters  and  quotas'
  )
  _describe -t monitor_keywords 'nft monitor' monitor_keywords -J monitor_keywords "$@"
  _nft_out_format
}

_nft_out_format(){
  local monitor_format=(
    'json:format output to JSON'
    'xml:format output to XML'
  )
  _describe -t monitor_format "output format" monitor_format -J monitor_format "$@"
}

_nft_add(){
  local commands=(
    'table:add a new table'
    'flowtable:add a new flowtable'
    'chain:add a chain to a table'
    'rule:add a rule to an existing chain'
    'set:add a set to a table'
    'map:add a map to a table'
    'element:add one or multiple element(s) to a set or map'
    'ct\ helper:add a ct helper to a table'
    'counter:add a named counter to a table'
    'quota:add a named quota helper to a table'
  )
  _describe -t commands 'nft add' commands "$@"
}

_nft_create(){
  local commands=(
    "table:add a table, but return an error if it already exists"
    "chain:add a chain to a table, but return an error if it already exists"
    "flowtable:add a flowtable, but return an error if it already exists"
  )
  _describe -t commands 'nft create' commands "$@"
}

_nft_delete(){
  local commands=(
    "table:delete the specified table"
    "chain:delete the specified chain, chain must be empty and mustn't be used as jump target"
    "rule:delete the specified rule, rule must be referable to by a handle"
    "set:delete the specified set"
    "map:delete the specified map"
    "element:delete element(s) from the specified set/map"
    "flowtable:delete the specified flowtable"
    "ct\ helper:delete the specified ct helper"
    "counter:delete the specified counter"
    "quota:delete the specified quota"
  )
  _describe -t commands 'nft delete' commands "$@"
}

_nft_flush(){
  local commands=(
    "ruleset:clear the whole ruleset, including removing all tables and containing objects"
    "table:flush all chains and rules of the specified table"
    "chain:flush all rules of the specified chain"
    "set:remove all elements from the specified set"
    "map:remove all elements from the specified map"
  )
  _describe -t commands 'nft flush' commands "$@"
}

_nft_insert(){
  local commands=(
    "rule:prepend a rule to the beginning of the chain or before the rule with the given handle"
  )
  _describe -t commands 'nft insert' commands "$@"
}

_nft_list(){
  local commands=(
    "ruleset:print the ruleset in human-readable format"
    "tables:list all tables (undocumented)"
    "table:list all chains and rules of the specified table"
    "chain:list all rules of the specified chain"
    "set:display the elements in the specified set"
    "map:display the elements in the specified map"
    "flowtable:list all flowtables"
    "ct\ helper:display stateful information the ct helper holds"
    "counter:display stateful information the counter holds"
    "quota:display stateful information the quota holds"
  )
  _describe -t commands 'nft list' commands "$@"
}

_nft_rename(){
  local commands=(
    "chain:replace a chain"
  )
  _describe -t commands 'nft rename' commands "$@"
}

_nft_replace(){
  local commands=(
    "rule:replace a rule"
  )
  _describe -t commands 'nft replace' commands "$@"
}

_nft_reset(){
  local commands=(
    'ct\ helper:reset and list a ct helper to a table'
    'counter:reset and list a counter from a table'
    'quota:reset and list a quota object a table'
  )
  _describe -t commands 'nft reset' commands "$@"
}
_nft_families(){
  local families=(
    "ip:IPv4 address family"
    "ip6:IPv6 address family"
    "inet:internet (IPv4+IPv6) address family"
    "arp:ARP address family, handling IPv4 ARP packets"
    "bridge:Bridge address family, handling packets which traverse a bridge device"
    "netdev:Netdev address family, handling packets from ingress"
  )
  _describe -t families 'nft families' families "$@"
}

_nft_table(){
  # complete the names of tables and the families of existing tables
  #$1 can be: all all-handle <family> <family>-handle
  local tables=()
  if [[ "$1" =~ "^all" ]]; then
    local families=( ${(f)"$(_call_program -p tables nft list tables 2>/dev/null \
        | cut -d\  -f2 )"} )
    # ip is the default family, search also for table names there
    1="${1/all/ip}"
    _describe -t families "family" families -J "family"
  fi
  if [[ "$1" =~ "-handle$" ]]; then
    tables=("handle:address the table by handle")
    #remove -handle from $1 to be able to complete table names
    1="${1/-handle/}"
    _describe -t tables "table" tables -V "handle"
  fi
  case $1 in
  (arp | bridge | inet | ip | ip6 | netdev)
    tables=( ${(f)"$(_call_program -p tables nft -a list ruleset 2>/dev/null \
      | grep '^table '"$1" | sed 's/table // ;s/{ # handle // ;s/\(\S*\) \(\S*\) \(\S*\)/\2:type \1, handle \3/' )"} )
    _describe -t tables "table" tables -V "table-name"
  ;;
  esac
}

_nft_table_handle(){
  # complete the handles of tables with the specified family (with the table name in the description)
  #$1:protocol family
  local tables=( ${(f)"$(_call_program -p tables nft -a list ruleset 2>/dev/null \
      | grep '^table '"$1" | sed 's/table // ;s/{ # handle // ;s/\(\S*\) \(\S*\) \(\S*\)/\3:\2(type \1)/' )"} )
  _describe -t tables "table handle" tables
}

_nft_table_handle_all(){
  # complete the handles of tables of all families (with the table name in the description)
  local tables=( ${(f)"$(_call_program -p tables nft -a list ruleset 2>/dev/null \
      | grep '^table' | sed 's/table // ;s/{ # handle // ;s/\(\S*\) \(\S*\) \(\S*\)/\3:\2(type \1)/' )"} )
  _describe -t tables "table handle" tables
}

_nft_object(){
  # complete the names of objects contained directly in a table (with the handle number in the description)
  #$1:protocol family
  #$2:table
  #$3:object type (chain/set/map/flowtable/ct helper/counter/quota/meter)
  #$4:include 'handle'?
  local objects=( ${(f)"$(_call_program -p objects nft -a list table $1 $2 2>/dev/null\
      | grep ""\^\\s\*$3"" | sed 's/\s*'"$3"' // ;s/ { # \(.*\)/:(\1)/' )"} )
  if $4 ;then
    objects+=( "handle:address $3 by handle")
  fi
  _describe -t objects "$3" objects
}

_nft_object_handle(){
  # complete handles of objects contained directly in a table (with the name in the description)
  #$1:protocol family
  #$2:table
  #$3:object type (chain/set/ct helper/counter/quota)
  local handles=( ${(f)"$(_call_program -p handles nft -a list table $1 $2 2>/dev/null\
      | grep ""\^\\s\*$3"" | sed 's/\s*'"$3"' // ;s/ { # handle// ;s/\(\S*\) \(\S*\)/\2:\1/' )"} )
  _describe -t handles "$3-handle" handles
}

_nft_rule_handle(){
  # complete the handles of rules (and put the rule into the description)
  #$1:protocol family
  #$2:table
  #$3:chain name
  local rules=( ${(f)"$(_call_program -p nft-rule-handle nft -a list chain $1 $2 $3 2>/dev/null \
      |grep -v '^\s*\(table\|chain\|type\|\}\)'|sed 's/^\s*\(.*\) # handle \(\S*\)$/\2:\1/' )"} )
  # don't sort those entries alphabetically, so they get shown in the order they are executed in nftables
  _describe -t rules "rule" rules -V "rules"
}

#currently, only the `nft` command is covered by this script.
_nft "$@"

# Local Variables:
# mode: Shell-Script
# sh-indentation: 2
# indent-tabs-mode: nil
# sh-basic-offset: 2
# End:
# vim: ft=zsh sw=2 ts=2 et
