#!/usr/bin/env bash

# -----------------------------------------------------------------
# ipset set listing wrapper script
#
# https://github.com/AllKind/ipset_list
# https://sourceforge.net/projects/ipset-list/
# -----------------------------------------------------------------

# Copyright (C) 2013-2019 Mart Frauenlob aka AllKind (AllKind@fastest.cc)
#
# 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, either version 3 of the License, or
# (at your option) any later version.
#
# 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 General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# -----------------------------------------------------------------
#
# This is the bash programmable completion for ipset_list
#
# -----------------------------------------------------------------
# Requirements:
#
# The bash completion package version 2.0 or greater is recommended.
# https://github.com/scop/bash-completion
#
# If the package is not available, things might not be so reliable.
# Also the colon (if there) is removed from COMP_WORDBREAKS.
# This alteration is globally. Which might affect other completions
# if they don't take care of it themselves.
#
# -----------------------------------------------------------------
# Installation (quote from bash-completion README):
#
# Install it in one of the directories pointed to by
# bash-completion's pkgconfig file variables.  There are two
# alternatives: the recommended one is 'completionsdir' (get it with
# "pkg-config --variable=completionsdir bash-completion") from which
# completions are loaded on demand based on invoked commands' names,
# so be sure to name your completion file accordingly, and to include
# for example symbolic links in case the file provides completions
# for more than one command.  The other one which is present for
# backwards compatibility reasons is 'compatdir' (get it with
# "pkg-config --variable=compatdir bash-completion") from which files
# are loaded when bash_completion is loaded.
#
# For packages using GNU autotools the installation can be handled
# for example like this in configure.ac:
#
#  PKG_CHECK_VAR(bashcompdir, [bash-completion], [completionsdir], ,
#    bashcompdir="${sysconfdir}/bash_completion.d")
#  AC_SUBST(bashcompdir)
#
# ...accompanied by this in Makefile.am:
#
#  bashcompdir = @bashcompdir@
#  dist_bashcomp_DATA = # completion files go here
#
# For cmake we ship the bash-completion-config.cmake and
# bash-completion-config-version.cmake files. Example usage:
#
#  find_package(bash-completion)
#  if(BASH_COMPLETION_FOUND)
#    message(STATUS
#      "Using bash completion dir ${BASH_COMPLETION_COMPLETIONSDIR}")
#  else()
#    set (BASH_COMPLETION_COMPLETIONSDIR "/etc/bash_completion.d")
#    message (STATUS
#      "Using fallback bash completion dir ${BASH_COMPLETION_COMPLETIONSDIR}")
#  endif()
#
#  install(FILES your-completion-file DESTINATION
#    ${BASH_COMPLETION_COMPLETIONSDIR})
#
# For backwards compatibility it is still possible to put it into
# ~/.bash_completion or /etc/bash_completion.d/.
# -----------------------------------------------------------------

# Name may be modified
ipset_list=ipset_list

# -----------------------------------------------------------------

# -----------------------------------------------------------------
# DO NOT MODIFY ANYTHING BEYOND THIS LINE!
# -----------------------------------------------------------------

shopt -s extglob

# -----------------------------------------------------------------
# Functions
# -----------------------------------------------------------------

_ipset_list_show_sets() {
COMPREPLY=( $( compgen -W '${sets[@]}' -- "$cur" ) )
# dedupe sets listing
for ((i=set_index; i < ${#words[@]}-1; i++)); do
    _ipset_list_remove_reply_entry "${words[i]}"
done
}

_ipset_list_remove_reply_entry() {
local -i x
while (($#)); do
    for x in ${!COMPREPLY[@]}; do
        if [[ ${COMPREPLY[x]} = $1 ]]; then
            if [[ $_DEBUG_NF_COMPLETION ]]; then
                printf "removing dupe entry COMPREPLY[$x]: %s\n" \
                    "${COMPREPLY[x]}"
            fi
            unset COMPREPLY[x]
            break
        fi
    done
    shift
done
}

# -----------------------------------------------------------------
# Main
# -----------------------------------------------------------------

_ipset_list_complete() {
local -i i=x=got_bashcompl=iactive=0
local -i show_all=isolate=show_members=resolve=headers_only=names_only=0
local -i header_operation=set_index=0
local cur prev cword words str_tmp
local sets=( $(command "${ipset_list:-ipset_list}" -n 2>/dev/null) ) 2>/dev/null
local opts=(- -- -? -a -c -d -h -i -m -n -r -s -t -v)
local Copts=(-Ca -Cs -Co)
local Fopts=(-Fh -Fi -Fg -Fr -Oi)
local Gopts=(-Gp -Gs -Gx)
local Iopts=(-- -d -r -s -G ${Gopts[@]} -To)
local Hopts=(-Hi -Hr -Hs -Ht -Hv)
local Topts=(-T -Tm -To -Ts)
local Xopts=(-Xh -Xg -Xr -Xs -Xo)
local arr_types=() arr_tmp=()

: ${PATH:=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin}

COMPREPLY=()

# expecting _get_comp_words_by_ref() to exist from bash_completion
if declare -f _get_comp_words_by_ref &>/dev/null; then got_bashcompl=1
    _get_comp_words_by_ref -n : cur prev cword words || return
else got_bashcompl=0 # not so neat, but a workaround
    COMP_WORDBREAKS="${COMP_WORDBREAKS//:/}"
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"
    cword=$COMP_CWORD
    for i in ${!COMP_WORDS[@]}; do words[i]="${COMP_WORDS[i]}"; done
fi

#_DEBUG_NF_COMPLETION=Y
if [[ $_DEBUG_NF_COMPLETION ]]; then
    printf "\nCOMP_WORDBREAKS: <%s>\n" "$COMP_WORDBREAKS"
    printf "COMP_LINE: <%s>\n" "$COMP_LINE"
    printf "COMP_TYPE: <%s>\n" "$COMP_TYPE"
    printf "COMP_POINT: <%s>\n" "$COMP_POINT"
    printf "COMP_KEY: <%s>\n" "$COMP_KEY"
    printf "COMP_CWORD: <%s>\n" "$COMP_CWORD"
    printf "\ncur: <%s> prev: <%s>\n" "$cur" "$prev"
    printf "words:\n"
    printf "<%s>\n" "${words[@]}"
fi

# collect info of cmdline
for ((i=1; i < ${#words[@]}-1; i++)); do
    case "${words[i]}" in
        -)  ((set_index)) && break || iactive=1 ;;
        -a) ((set_index)) && break || show_all=1 ;;
        -i) ((set_index)) && break || isolate=1 Copts=(-Co) ;;
        -m) ((set_index)) && break || show_members=1 ;;
        -n) ((set_index)) && break || names_only=1 ;;
        -r) ((set_index)) && break || resolve=1 ;;
        -t) ((set_index)) && break || headers_only=1 Xopts=(-Xh -Xs) Fopts=(${Fopts[*]/-Oi/}) ;;
        --) ((set_index)) && break || set_index=$((i+1)) ;;
        -\?|-h|-v)
            ((set_index)) || return 0
        ;;
        @(-Fh|-Fi|-Hi|-Xh)) ((set_index)) && break || header_operation=1 ;;
        *)
        ((set_index)) && break
        # options expecting an opt arg
        str_tmp="-@(d|Fg|Fh|Fi|Fr|Gp|Gs|Gx|Hi|Ht|Hr|Hs|Hv|Mc|Oi|T|To|Xg|Xh|Xr|Xs)"
        if [[ ${words[i-1]} = $str_tmp ]]; then
            continue
		elif [[ ${words[i-1]} = @(-Gp|-Gs|-Gx) && ${words[i]} != -* ]]; then
            continue
        fi
        # if not an option, register set index
        str_tmp="-@(-|?|a|c|d|h|i|m|n|r|s|t|v|Ca|Cs|Co|Fg|Fh|Fi|Fr|Gp|Gs|Gx|Hi|Ht|Hr|Hs|Hv|Mc|Oi|T|To|Tm|Ts|Xg|Xh|Xo|Xr)"
        if [[ ${words[i]} != $str_tmp ]]; then
            for x in ${!sets[@]}; do
                if [[ ${sets[x]} = ${words[i]} ]]; then
                    set_index=$i
                    break
                fi
            done
        fi
    esac
done

# invalid combinations of options
if ((names_only)); then
    if ((headers_only)); then
        return 0
    fi
fi
if ((headers_only||names_only)); then
    if ((show_all || show_members || isolate || resolve)); then
        return 0
    fi
elif ((isolate)); then
    if ((show_all || header_operation)); then
        return 0
    fi
fi

# start setting compreply
# all depends on $set_index
if ((set_index)); then
    if ((isolate && cword > set_index)); then
        return 0 # allow only one set with isolate
    fi
    # dont' allow an option after the set name(s)
    # allows to list sets which start with an hyphen
    # and also handles those who have the name of ipset_list options
    _ipset_list_show_sets
else
if [[ $prev = -@(\?|d|h|v|Fg|Fi|Fr|Hi|Oi|T|To|Xg|Xr) ]]; then
    return 0
elif [[ $prev = -Xs ]]; then
    # list sets if user does not want to enter a glob
    _ipset_list_show_sets
elif [[ $prev = -Ht ]]; then i=0
    # show supported set types
    while read -r; do
        [[ $REPLY = "Supported set types:"* ]] && ((!i)) && \
            i=1 && continue
        ((i)) || continue
        if [[ $REPLY = *:* ]]; then
            set -- $REPLY
            arr_types[${#arr_types[@]}]="$1"
        fi
    done < <(ipset help)
    for i in ${!arr_types[@]}; do # remove dupe entries
        for x in ${!arr_tmp[@]}; do
            [[ ${arr_tmp[x]} = ${arr_types[i]} ]] && continue 2
        done
        arr_tmp[${#arr_tmp[@]}]="${arr_types[i]}"
    done
    arr_types=( "${arr_tmp[@]}" )
    COMPREPLY=( $( compgen -W '${arr_types[@]}' -- "$cur" ) )
elif [[ $prev = @(-Fh|-Xh) ]]; then
    # retrieve list of headers
    if ((${#sets[*]} > 0)); then
        while read -r; do
            [[ $REPLY = Name ]] && continue
            COMPREPLY[${#COMPREPLY[@]}]="$REPLY"
        done < <(command "$ipset_list" -t "${sets[0]}" 2>/dev/null|command awk -F: '{print $1}')
        compopt -o nospace
        local IFS=$'\n'
        if [[ $prev = -Xh ]]; then
            COMPREPLY=( $( compgen -P '"' -S ':*"' \
                -W '${COMPREPLY[@]}' -- "$cur" ) )
        elif [[ $prev = -Fh ]]; then
            COMPREPLY=( $( compgen -P '"' -S ':"' \
                -W '${COMPREPLY[@]}' -- "$cur" ) )
        fi
    fi
elif [[ $prev = @(-@(Hr|Hs|Hv|Mc)) ]]; then
    # options making use of arithmetic comparison
    compopt -o nospace
    COMPREPLY=( $( compgen -P '\' -W '\! \< \> \<= \>=' -- "$cur" ) )
elif [[ $prev = @(-Gp|-Gs|-Gx) && $cur != -* ]]; then # fails on files starting with a dash. I don't care.
	if ((got_bashcompl)); then
		_filedir
		COMPREPLY=( $( compgen -W 'auto none ${COMPREPLY[@]}' -- "$cur" ) )
	else
		COMPREPLY=( $( compgen -f -- "$cur" ) )
	fi
elif [[ $cur = -* ]]; then
    # any option is requested
    case "$prev" in
        -@(-|\?|d|h|v|Fg|Fh|Fi|Fr|Hi|Ht|Hr|Hs|Hv|Mc|Oi|T|To|Xg|Xh|Xr))
        # options that exclude any other option,
        # or need a value we can't predict
            return 0
        ;;
    esac
    if ((${#words[@]} > 2)); then
        # these options don't allow any other - remove them
        opts=("${opts[@]/@(-v|-h|-\?)/}")
    fi
    # some options allow only a subset of other options
    if ((iactive)); then
        COMPREPLY=( $(compgen -W '${Iopts[@]} ${Gopts[@]}' -- $cur ) )
    elif ((isolate)); then
        COMPREPLY=( $(compgen -W '-- -Co -d -r -s -Fg -Fr ${Gopts[@]} -Oi -T -To -Xg -Xo -Xr' -- $cur ) )
    elif ((names_only)); then
        COMPREPLY=( $(compgen -W \
            '-- -c ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} -Xs' \
            -- $cur ) )
    elif ((headers_only)); then
        COMPREPLY=( $(compgen -W \
            '-- -c ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} ${Xopts[@]}' \
            -- $cur ) )
    elif ((show_members)); then
        COMPREPLY=( $(compgen -W \
            '-- -c -d -r -s ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} -Xg -Xr -Xo' \
            -- $cur ) )
    elif ((show_all)); then
        COMPREPLY=( $(compgen -W \
            '-- -c -d -r -s ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} ${Xopts[@]}' \
            -- $cur ) )
    elif ((resolve)); then
        COMPREPLY=( $(compgen -W \
            '-- -a -c -d -s -m ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} ${Xopts[@]}' \
            -- $cur ) )
    elif ((header_operation)); then
        COMPREPLY=( $(compgen -W \
            '-- -a -c -d -s -m -t ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} ${Xopts[@]}' \
            -- $cur ) )
    else
        COMPREPLY=( $(compgen -W \
            '${opts[@]} ${Copts[@]} ${Fopts[@]} ${Gopts[@]} ${Hopts[@]} -Mc ${Topts[@]} ${Xopts[@]}' \
            -- $cur ) )
    fi
    # post process the reply
    if ((${#COMPREPLY[@]})); then
        x=$((set_index ? set_index : ${#words[*]}-1))
        # mutual exclusive options
        for ((i=1; i < x; i++)); do
            case "${words[i]}" in
                -Fg) _ipset_list_remove_reply_entry "-Fr" ;;
                -Fr) _ipset_list_remove_reply_entry "-Fg" ;;
                -Xg) _ipset_list_remove_reply_entry "-Xr" ;;
                -Xr) _ipset_list_remove_reply_entry "-Xg" ;;
            esac
            # options allowed multiple times
            if [[ ${words[i]} = @(""|-|-@(Fh|Fi|Hi|Mc|Oi|T|Xh|Xs)) ]]; then
                continue
            else # remove dupe
                _ipset_list_remove_reply_entry "${words[i]}"
            fi
        done
    fi
elif [[ $cur = * ]]; then
    # non option request
    # default to sets listing
    # except we are in interactive mode
    if ! ((iactive)); then
        _ipset_list_show_sets
    fi
fi
fi

((got_bashcompl)) && __ltrim_colon_completions "$cur"

if [[ $_DEBUG_NF_COMPLETION ]]; then
    printf "COMPREPLY:\n"
    printf "<%s>\n" "${COMPREPLY[@]}"
fi
}
complete -F _ipset_list_complete "${ipset_list:-ipset_list}"

