Merge pull request #53 from dustymabe/localtermoutput

Local term output

docs/CLI.rst

@@ -3,19 +3,31 @@
 
 ::
 
-    usage: pyrasite [-h] [--gdb-prefix GDB_PREFIX] [--verbose] pid [filename]
+    usage: pyrasite [-h] [--gdb-prefix GDB_PREFIX] [--verbose] [--output OUTPUT_TYPE] pid [filepath|payloadname]
+           pyrasite --list-payloads
 
     pyrasite - inject code into a running python process
 
     positional arguments:
       pid                   The ID of the process to inject code into
-      filename              The second argument must be a filename
+      filepath|payloadname  The second argument must be a path to a
+                            file that will be sent as a payload to the
+                            target process or it must be the name of
+                            an existing payload (see --list-payloads).
 
     optional arguments:
       -h, --help            show this help message and exit
       --gdb-prefix GDB_PREFIX
                             GDB prefix (if specified during installation)
       --verbose             Verbose mode
+      --output OUTPUT_TYPE  This option controls where the output from
+                            the executed payload will be printed. If
+                            the value is 'procstreams' (the default) then 
+                            the output is sent to the stdout/stderr of the
+                            process. If the value is 'localterm' then the
+                            output is piped back and printed on the local
+                            terminal where pyrasite is being run.
+      --list-payloads       List payloads that are delivered by pyrasite
 
     For updates, visit https://github.com/lmacken/pyrasite
 

pyrasite/main.py

@@ -87,6 +87,11 @@ def main():
                         default="")
     parser.add_argument('--verbose', dest='verbose', help='Verbose mode',
                         default=False, action='store_const', const=True)
+    parser.add_argument('--output', dest='output_type', default='procstreams',
+                        action='store',
+                        help="Set where output is to be printed. 'procstreams'" 
+                             " prints output in stdout/stderr of running process"
+                             " and 'localterm' prints output in local terminal.")
 
     if len(sys.argv) == 1:
         parser.print_help()
@@ -100,6 +105,11 @@ def main():
             print("  %s" % payload)
         sys.exit()
 
+    # Make sure the output type is valid (procstreams || localterm)
+    if args.output_type != 'procstreams' and args.output_type != 'localterm':
+        print("Error: --output arg must be 'procstreams' or 'localterm'")
+        sys.exit(5)
+
     try:
         pid = int(args.pid)
     except ValueError:
@@ -115,8 +125,26 @@ def main():
         print("Error: The second argument must be a filename or a payload name")
         sys.exit(4)
 
-    pyrasite.inject(pid, filename, verbose=args.verbose,
-                    gdb_prefix=args.gdb_prefix)
+        
+    
+    if args.output_type == 'localterm':
+        # Create new IPC connection to the process.
+        ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonConnection')
+        ipc.connect()
+        print("Pyrasite Shell %s" % pyrasite.__version__)
+        print("Connected to '%s'" % ipc.title)
+
+        # Read in the payload
+        fd = open(filename)
+        payload = fd.read()
+        fd.close
+
+        # Run the payload, print output, close ipc connection
+        print(ipc.cmd(payload))
+        ipc.close()
+    else:
+        pyrasite.inject(pid, filename, verbose=args.verbose,
+                        gdb_prefix=args.gdb_prefix)
 
 
 if __name__ == '__main__':