Refactor command-line tools to use a single "pwn" entry point.

docs/source/commandline.rst

@@ -14,45 +14,5 @@ pwntools comes with a handful of useful command-line utilities which serve as wr
 
 .. toctree::
 
-.. autoprogram:: pwnlib.commandline.asm:parser
-   :prog: asm
-
-.. autoprogram:: pwnlib.commandline.checksec:parser
-   :prog: checksec
-
-.. autoprogram:: pwnlib.commandline.constgrep:p
-   :prog: constgrep
-
-.. autoprogram:: pwnlib.commandline.cyclic:parser
-   :prog: cyclic
-
-.. autoprogram:: pwnlib.commandline.disasm:parser
-   :prog: disasm
-
-.. autoprogram:: pwnlib.commandline.elfdiff:p
-   :prog: elfdiff
-
-.. autoprogram:: pwnlib.commandline.elfpatch:p
-   :prog: elfpatch
-
-.. autoprogram:: pwnlib.commandline.errno:parser
-   :prog: errno
-
-.. autoprogram:: pwnlib.commandline.hex:parser
-   :prog: hex
-
-.. autoprogram:: pwnlib.commandline.phd:parser
-   :prog: phd
-
-.. autoprogram:: pwnlib.commandline.pwnstrip:p
-   :prog: pwnstrip
-
-.. autoprogram:: pwnlib.commandline.scramble:parser
-   :prog: scramble
-
-.. autoprogram:: pwnlib.commandline.shellcraft:p
-   :prog: shellcraft
-
-.. autoprogram:: pwnlib.commandline.unhex:parser
-   :prog: unhex
-
+.. autoprogram:: pwnlib.commandline.main:parser
+   :prog: pwn

extra/bash_completion.d/shellcraft

@@ -43,4 +43,4 @@ _shellcraft()
     COMPREPLY=( $(compgen -W "${_shellcraft_shellcodes}" -- ${cur}) )
 }
 
-complete -F _shellcraft shellcraft
+complete -F _shellcraft pwn shellcraft

pwnlib/commandline/__init__.py

@@ -1 +1 @@
-# empty
+#!/usr/bin/env python2

pwnlib/commandline/asm.py

@@ -6,8 +6,9 @@
 
 from . import common
 
-parser = argparse.ArgumentParser(
-    description = 'Assemble shellcode into bytes'
+parser = common.parser_commands.add_parser(
+    'asm',
+    help = 'Assemble shellcode into bytes'
 )
 
 parser.add_argument(
@@ -91,8 +92,7 @@
     action='store_true'
 )
 
-def main():
-    args   = parser.parse_args()
+def main(args):
     tty    = args.output.isatty()
 
     data   = '\n'.join(args.lines) or args.infile.read()
@@ -123,4 +123,5 @@ def main():
     if tty and fmt is not 'raw':
         args.output.write('\n')
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/checksec.py

@@ -6,11 +6,10 @@
 
 from . import common
 
-parser = argparse.ArgumentParser(
-    description = 'Check binary security settings'
+parser = common.parser_commands.add_parser(
+    'checksec',
+    help = 'Check binary security settings'
 )
-
-
 parser.add_argument(
     'elf',
     nargs='*',
@@ -26,8 +25,7 @@
     help='File to check (for compatibility with checksec.sh)'
 )
 
-def main():
-    args   = parser.parse_args()
+def main(args):
     files  = args.elf or args.elf2 or []
 
     if not files:
@@ -37,4 +35,5 @@ def main():
     for f in files:
         e = ELF(f.name)
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/common.py

@@ -1,3 +1,5 @@
+import argparse
+import os
 import sys
 
 import pwnlib
@@ -20,3 +22,12 @@ def context_arg(arg):
     try: context.endian = arg
     except Exception: pass
     return arg
+
+parser = argparse.ArgumentParser(description='Pwntools Command-line Interface')
+parser_commands = parser.add_subparsers(dest='command')
+
+def main(file=sys.argv[0]):
+    import pwnlib.commandline.main
+    sys.argv.insert(0, 'pwn')
+    sys.argv[1] = os.path.splitext(os.path.basename(file))[0]
+    pwnlib.commandline.main.main()

pwnlib/commandline/constgrep.py

@@ -7,8 +7,9 @@
 from pwn import *
 from . import common
 
-p = argparse.ArgumentParser(
-    description = "Looking up constants from header files.\n\nExample: constgrep -c freebsd -m  ^PROT_ '3 + 4'",
+p = common.parser_commands.add_parser(
+    'constgrep',
+    help = "Looking up constants from header files.\n\nExample: constgrep -c freebsd -m  ^PROT_ '3 + 4'",
     formatter_class = argparse.RawDescriptionHelpFormatter,
 )
 
@@ -55,9 +56,7 @@
     help = 'The os/architecture/endianness/bits the shellcode will run in (default: linux/i386), choose from: %s' % common.choices,
 )
 
-def main():
-    args = p.parse_args()
-
+def main(args):
     if args.exact:
         # This is the simple case
         print cpp(args.exact).strip()
@@ -134,4 +133,5 @@ def main():
                 print
                 print '(%s) == %s' % (' | '.join(k for v, k in good), args.constant)
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/cyclic.py

@@ -7,8 +7,9 @@
 
 from . import common
 
-parser = argparse.ArgumentParser(
-    description = "Cyclic pattern creator/finder"
+parser = common.parser_commands.add_parser(
+    'cyclic',
+    help = "Cyclic pattern creator/finder"
 )
 
 parser.add_argument(
@@ -50,8 +51,7 @@
     help = 'Number of characters to print'
 )
 
-def main():
-    args = parser.parse_args()
+def main(args):
     alphabet = args.alphabet
     subsize  = args.length
 
@@ -90,4 +90,5 @@ def main():
         if sys.stdout.isatty():
             sys.stdout.write('\n')
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/disasm.py

@@ -7,8 +7,9 @@
 
 from . import common
 
-parser = argparse.ArgumentParser(
-    description = 'Disassemble bytes into text format'
+parser = common.parser_commands.add_parser(
+    'disasm',
+    help = 'Disassemble bytes into text format'
 )
 
 parser.add_argument(
@@ -52,9 +53,7 @@
 )
 
 
-def main():
-    args = parser.parse_args()
-
+def main(args):
     if len(args.hex) > 0:
         dat = ''.join(args.hex)
         dat = dat.translate(None, string.whitespace)
@@ -88,4 +87,5 @@ def main():
 
     print disasm(dat, vma=safeeval.const(args.address))
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/elfdiff.py

@@ -8,6 +8,8 @@
 
 from pwn import *
 
+from . import common
+
 def dump(objdump, path):
     n = NamedTemporaryFile(delete=False)
     o = check_output([objdump,'-d','-x','-s',path])
@@ -20,13 +22,15 @@ def diff(a,b):
     except CalledProcessError as e:
         return e.output
 
-p = ArgumentParser()
+p = common.parser_commands.add_parser(
+    'elfdiff',
+    help = 'Compare two ELF files'
+)
+
 p.add_argument('a')
 p.add_argument('b')
 
-def main():
-    a = p.parse_args()
-
+def main(a):
     with context.silent:
         x = ELF(a.a)
         y = ELF(a.b)
@@ -49,4 +53,5 @@ def main():
 
     print diff(x, y)
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/elfpatch.py

@@ -4,15 +4,20 @@
 
 from pwn import *
 
+from . import common
+
+p = common.parser_commands.add_parser(
+    'elfpatch',
+    help = 'Patch an ELF file'
+)
+
 p = argparse.ArgumentParser()
 p.add_argument('elf',help="File to patch")
 p.add_argument('offset',help="Offset to patch in virtual address (hex encoded)")
 p.add_argument('bytes',help='Bytes to patch (hex encoded)')
 
 
-def main():
-    a = p.parse_args()
-
+def main(a):
     if not a.offset.startswith('0x'):
         a.offset = '0x' + a.offset
 
@@ -25,4 +30,5 @@ def main():
     elf.write(offset, bytes)
     sys.stdout.write(elf.get_data())
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/errno.py

@@ -1,17 +1,18 @@
 import argparse
 import os
 
-parser = argparse.ArgumentParser(
-    description = 'Prints out error messages'
+from . import common
+
+parser = common.parser_commands.add_parser(
+    'errno',
+    help = 'Prints out error messages'
 )
 
 parser.add_argument(
     'error', help='Error message or value', type=str
 )
 
-def main():
-  args, unknown = parser.parse_known_args()
-
+def main(args):
   try:
     value = int(args.error, 0)
 
@@ -40,4 +41,5 @@ def main():
   print '#define', name, value
   print os.strerror(value)
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/hex.py

@@ -2,17 +2,21 @@
 import argparse
 import sys
 
-parser = argparse.ArgumentParser(description='''
-Hex-encodes data provided on the command line or via stdin.
+from . import common
+
+parser = common.parser_commands.add_parser(
+    'hex',
+    help = '''
+Hex-encodes data provided on the command line or stdin
 ''')
 parser.add_argument('data', nargs='*',
     help='Data to convert into hex')
 
-def main():
-    args = parser.parse_args()
+def main(args):
     if not args.data:
         print sys.stdin.read().encode('hex')
     else:
         print ' '.join(args.data).encode('hex')
 
-if __name__ == '__main__': main()
+if __name__ == '__main__':
+    pwnlib.common.main(__file__)

pwnlib/commandline/main.py

@@ -0,0 +1,40 @@
+from .common import parser
+from . import asm
+from . import checksec
+from . import common
+from . import constgrep
+from . import cyclic
+from . import disasm
+from . import elfdiff
+from . import elfpatch
+from . import errno
+from . import hex
+from . import phd
+from . import pwnstrip
+from . import scramble
+from . import shellcraft
+from . import unhex
+
+commands = {
+    'asm': asm.main,
+    'checksec': checksec.main,
+    'constgrep': constgrep.main,
+    'cyclic': cyclic.main,
+    'disasm': disasm.main,
+    'elfdiff': elfdiff.main,
+    'elfpatch': elfpatch.main,
+    'errno': errno.main,
+    'hex': hex.main,
+    'phd': phd.main,
+    'pwnstrip': pwnstrip.main,
+    'scramble': scramble.main,
+    'shellcraft': shellcraft.main,
+    'unhex': unhex.main,
+}
+
+def main():
+    args = parser.parse_args()
+    commands[args.command](args)
+
+if __name__ == '__main__':
+    main()