Merge branch 'std-payloads' of https://github.com/mgedmin/pyrasite into mgedmin-std-payloads

pyrasite/main.py

@@ -47,17 +47,41 @@ def ptrace_check():
                 print("")
 
 
+def get_payload_dir():
+    return os.path.join(os.path.dirname(pyrasite.__file__), 'payloads')
+
+
+def list_payloads():
+    return sorted(fn for fn in os.listdir(get_payload_dir())
+                  if fn.endswith('.py') and not fn.startswith('_'))
+
+def expand_payload(payload):
+    """If a standard payload with this name exists, return its full path.
+
+    Otherwise return the input value unchanged.
+    """
+    if os.path.sep not in payload:
+        fn = os.path.join(get_payload_dir(), payload)
+        if os.path.isfile(fn):
+            return fn
+    return payload
+
+
 def main():
     ptrace_check()
 
     parser = argparse.ArgumentParser(
             description='pyrasite - inject code into a running python process',
             epilog="For updates, visit https://github.com/lmacken/pyrasite")
-    parser.add_argument('pid',
+
+    parser.add_argument('pid', nargs='?',
                         help="The ID of the process to inject code into")
-    parser.add_argument('filename',
+    parser.add_argument('payload', nargs='?', default='',
                         help="The Python script to be executed inside the"
-                             " running process, also referred to as 'payload'")
+                             " running process.  Can be one of the standard"
+                             " payloads (see --list-payloads) or a filname.")
+    parser.add_argument('-l', '--list-payloads', help='List standard payloads',
+                        default=False, action='store_const', const=True)
     parser.add_argument('--gdb-prefix', dest='gdb_prefix',
                         help='GDB prefix (if specified during installation)',
                         default="")
@@ -70,19 +94,25 @@ def main():
 
     args = parser.parse_args()
 
+    if args.list_payloads:
+        print("Available payloads:")
+        for payload in list_payloads():
+            print("  %s" % payload)
+        sys.exit()
+
     try:
         pid = int(args.pid)
     except ValueError:
         print("Error: The first argument must be a pid")
         sys.exit(2)
 
-    filename = args.filename
+    filename = expand_payload(args.payload)
     if filename:
         if not os.path.exists(filename):
             print("Error: Invalid path or file doesn't exist")
             sys.exit(3)
     else:
-        print("Error: The second argument must be a filename")
+        print("Error: The second argument must be a filename or a payload name")
         sys.exit(4)
 
     pyrasite.inject(pid, filename, verbose=args.verbose,