Commit 016e2d80 authored by Marius Gedminas's avatar Marius Gedminas

Make it easier to use standard payloads

This adds a command line option -l (--list-payloads) to list all the
standard payloads available in the 'pyrasite.payloads' package.
It also allows the user to refer to those using just the file name without
a directory, e.g.

    pyrasite $pid dump_stacks.py

It's possible to disambiguate standard payload names from local user files
with the same name by using explicit paths, e.g.

    pyrasite $pid ./dump_stacks.py  # runs user script, not standard payload
parent 92d46fa9
...@@ -47,17 +47,41 @@ def ptrace_check(): ...@@ -47,17 +47,41 @@ def ptrace_check():
print("") 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(): def main():
ptrace_check() ptrace_check()
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description='pyrasite - inject code into a running python process', description='pyrasite - inject code into a running python process',
epilog="For updates, visit https://github.com/lmacken/pyrasite") 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") 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" 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', parser.add_argument('--gdb-prefix', dest='gdb_prefix',
help='GDB prefix (if specified during installation)', help='GDB prefix (if specified during installation)',
default="") default="")
...@@ -70,19 +94,25 @@ def main(): ...@@ -70,19 +94,25 @@ def main():
args = parser.parse_args() args = parser.parse_args()
if args.list_payloads:
print("Available payloads:")
for payload in list_payloads():
print(" %s" % payload)
sys.exit()
try: try:
pid = int(args.pid) pid = int(args.pid)
except ValueError: except ValueError:
print("Error: The first argument must be a pid") print("Error: The first argument must be a pid")
sys.exit(2) sys.exit(2)
filename = args.filename filename = expand_payload(args.payload)
if filename: if filename:
if not os.path.exists(filename): if not os.path.exists(filename):
print("Error: Invalid path or file doesn't exist") print("Error: Invalid path or file doesn't exist")
sys.exit(3) sys.exit(3)
else: 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) sys.exit(4)
pyrasite.inject(pid, filename, verbose=args.verbose, pyrasite.inject(pid, filename, verbose=args.verbose,
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment