Make 'pwn template' better

pwnlib/commandline/template.py

@@ -19,6 +19,7 @@
 parser.add_argument('--user', help='SSH Username')
 parser.add_argument('--pass', help='SSH Password', dest='password')
 parser.add_argument('--path', help='Remote path of file on SSH server')
+parser.add_argument('--quiet', help='Less verbose template comments', action='store_true')
 
 def main(args):
     cache = None
@@ -49,7 +50,8 @@ def main(args):
                              args.port,
                              args.user,
                              args.password,
-                             args.path)
+                             args.path,
+                             args.quiet)
 
     # Fix Mako formatting bs
     output = re.sub('\n\n\n', '\n\n', output)

pwnlib/data/templates/pwnup.mako

@@ -1,13 +1,19 @@
-<%page args="binary, host=None, port=None, user=None, password=None, remote_path=None"/>\
+<%page args="binary, host=None, port=None, user=None, password=None, remote_path=None, quiet=False"/>\
 <%
+import os
+import sys
+
 from pwnlib.context import context as ctx
 from pwnlib.elf.elf import ELF
+from pwnlib.util.sh_string import sh_string
 from elftools.common.exceptions import ELFError
 
-import os
+argv = list(sys.argv)
+argv[0] = os.path.basename(argv[0])
+
 try:
     if binary:
-        ctx.binary = ELF(binary, checksec=False)
+       ctx.binary = ELF(binary, checksec=False)
 except ELFError:
     pass
 
@@ -28,9 +34,15 @@ binary_repr = repr(binary)
 %>\
 #!/usr/bin/env python2
 # -*- coding: utf-8 -*-
+%if not quiet:
+# This exploit template was generated via:
+# $ ${' '.join(map(sh_string, argv))}
+%endif
 from pwn import *
 
+%if not quiet:
 # Set up pwntools for the correct architecture
+%endif
 %if ctx.binary:
 exe = context.binary = ELF(${binary_repr})
 <% binary_repr = 'exe.path' %>
@@ -40,13 +52,15 @@ exe = ${binary_repr}
 <% binary_repr = 'exe' %>
 %endif
 
+%if not quiet:
 # Many built-in settings can be controlled on the command-line and show up
 # in "args".  For example, to dump all data sent/received, and disable ASLR
 # for all created processes...
 # ./exploit.py DEBUG NOASLR
 %if host or port or user:
 # ./exploit.py GDB HOST=example.com PORT=4141
 %endif
+%endif
 %if host:
 host = args.HOST or ${repr(host)}
 %endif
@@ -61,22 +75,6 @@ password = args.PASSWORD or ${repr(password)}
 remote_path = ${repr(remote_path)}
 %endif
 
-%if exe or remote_path:
-# Specify your GDB script here for debugging
-# GDB will be launched if the exploit is run via e.g.
-# ./exploit.py GDB
-gdbscript = '''
-%if ctx.binary:
-  %if 'main' in ctx.binary.symbols:
-break *0x{exe.symbols.main:x}
-  %else:
-break *0x{exe.entry:x}
-  %endif
-%endif
-continue
-'''.format(**locals())
-%endif
-
 %if ssh:
 # Connect to the remote SSH server
 shell = None
@@ -86,25 +84,22 @@ if not args.LOCAL:
 %endif
 
 %if host:
-# Execute the target binary locally
 def local(argv=[], *a, **kw):
+    '''Execute the target binary locally'''
     if args.GDB:
-        return gdb.debug([exe.path] + argv, gdbscript=gdbscript, *a, **kw)
+        return gdb.debug([${binary_repr}] + argv, gdbscript=gdbscript, *a, **kw)
     else:
-        return process([exe.path] + argv, *a, **kw)
+        return process([${binary_repr}] + argv, *a, **kw)
 
-%if ssh:
-# Execute the target binary on the remote host
-%else:
-# Connect to the process on the remote host
-%endif
 def remote(argv=[], *a, **kw):
   %if ssh:
+    '''Execute the target binary on the remote host'''
     if args.GDB:
         return gdb.debug([remote_path] + argv, gdbscript=gdbscript, ssh=shell, *a, **kw)
     else:
         return shell.process([remote_path] + argv, *a, **kw)
   %else:
+    '''Connect to the process on the remote host'''
     io = connect(host, port)
     if args.GDB:
         gdb.attach(io, gdbscript=gdbscript)
@@ -113,20 +108,60 @@ def remote(argv=[], *a, **kw):
 %endif
 
 %if host:
-start = local if args.LOCAL else remote
+def start(argv=[], *a, **kw):
+    '''Start the exploit against the target.'''
+    if args.LOCAL:
+        return local(argv, *a, **kw)
+    else:
+        return remote(argv, *a, **kw)
 %else:
 def start(argv=[], *a, **kw):
+    '''Start the exploit against the target.'''
     if args.GDB:
         return gdb.debug([${binary_repr}] + argv, gdbscript=gdbscript, *a, **kw)
     else:
         return process([${binary_repr}] + argv, *a, **kw)
 %endif
 
+%if exe or remote_path:
+%if not quiet:
+# Specify your GDB script here for debugging
+# GDB will be launched if the exploit is run via e.g.
+# ./exploit.py GDB
+%endif
+gdbscript = '''
+%if ctx.binary:
+  %if 'main' in ctx.binary.symbols:
+break *0x{exe.symbols.main:x}
+  %else:
+break *0x{exe.entry:x}
+  %endif
+%endif
+continue
+'''.format(**locals())
+%endif
+
+
+%if not quiet:
 #===========================================================
 #                    EXPLOIT GOES HERE
 #===========================================================
+%else:
+# -- Exploit goes here --
+%endif
+%if ctx.binary and not quiet:
+# ${'%-10s%s-%s-%s' % ('Arch:',
+                       ctx.binary.arch,
+                       ctx.binary.bits,
+                       ctx.binary.endian)}
+%for line in ctx.binary.checksec(color=False).splitlines():
+# ${line}
+%endfor
+%endif
+
 io = start()
 
+%if not quiet:
 # shellcode = asm(shellcraft.sh())
 # payload = fit({
 #     32: 0xdeadbeef,
@@ -135,5 +170,6 @@ io = start()
 # io.send(payload)
 # flag = io.recv(...)
 # log.success(flag)
+%endif
 
 io.interactive()

pwnlib/elf/elf.py

@@ -408,10 +408,10 @@ def debug(self, argv=[], *a, **kw):
         import pwnlib.gdb
         return pwnlib.gdb.debug([self.path] + argv, *a, **kw)
 
-    def _describe(self):
+    def _describe(self, *a, **kw):
         log.info_once('\n'.join((repr(self.path),
                                 '%-10s%s-%s-%s' % ('Arch:', self.arch, self.bits, self.endian),
-                                self.checksec())))
+                                self.checksec(*a, **kw))))
 
     def __repr__(self):
         return "ELF(%r)" % self.path
@@ -1490,17 +1490,17 @@ def runpath(self):
 
         return self.dynamic_string(dt_rpath.entry.d_ptr)
 
-    def checksec(self, banner=True):
+    def checksec(self, banner=True, color=True):
         """checksec(banner=True)
 
         Prints out information in the binary, similar to ``checksec.sh``.
 
         Arguments:
             banner(bool): Whether to print the path to the ELF binary.
         """
-        red    = text.red
-        green  = text.green
-        yellow = text.yellow
+        red    = text.red if color else str
+        green  = text.green if color else str
+        yellow = text.yellow if color else str
 
         res = []
 

pwnlib/util/sh_string.py

@@ -393,7 +393,7 @@ def sh_string(s):
         return "''"
 
     chars = set(s)
-    very_good = set(string.ascii_letters + string.digits + "_+.,/")
+    very_good = set(string.ascii_letters + string.digits + "_+.,/-")
 
     # Alphanumeric can always just be used verbatim.
     if chars <= very_good: