_msfvenom

#compdef msfvenom
#autoload
#
# zsh completion for msfvenom in Metasploit Framework Project (https://www.metasploit.com)
#
# github:  https://github.com/Green-m/msfvenom-zsh-completion
#
# author:  Green-m (greenm.xxoo@gmail.com)
#
# license: GNU General Public License v3.0
#
# Copyright (c) 2018, Green-m
# All rights reserved.
#

_msfvenom() {
  zmodload zsh/parameter 2>/dev/null
  
  local CACHE_DIR="$HOME/.zsh/venom-cache"
  local OPT_CACHE_DIR="$CACHE_DIR/options"
  local PAYLOAD_CACHE="$CACHE_DIR/payloads"
  local ARCH_CACHE="$CACHE_DIR/archs"
  local FORMAT_CACHE="$CACHE_DIR/formats"
  local ENCODER_CACHE="$CACHE_DIR/encoders"
  local PLATFORM_CACHE="$CACHE_DIR/platforms"
  local NOPS_CACHE="$CACHE_DIR/nops"

  [[ ! -d "$OPT_CACHE_DIR" ]] && mkdir -p "$OPT_CACHE_DIR"

  local curcontext="$curcontext" state line ret=1
  typeset -A opt_args

_arguments -C \
    '(-h --help)'{-h,--help}'[show help]' \
    '(-l --list)'{-l,--list}'[list modules]:type:(payloads encoders nops platforms archs encrypt formats all)' \
    '(-p --payload)'{-p,--payload}'[payload to use]:payload:->payloads' \
    '(--list-options)--list-options[list payload options]' \
    '(-f --format)'{-f,--format}'[output format]:format:->formats' \
    '(-e --encoder)'{-e,--encoder}'[the encoder to use]:encoder:->encoders' \
    '(--smallest)--smallest[generate the smallest possible payload]' \
    '(--encrypt)--encrypt[type of encryption]:encryption:(aes256 base64 rc4 xor)' \
    '(--encrypt-key)--encrypt-key[key for --encrypt]:key' \
    '(--encrypt-iv)--encrypt-iv[IV for --encrypt]:iv' \
    '(-a --arch)'{-a,--arch}'[architecture to use]:arch:->archs' \
    '(--platform)--platform[platform for payload]:platform:->platforms' \
    '(-o --out)'{-o,--out}'[save payload to file]:file:_files' \
    '(-b --bad-chars)'{-b,--bad-chars}'[characters to avoid (e.g. \\x00\\xff)]:badchars' \
    '(-n --nops --nopsled)'{-n,--nops,--nopsled}'[nopsled length and optional generator]:length: :nop:->nops' \
    '(--encoder-space)--encoder-space[max size of encoded payload]:size' \
    '(-i --iterations)'{-i,--iterations}'[number of times to encode]:count' \
    '(-c --add-code)'{-c,--add-code}'[include additional win32 shellcode file]:file:_files' \
    '(-x --template)'{-x,--template}'[custom executable template]:file:_files' \
    '(-k --keep)'{-k,--keep}'[preserve template behavior and inject as new thread]' \
    '(-v --var-name)'{-v,--var-name}'[custom variable name for certain formats]:name' \
    '(-t --timeout)'{-t,--timeout}'[seconds to wait for STDIN (default 30)]:seconds' \
    '*: :->msf_options' && ret=0

  case "$state" in
    payloads)
      local -a p_list
      p_list=(${(f)"$(cat_or_build "$PAYLOAD_CACHE" "payloads")"})
      [[ ${#p_list} -gt 0 ]] && _describe 'payloads' p_list && ret=0
      ;;
    nops)
      local -a n_list
      n_list=(${(f)"$(cat_or_build "$NOPS_CACHE" "nops")"})
      [[ ${#n_list} -gt 0 ]] && _values 'nops' "${(@)n_list}" && ret=0
      ;;
    archs)
      local -a a_list
      a_list=(${(f)"$(cat_or_build "$ARCH_CACHE" "archs")"})
      [[ ${#a_list} -gt 0 ]] && _values 'architectures' "${(@)a_list}" && ret=0
      ;;
    encoders)
      local -a e_list
      e_list=(${(f)"$(cat_or_build "$ENCODER_CACHE" "encoders")"})
      [[ ${#e_list} -gt 0 ]] && _values 'encoders' "${(@)e_list}" && ret=0
      ;;
    formats)
      local -a f_list
      f_list=(${(f)"$(cat_or_build "$FORMAT_CACHE" "formats")"})
      [[ ${#f_list} -gt 0 ]] && _values 'formats' "${(@)f_list}" && ret=0
      ;;
    platforms)
      local -a plat_list
      plat_list=(${(f)"$(cat_or_build "$PLATFORM_CACHE" "platforms")"})
      [[ ${#plat_list} -gt 0 ]] && _values 'platforms' "${(@)plat_list}" && ret=0
      ;;
    msf_options)
      if [[ ! "$words[$CURRENT]" == -* ]]; then
        local selected_payload=""
        local i
        for (( i=1; i < $#words; i++ )); do
          if [[ "$words[$i]" == "-p" || "$words[$i]" == "--payload" ]]; then
            selected_payload="${words[$i+1]}"
            break
          fi
        done

        if [[ -n "$selected_payload" ]]; then
          local -a dynamic_opts
          dynamic_opts=(${(f)"$(__msfvenom_dynamic_opts "$selected_payload" "$OPT_CACHE_DIR")"})
          [[ ${#dynamic_opts} -gt 0 ]] && _values "options for $selected_payload" "${(@)dynamic_opts}" && ret=0
        else
          _values 'general options' 'LHOST=' 'LPORT=' 'RHOST=' 'RPORT=' && ret=0
        fi
      fi
      ;;
  esac
  return ret
}

cat_or_build() {
  local cache_file=$1
  local type=$2
  if [[ ! -f "$cache_file" ]]; then
    mkdir -p "${cache_file:h}"
    msfvenom --list "$type" 2>/dev/null | \
      awk '/^[[:space:]]{4}[a-z0-9]/ {print $1}' | \
      grep -vE '^(Name|----|Framework|^$)' > "$cache_file"
  fi
  cat "$cache_file"
}

__msfvenom_dynamic_opts() {
  local payload=$1
  local cache_dir=$2
  local cache_name="${payload//\//_}" 
  local cache_path="$cache_dir/$cache_name"

  if [[ ! -f "$cache_path" ]]; then
    msfvenom -p "$payload" --list-options 2>/dev/null | \
      awk '$0 ~ /yes/ || $0 ~ /no/ {print $1}' | \
      grep -vE '^(Name|----|^$|Current|Required|Description)' | \
      sort -u | sed 's/$/=/' > "$cache_path"
  fi
  cat "$cache_path"
}

_msfvenom "$@"