Merge pull request #61 from TheJakeSchmidt/develop

Add the ability to control the IPC timeout.

pyrasite/ipc.py

@@ -64,14 +64,15 @@ class PyrasiteIPC(object):
     # shell payloads with netcat.
     reliable = True
 
-    def __init__(self, pid, reverse='ReversePythonConnection'):
+    def __init__(self, pid, reverse='ReversePythonConnection', timeout=5):
         super(PyrasiteIPC, self).__init__()
         self.pid = pid
         self.sock = None
         self.server_sock = None
         self.hostname = None
         self.port = None
         self.reverse = reverse
+        self.timeout = float(timeout)
 
     def __enter__(self):
         self.connect()
@@ -165,7 +166,7 @@ def wait(self):
         """Wait for the injected payload to connect back to us"""
         (clientsocket, address) = self.server_sock.accept()
         self.sock = clientsocket
-        self.sock.settimeout(5)
+        self.sock.settimeout(self.timeout)
         self.address = address
 
     def cmd(self, cmd):

pyrasite/main.py

@@ -92,6 +92,10 @@ def main():
                         help="Set where output is to be printed. 'procstreams'" 
                              " prints output in stdout/stderr of running process"
                              " and 'localterm' prints output in local terminal.")
+    parser.add_argument('--ipc-timeout', dest='ipc_timeout', default=5,
+                        action='store', type=int,
+                        help="The number of seconds to wait for the injected"
+                             " code to reply over IPC before giving up.")
 
     if len(sys.argv) == 1:
         parser.print_help()
@@ -129,7 +133,8 @@ def main():
 
     if args.output_type == 'localterm':
         # Create new IPC connection to the process.
-        ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonConnection')
+        ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonConnection',
+                                   timeout=ipc_timeout)
         ipc.connect()
         print("Pyrasite Shell %s" % pyrasite.__version__)
         print("Connected to '%s'" % ipc.title)

pyrasite/tools/shell.py

@@ -15,6 +15,7 @@
 #
 # Copyright (C) 2011-2013 Red Hat, Inc., Luke Macken <[email protected]>
 
+import os
 import sys
 import pyrasite
 
@@ -32,7 +33,8 @@ def shell():
         print(usage)
         sys.exit(1)
 
-    ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonShell')
+    ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonShell',
+                               timeout=os.getenv('PYRASITE_IPC_TIMEOUT') or 5)
     ipc.connect()
 
     print("Pyrasite Shell %s" % pyrasite.__version__)