Allow pyrasite-shell to work across 2 docker containers

.travis.yml

@@ -1,16 +1,21 @@
 language: python
+sudo: true
 
 python:
   - 2.6
   - 2.7
-  - 3.2
   - 3.3
   - 3.4
+  - 3.5
+  - 3.6
+
+before_install:
+  - echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope
 
 install:
   - sudo apt-get install -qq gdb python-all-dbg
   - python setup.py install
-  - if [[ $TRAVIS_PYTHON_VERSION != '2.5' ]]; then pip install coveralls --use-mirrors && export HAS_COVERALLS=1; fi
+  - if [[ $TRAVIS_PYTHON_VERSION != '2.5' ]]; then pip install coveralls && export HAS_COVERALLS=1; fi
 
 script:
   - nosetests -v --with-coverage --cover-erase --cover-package=pyrasite --exclude=test_injecting_into_all_interpreters

pyrasite/ipc.py

@@ -64,15 +64,18 @@ class PyrasiteIPC(object):
     # shell payloads with netcat.
     reliable = True
 
-    def __init__(self, pid, reverse='ReversePythonConnection', timeout=5):
+    def __init__(self, pid, reverse='ReversePythonConnection', timeout=5, port=None, server_host='localhost', client_host='localhost', tmpdir=None):
         super(PyrasiteIPC, self).__init__()
         self.pid = pid
         self.sock = None
         self.server_sock = None
         self.hostname = None
-        self.port = None
+        self.port = port
         self.reverse = reverse
         self.timeout = float(timeout)
+        self.server_host = server_host
+        self.client_host = client_host
+        self.tmpdir = tmpdir
 
     def __enter__(self):
         self.connect()
@@ -109,7 +112,7 @@ def connect(self):
 
     def listen(self):
         """Listen on a random port"""
-        for res in socket.getaddrinfo('localhost', None, socket.AF_UNSPEC,
+        for res in socket.getaddrinfo(self.server_host, self.port, socket.AF_UNSPEC,
                                       socket.SOCK_STREAM, 0, 0):
             af, socktype, proto, canonname, sa = res
             try:
@@ -134,15 +137,18 @@ def listen(self):
 
     def create_payload(self):
         """Write out a reverse python connection payload with a custom port"""
-        (fd, filename) = tempfile.mkstemp()
+        (fd, filename) = tempfile.mkstemp(dir=self.tmpdir)
+        if platform.system() != 'Windows':
+            os.fchmod(fd, stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH)
         tmp = os.fdopen(fd, 'w')
         path = dirname(abspath(pyrasite.__file__))
         payload = open(join(path, 'reverse.py'))
 
         for line in payload.readlines():
             if line.startswith('#'):
                 continue
-            line = line.replace('port = 9001', 'port = %d' % self.port)
+            line = line.replace('port = 9001', 'port = {0}'.format(self.port))
+            line = line.replace('host = localhost', 'host = {0}'.format(self.client_host))
             if not self.reliable:
                 line = line.replace('reliable = True', 'reliable = False')
             tmp.write(line)
@@ -151,9 +157,6 @@ def create_payload(self):
         tmp.close()
         payload.close()
 
-        if platform.system() != 'Windows':
-            os.chmod(filename, stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH)
-
         return filename
 
     def inject(self):

pyrasite/tools/shell.py

@@ -15,26 +15,32 @@
 #
 # Copyright (C) 2011-2013 Red Hat, Inc., Luke Macken <[email protected]>
 
+import argparse
 import os
-import sys
 import pyrasite
 
 
 def shell():
-    """Open a Python shell in a running process"""
+    parser = argparse.ArgumentParser(description='Open a Python shell in a running process')
+    parser.add_argument('pid', type=int, help='PID of the running process to attach to.')
+    parser.add_argument('--timeout', type=int, default=5, help='IPC Timeout, applies for each command.')
+    parser.add_argument('--server-host', default='localhost', help='The hostname to use for the listening (server) side of the IPC communication. Sometimes, (docker), localhost does not work.')
+    parser.add_argument('--client-host', default='localhost', help='The hostname to use for the connecting (client) side of the IPC communication. Sometimes, (docker), localhost does not work.')
+    parser.add_argument('--port', type=int, default=None, help='Optionally specify a port that will be listened on for the remote process to attach back.')
+    parser.add_argument('--tmpdir', default=None, help='Use the following tmp directory for transferring the payload')
 
-    usage = "Usage: pyrasite-shell <PID>"
-    if not len(sys.argv) == 2:
-        print(usage)
-        sys.exit(1)
-    try:
-        pid = int(sys.argv[1])
-    except ValueError:
-        print(usage)
-        sys.exit(1)
+    args = parser.parse_args()
+
+    ipc = pyrasite.PyrasiteIPC(
+        args.pid,
+        'ReversePythonShell',
+        timeout=args.timeout,
+        server_host=args.server_host,
+        client_host=args.client_host,
+        port=args.port,
+        tmpdir=args.tmpdir,
+    )
 
-    ipc = pyrasite.PyrasiteIPC(pid, 'ReversePythonShell',
-                               timeout=os.getenv('PYRASITE_IPC_TIMEOUT') or 5)
     ipc.connect()
 
     print("Pyrasite Shell %s" % pyrasite.__version__)