#compdef pwn
# ----------------------------------------------------------------------------------------
# Zsh completion file for pwn
#
# Pwntools Command-line Interface
#
# This file was generated with help from Zargparse (https://github.com/ctil/zargparse)
# NOTE: Don't use Zargparse any more as it may bring known errors!
# ----------------------------------------------------------------------------------------

_comp_pid() {
    if [[ $1 == 'pid' ]] {
        local line pids lines
        pids=()
        # extract command output line by line
        lines=("${(@f)$(ps --ppid 2 -p 2 -N -o pid=,tty=,user=,comm=)}")
        for line ($lines) {
            pids+=${line[(w)1]} # extract first word (pid)
        }
        _wanted nonk-pids expl 'non-kernel process ID' \
            compadd -o nosort -ld lines -a pids
    } elif [[ $1 == 'procs' ]] {
        local line names lines
        names=()
        # extract command output line by line
        lines=("${(@f)$(ps --ppid 2 -p 2 -N -o comm=,pid=,tty=,user=)}")
        for line ($lines) {
            names+=${line[(w)1]} # extract first word (pid)
        }
        _wanted nonk-comm expl 'non-kernel process names' \
            compadd -o nosort -ld lines -a names
    }
}

_comp_funcs() {
    local names
    output=$(readelf -Ws $1 2>/dev/null)
    if [[ $? -ne 0 ]] {
        return
    }
    names=("${(f)$(awk '$4=="FUNC" && $7!="UND" {print $8}' <(echo "$output"))}")
    if [[ $#names -gt 0 ]] {
      _values "FUNCTIONS" $names
    }
}

_comp_error() {
    local -a names
    which errno > /dev/null
    if [[ $? -ne 0 ]] {
        return
    }
    names=("${(f)$(LC_ALL=C errno -l | awk '{printf "%s:(%d) %s\n", $1, $2, substr($0, length($1 $2)+3)}')}")
    _describe "ERRORS" names
}

typeset -A opt_args

_pwn() {
  setopt localoptions
  setopt extendedglob
  # Define the subcommands
  local -a commands
  
  commands=(
    'asm:Assemble shellcode into bytes'
    'checksec:Check binary security settings'
    'constgrep:Looking up constants from header files.'
    'cyclic:Cyclic pattern creator/finder'
    'debug:Debug a binary in GDB'
    'disasm:Disassemble bytes into text format'
    'disablenx:Disable NX for an ELF binary'
    'elfdiff:Compare two ELF files'
    'elfpatch:Patch an ELF file'
    'errno:Prints out error messages'
    'hex:Hex-encodes data provided on the command line or stdin'
    'libcdb:Print various information about a libc binary'
    'phd:Pretty hex dump'
    'pwnstrip:Strip binaries for CTF usage'
    'scramble:Shellcode encoder'
    'shellcraft:Microwave shellcode -- Easy, fast and delicious'
    'template:Generate an exploit template'
    'unhex:Decodes hex-encoded data provided on the command line or via stdin.'
    'update:Check for pwntools updates'
    'version:Pwntools version'
  )

  # Global flags (i.e. ones not associated with a subcommand)
  _arguments \
    "1:COMMAND: _describe 'command' commands" \
    '(-h --help)'{-h,--help}'[show help message and exit]' \
    '*:: :->args'

  # Flags for each subcommand
  case $state in
    args)
      case $words[1] in
        asm)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-f --format)'{-f,--format}'[Output format]:FORMAT:(raw hex string elf)' \
            '(-o --output)'{-o,--output}'[Output file (defaults to stdout)]:OUTPUT:_files' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '(-v --avoid)'{-v,--avoid}'[Encode the shellcode to avoid the listed bytes (provided as hex)]:HEXBYTES: ' \
            '(-n --newline)'{-n,--newline}'[Encode the shellcode to avoid newlines]' \
            '(-z --zero)'{-z,--zero}'[Encode the shellcode to avoid NULL bytes]' \
            '(-d --debug)'{-d,--debug}'[Debug the shellcode with GDB]' \
            '(-i --infile)'{-i,--infile}'[Specify input file]:INPUT:_files' \
            '(-r --run)'{-r,--run}'[Run output]' \
            '*:LINE: '
            ;;
        checksec)
            _arguments \
              '(-h --help)'{-h,--help}'[show help message and exit]' \
              '--file[File to check]:ELF:_files' \
              '*:ELF:_files'
            ;;
        constgrep)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-e --exact)'{-e,--exact}'[Do an exact match for a constant instead of searching for a regex]' \
            '(-i --case-insensitive)'{-i,--case-insensitive}'[Search case insensitive]' \
            '(-m --mask-mode)'{-m,--mask-mode}'[Instead of searching for a specific constant value, search for values not containing strictly less bits that the given value.]' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '1:REGEX: ' \
            '2:CONST: ' \
            '*: :'
            ;;
        cyclic)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-a --alphabet)'{-a,--alphabet}'[The alphabet to use in the cyclic pattern (defaults to all lower case letters)]:ALPHABETS: ' \
            '(-n --length)'{-n,--length}'[Size of the unique subsequences (defaults to 4)]:LENGTH: ' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '(-l -o --offset --lookup)'{-l,-o,--offset,--lookup}'[Do a lookup instead printing the alphabet]:LOOKUP: ' \
            '1:COUNT: ' \
            '*: :'
            ;;
        debug)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '-x[Execute GDB commands from this file]:GDBSCRIPT:_files' \
            '--pid[PID to attach to]:PID: _comp_pid pid' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '--exec[File to debug]:EXE:_files' \
            '--process[Name of the process to attach to (e.g. "bash")]:NAME: _comp_pid procs' \
            '--sysroot[GDB sysroot path]:PATH:_files -/'
          ;;
        disasm)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '(-a --address)'{-a,--address}'[Base address]:ADDR: ' \
            '--color[Color output]' \
            '--no-color[Disable color output]' \
            '*:HEX: '
            ;;
        disablenx)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '*:ELF:_files'
            ;;
        elfdiff)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '1:ELF_A:_files' \
            '2:ELF_B:_files' \
            '*: :'
            ;;
        elfpatch)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '1:ELF:_files' \
            '2:OFFSET: ' \
            '3:PATCH: ' \
            '*: :'
            ;;
        errno)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '1:ERROR:_comp_error' \
            '*: :'
            ;;
        hex)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-p --prefix)'{-p,--prefix}'[Insert a prefix before each byte]:PREFIX: ' \
            '(-s --separator)'{-s,--separator}'[Add a separator between each byte]:SEPARATOR: ' \
            '*:DATA: '
            ;;
        libcdb)
          local -a libc_cmds=(
            'lookup:Lookup a libc version by function offsets'
            'hash:Display information of a libc version given an unique hash'
            'file:Dump information about a libc binary'
            'fetch:Fetch libc database'
          )
          _arguments \
            "1:LIBC_COMMAND: _describe 'libc command' libc_cmds" \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '*:: :->args'
          if [[ $state != 'args' ]] {
              return
          }
          case $words[1] {
            (lookup)
              _arguments \
                '(-h --help)'{-h,--help}'[show help message and exit]' \
                '(-d --download-libc)'{-d,--download-libc}'[Attempt to download the matching libc.so]' \
                '--no-unstrip[Do NOT attempt to unstrip the libc binary with debug symbols from a debuginfod server]' \
                '--offline-only[Attempt to searching with offline only mode]' \
                '*:SYMBOL_OFFSET_PAIRS: _message -r "SYMBOL_OFFSET_PAIRS Hint: read 3e0 write 520"'
              ;;
            (hash)
              _arguments \
                '(-h --help)'{-h,--help}'[show help message and exit]' \
                '(-t --hash_type)'{-t,--hash_type}'[The type of the provided hash value]:HASH TYPE:(id buildid md5 sha1 sha256)' \
                '(-d --download-libc)'{-d,--download-libc}'[Attempt to download the matching libc.so]' \
                '--no-unstrip[Do NOT attempt to unstrip the libc binary with debug symbols from a debuginfod server]' \
                '--offline-only[Attempt to searching with offline only mode]' \
                '*:HASH VALUE: '
              ;;
            (file)
              _arguments \
                '(-h --help)'{-h,--help}'[show help message and exit]' \
                '(-s --symbols)'{-s,--symbols}'[List of symbol offsets to dump in addition to the common ones]:SYMBOLS: ' \
                '(-o --offset)'{-o,--offset}'[Display all offsets relative to this symbol]:OFFSET: ' \
                '--unstrip[Attempt to unstrip the libc binary inplace with debug symbols from a debuginfod server]' \
                '*:FILE:_files'
              ;;
            (fetch)
              _arguments \
                '(-h --help)'{-h,--help}'[show help message and exit]' \
                '(-u --update)'{-u,--update}'[Fetch the desired libc categories]:UPDATE: ' \
                '1:PATH:_files' \
                '*: :'
              ;;
          }
          ;;
        phd)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-w --width)'{-w,--width}'[Number of bytes per line]:WIDTH INT: ' \
            '(-l --highlight)'{-l,--highlight}'[Byte to highlight]:*^<->:HIGHLIGHTS INT: ' \
            '(-s --skip)'{-s,--skip}'[Skip this many initial bytes]:SKIP INT: ' \
            '(-c --count)'{-c,--count}'[Only show this many bytes]:COUNT INT: ' \
            '(-o --offset)'{-o,--offset}'[Addresses in left hand column starts at this address]:OFFSET INT: ' \
            '--color[Colorize the output.  When "auto" output is colorized exactly when stdout is a TTY; default is auto]:COLOR:(always never auto)' \
            '1:FILE:_files' \
            '*: :'
            ;;
        pwnstrip)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-b --build-id)'{-b,--build-id}'[Strip build ID]' \
            '(-p --patch)'{-p,--patch}'[Patch function]:FUNCTION:->funcs' \
            '(-o --output)'{-o,--output}'[Output file]:OUTPUT:_files' \
            '1:FILE:_files' \
            '*: :'
          if [[ $state == 'funcs' ]] {
            _message "FUNCTION"
            if [[ $#line -ne 0 ]] {
              filename=$~line[1]
            } elif [[ -f $~words[-1] ]] {
              filename=$~words[-1]
            } else {
                return
            }
            _comp_funcs $filename
          }
          ;;
        scramble)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-f --format)'{-f,--format}'[Output format (defaults to hex for ttys, otherwise raw)]:FORMAT:(raw hex string elf)' \
            '(-o --output)'{-o,--output}'[Output file (defaults to stdout)]:OUTPUT:_files' \
            '(-c --context)'{-c,--context}'[The os/architecture/endianness/bits the shellcode will run in (default: linux/i386)]:CONTEXT:(16 32 64 baremetal freebsd windows android darwin linux cgc loongarch64 powerpc64 aarch64 powerpc riscv32 riscv64 sparc64 mips64 msp430 alpha amd64 sparc thumb cris i386 ia64 m68k mips s390 none avr arm vax little big be eb le el)' \
            '(-p --alphanumeric)'{-p,--alphanumeric}'[Encode the shellcode with an alphanumeric encoder]' \
            '(-v --avoid)'{-v,--avoid}'[Encode the shellcode to avoid the listed bytes]:AVOID: ' \
            '(-n --newline)'{-n,--newline}'[Encode the shellcode to avoid newlines]' \
            '(-z --zero)'{-z,--zero}'[Encode the shellcode to avoid NULL bytes]' \
            '(-d --debug)'{-d,--debug}'[Debug the shellcode with GDB]' \
            ;;
        shellcraft)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '(-? --show)'{-\?,--show}'[Show shellcode documentation]' \
            '(-o --out)'{-o,--out}'[Output file (default: stdout)]:OUTPUT:_files' \
            '(-f --format)'{-f,--format}'[Output format (default: hex)]:FORMAT:(e r s c h i a p d)' \
            '(-d --debug)'{-d,--debug}'[Debug the shellcode with GDB]' \
            '--delim[Set the delimiter between multiple shellcodes]:DELIM: ' \
            '(-b --before)'{-b,--before}'[Insert a debug trap before the code]' \
            '(-a --after)'{-a,--after}'[Insert a debug trap after the code]' \
            '(-v --avoid)'{-v,--avoid}'[Encode the shellcode to avoid the listed bytes]:AVOID: ' \
            '(-n --newline)'{-n,--newline}'[Encode the shellcode to avoid newlines]' \
            '(-z --zero)'{-z,--zero}'[Encode the shellcode to avoid NULL bytes]' \
            '(-r --run)'{-r,--run}'[Run output]' \
            '--color[Color output]' \
            '--no-color[Disable color output]' \
            '--syscalls[List syscalls]' \
            '--address[Load address]:ADDRESS: ' \
            '(-l --list)'{-l,--list}'[List available shellcodes, optionally provide a filter]' \
            '(-s --shared)'{-s,--shared}'[Generated ELF is a shared library]' \
            '*:SHELLCODE: '
            ;;
        template)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '--host[Remote host / SSH server]:HOST:_ssh_hosts ' \
            '--port[Remote port / SSH port]:PORT: ' \
            '--user[SSH Username]:USER:_ssh_users' \
            '(--pass --password)'{--pass,--password}'[SSH Password]:PASSWORD: ' \
            '--libc[Path to libc binary to use. If not given, the current directory is searched for a libc binary]:LIBC_PATH:_files' \
            '--path[Remote path of file on SSH server]:PATH:_files' \
            '--quiet[Less verbose template comments]' \
            '--color[Print the output in color]:COLOR:(never always auto)' \
            '--template[Path to a custom template]:TEMPLATE_PATH:_files' \
            '--no-auto[Do not automatically detect missing binaries]' \
            '1:EXE:_files' \
            '*: :'
            ;;
        unhex)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '*:HEX: '
            ;;
        update)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            '--install[Install the update automatically]' \
            '--pre[Check for pre-releases]' \
            ;;
        version)
          _arguments \
            '(-h --help)'{-h,--help}'[show help message and exit]' \
            ;;
      esac
      ;;
  esac
}
_pwn
