#!/bin/bash

type _comp_filter &> /dev/null || return
type _comp_contains &> /dev/null || return
type _comp_filter_shorts &> /dev/null || return
type _comp_get_parameter &> /dev/null || return

function _rmg() {

    local cur prev words cword opts arg args gadgets value_options operations enum_operations components
    _init_completion || return

    value_options="--argument-position --bound-name --component --config --objid --plugin --sample-folder"
    value_options="$value_options --signature --template-folder --threads --wordlist-file --wordlist-folder --yso"
    value_options="$value_options --serial-version-uid --payload-serial-version-uid"

    _count_args "" "@(${value_options// /|})"
    _comp_get_parameter 3
    COMPREPLY=()

    gadgets="BeanShell1 C3P0 Clojure CommonsBeanutils1 CommonsCollections1 CommonsCollections2"
    gadgets="$gadgets CommonsCollections3 CommonsCollections4 CommonsCollections5 CommonsCollections6"
    gadgets="$gadgets CommonsCollections7 FileUpload1 Groovy1 Hibernate1 Hibernate2 JBossInterceptors1"
    gadgets="$gadgets JRMPClient JRMPListener JSON1 JavassistWeld1 Jdk7u21 Jython1 MozillaRhino1 MozillaRhino2"
    gadgets="$gadgets Myfaces1 Myfaces2 ROME Spring1 Spring2 URLDNS Vaadin1 Wicket1 JRMPClient2 AnTrinh"

    components="act dgc reg"
    operations="bind call codebase enum guess known listen objid rebind roguejmx scan serial unbind"
    enum_operations="list string-marshalling codebase localhost-bypass security-manager jep290 filter-bypass activator"


    if [[ $cword -eq 1 ]]; then

        if [[ $cur == -* ]]; then
            opts="--help"
        else
            opts="$operations"
        fi

    elif [[ ${words[1]} == "bind" || ${words[1]} == "rebind" ]]; then

        if _comp_contains "--config --plugin --yso" $prev; then
            _filedir
            return 0

        elif _comp_contains "--ssrf-response --bind-objid --gadget-cmd" $prev; then
            return 0

        elif _comp_contains "--gadget-name" $prev; then
            opts="$gadgets"

        elif [[ $cur == -* ]] || [[ $args -ge 8 ]]; then
            opts="--help"
            opts="$opts --bind-bound-name"
            opts="$opts --bind-objid"
            opts="$opts --config"
            opts="$opts --encode"
            opts="$opts --gadget-cmd"
            opts="$opts --gadget-name"
            opts="$opts --gopher"
            opts="$opts --localhost-bypass"
            opts="$opts --no-color"
            opts="$opts --plugin"
            opts="$opts --raw"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"
            opts="$opts --yso"

        elif [[ $args -eq 4 ]]; then
            _ip_addresses
            compopt -o nospace
            return 0

        else
            return 0
        fi


    elif [[ ${words[1]} == "call" ]]; then

        if _comp_contains "--config --plugin" $prev; then
            _filedir
            return 0

        elif _comp_contains "--component" $prev; then
            opts="$components"

        elif _comp_contains "--serial-version-uid --ssrf-response --objid --signature --bound-name" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -gt 7 ]]; then
            opts="--help"
            opts="$opts --bound-name"
            opts="$opts --component"
            opts="$opts --config"
            opts="$opts --encode"
            opts="$opts --follow"
            opts="$opts --force-activation"
            opts="$opts --gopher"
            opts="$opts --no-color"
            opts="$opts --objid"
            opts="$opts --plugin"
            opts="$opts --raw"
            opts="$opts --serial-version-uid"
            opts="$opts --signature"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"

        else
            return 0
        fi


    elif [[ ${words[1]} == "codebase" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif _comp_contains "--component" $prev; then
            opts="$components"

        elif _comp_contains "--payload-serial-version-uid --serial-version-uid --ssrf-response --position --objid --signature --bound-name" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -gt 7 ]]; then
            opts="--help"
            opts="$opts --bound-name"
            opts="$opts --component"
            opts="$opts --config"
            opts="$opts --encode"
            opts="$opts --follow"
            opts="$opts --force-activation"
            opts="$opts --gopher"
            opts="$opts --no-color"
            opts="$opts --objid"
            opts="$opts --payload-serial-version-uid"
            opts="$opts --position"
            opts="$opts --raw"
            opts="$opts --serial-version-uid"
            opts="$opts --signature"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"

        elif [[ $args -eq 5 ]]; then
            _ip_addresses
            compopt -o nospace
            return 0

        else
            return 0
        fi


    elif [[ ${words[1]} == "enum" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif _comp_contains "--serial-version-uid --ssrf-response --bound-name" $prev; then
            return 0

        elif _comp_contains "--scan-action" $prev; then
            opts="$enum_operations"

        elif _comp_contains "--dgc-method" $prev; then
            opts="clean dirty"

        elif _comp_contains "--registry-method" $prev; then
            opts="bind lookup rebind unbind"

        elif [[ $cur == -* ]] || [[ $args -ge 4 ]]; then
            opts="--help"
            opts="$opts --activate"
            opts="$opts --bound-name"
            opts="$opts --config"
            opts="$opts --dgc-method"
            opts="$opts --encode"
            opts="$opts --follow"
            opts="$opts --force-activation"
            opts="$opts --gopher"
            opts="$opts --localhost-bypass"
            opts="$opts --no-color"
            opts="$opts --raw"
            opts="$opts --registry-method"
            opts="$opts --scan-action"
            opts="$opts --serial-version-uid"
            opts="$opts --signature"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"

        else
            return 0
        fi


    elif [[ ${words[1]} == "guess" ]]; then

        if _comp_contains "--config --wordlist-file --wordlist-folder --sample-folder --template-folder" $prev; then
            _filedir
            return 0

        elif _comp_contains "--component" $prev; then
            opts="$components"

        elif _comp_contains "--serial-version-uid --bound-name --objid --signature --threads" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 4 ]]; then
            opts="--help"
            opts="$opts --bound-name"
            opts="$opts --component"
            opts="$opts --config"
            opts="$opts --create-samples"
            opts="$opts --follow"
            opts="$opts --force-activation"
            opts="$opts --force-guessing"
            opts="$opts --guess-duplicates"
            opts="$opts --no-color"
            opts="$opts --no-progress"
            opts="$opts --objid"
            opts="$opts --sample-folder"
            opts="$opts --serial-version-uid"
            opts="$opts --signature"
            opts="$opts --ssl"
            opts="$opts --stack-trace"
            opts="$opts --template-folder"
            opts="$opts --threads"
            opts="$opts --trusted"
            opts="$opts --update"
            opts="$opts --verbose"
            opts="$opts --wordlist-file"
            opts="$opts --wordlist-folder"
            opts="$opts --zero-arg"

        else
            return 0
        fi


    elif [[ ${words[1]} == "known" ]]; then

        if [[ $cur == -* ]] || [[ $args -ge 3 ]]; then
            opts="--help"
            opts="$opts --no-color"
            opts="$opts --stack-trace"

        else
            return 0
        fi


    elif [[ ${words[1]} == "listen" ]]; then

        if _comp_contains "--config --plugin --yso" $prev; then
            _filedir
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 6 ]]; then
            opts="--help"
            opts="$opts --config"
            opts="$opts --no-color"
            opts="$opts --plugin"
            opts="$opts --stack-trace"
            opts="$opts --verbose"
            opts="$opts --yso"

        elif [[ $args -eq 2 ]]; then
            _ip_addresses
            return 0

        elif [[ $args -eq 4 ]]; then
            opts="$gadgets"

        else
            return 0
        fi


    elif [[ ${words[1]} == "objid" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 3 ]]; then
            opts="--help"
            opts="$opts --config"
            opts="$opts --no-color"
            opts="$opts --stack-trace"
            opts="$opts --verbose"

        else
            return 0
        fi


    elif [[ ${words[1]} == "roguejmx" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif _comp_contains "--objid --forward-host --forward-port --forward-bound-name --forward-objid" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 4 ]]; then
            opts="--help"
            opts="$opts --config"
            opts="$opts --follow"
            opts="$opts --forward-bound-name"
            opts="$opts --forward-host"
            opts="$opts --forward-objid"
            opts="$opts --forward-port"
            opts="$opts --no-color"
            opts="$opts --objid"
            opts="$opts --ssl"
            opts="$opts --stack-trace"
            opts="$opts --verbose"

        elif [[ $args -eq 2 ]]; then
            _ip_addresses
            return 0

        else
            return 0
        fi


    elif [[ ${words[1]} == "scan" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif _comp_contains "--ports --threads --timeout-read --timeout-connect" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 3 ]]; then
            opts="--help"
            opts="$opts --config"
            opts="$opts --no-color"
            opts="$opts --no-progress"
            opts="$opts --ports"
            opts="$opts --stack-trace"
            opts="$opts --threads"
            opts="$opts --timeout-connect"
            opts="$opts --timeout-read"
            opts="$opts --verbose"

        else
            return 0
        fi


    elif [[ ${words[1]} == "serial" ]]; then

        if _comp_contains "--config --plugin --yso" $prev; then
            _filedir
            return 0

        elif _comp_contains "--serial-version-uid --ssrf-response --position --bound-name --objid --signature" $prev; then
            return 0

        elif _comp_contains "--component" $prev; then
            opts="$components"

        elif [[ $cur == -* ]] || [[ $args -ge 6 ]]; then
            opts="--help"
            opts="$opts --bound-name"
            opts="$opts --component"
            opts="$opts --config"
            opts="$opts --encode"
            opts="$opts --follow"
            opts="$opts --force-activation"
            opts="$opts --gopher"
            opts="$opts --no-color"
            opts="$opts --objid"
            opts="$opts --plugin"
            opts="$opts --position"
            opts="$opts --raw"
            opts="$opts --serial-version-uid"
            opts="$opts --signature"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"
            opts="$opts --yso"

        elif [[ $args -eq 4 ]]; then
            opts="$gadgets"

        else
            return 0
        fi


    elif [[ ${words[1]} == "unbind" ]]; then

        if _comp_contains "--config" $prev; then
            _filedir
            return 0

        elif _comp_contains "--ssrf-response" $prev; then
            return 0

        elif [[ $cur == -* ]] || [[ $args -ge 3 ]]; then
            opts="--help"
            opts="$opts --config"
            opts="$opts --encode"
            opts="$opts --gopher"
            opts="$opts --localhost-bypass"
            opts="$opts --no-color"
            opts="$opts --plugin"
            opts="$opts --raw"
            opts="$opts --ssl"
            opts="$opts --ssrf"
            opts="$opts --ssrf-response"
            opts="$opts --stack-trace"
            opts="$opts --stream-protocol"
            opts="$opts --verbose"

        else
            return 0
        fi
    fi

    _comp_filter "opts"
    _comp_filter_shorts "opts"

	mapfile -t COMPREPLY < <(compgen -W "${opts}" -- "${cur}") 
	return 0
}

complete -F _rmg rmg
