From e84c8fb3b0f97e8b43a3628328d26d91fe3151f7 Mon Sep 17 00:00:00 2001
From: Stephen Arthur <sarth@yelp.com>
Date: Mon, 11 Jun 2018 15:35:50 -0700
Subject: [PATCH] Make a few changes so that the server / client processes of
 pyrasite-shell can be in two different docker containers. We need a
 deterministic listen-back port so the connecting (debugging) container can
 provide it at launch. We need to write the payload file to a temporary
 location and it should have others read bit set so a container that dropped
 permissions can still read it.

There are a few settings required in the docker run command of the debugging container process, but not other requirements for the container under inspection.
---
 pyrasite/ipc.py         | 19 +++++++++++--------
 pyrasite/tools/shell.py | 32 +++++++++++++++++++-------------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/pyrasite/ipc.py b/pyrasite/ipc.py
index 780e8f4..8f689aa 100644
--- a/pyrasite/ipc.py
+++ b/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,7 +137,9 @@ 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'))
@@ -142,7 +147,8 @@ def create_payload(self):
         for line in payload.readlines():
             if line.startswith('#'):
                 continue
-            line = line.replace('port = 9001', 'port = %d' % self.port)
+            line = line.replace('port = 9001', 'port = {}'.format(self.port))
+            line = line.replace('host = localhost', 'host = {}'.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):
diff --git a/pyrasite/tools/shell.py b/pyrasite/tools/shell.py
index 1de50a0..8732e13 100644
--- a/pyrasite/tools/shell.py
+++ b/pyrasite/tools/shell.py
@@ -15,26 +15,32 @@
 #
 # Copyright (C) 2011-2013 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
 
+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__)