Merge pull request #35 from abeutot/develop

Interactive interpreter improved.

pyrasite/ipc.py

@@ -61,13 +61,14 @@ class PyrasiteIPC(object):
     # shell payloads with netcat.
     reliable = True
 
-    def __init__(self, pid):
+    def __init__(self, pid, reverse='ReversePythonConnection'):
         super(PyrasiteIPC, self).__init__()
         self.pid = pid
         self.sock = None
         self.server_sock = None
         self.hostname = None
         self.port = None
+        self.reverse = reverse
 
     def __enter__(self):
         self.connect()
@@ -133,7 +134,7 @@ class PyrasiteIPC(object):
                 line = line.replace('reliable = True', 'reliable = False')
             tmp.write(line)
 
-        tmp.write('ReversePythonConnection().start()\n')
+        tmp.write('%s().start()\n' % self.reverse)
         tmp.close()
         payload.close()
 

pyrasite/reverse.py

@@ -23,6 +23,7 @@ import sys
 import socket
 import traceback
 import threading
+from code import InteractiveConsole
 
 if sys.version_info[0] == 3:
     from io import StringIO
@@ -113,3 +114,78 @@ class ReversePythonConnection(ReverseConnection):
         buffer.close()
         self.send(output)
         return True
+
+
+class DistantInteractiveConsole(InteractiveConsole):
+    def __init__(self, ipc):
+        InteractiveConsole.__init__(self)
+
+        self.ipc = ipc
+        self.set_buffer()
+
+    def set_buffer(self):
+        self.out_buffer = StringIO()
+        sys.stdout = sys.stderr = self.out_buffer
+
+    def unset_buffer(self):
+        sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__
+        value = self.out_buffer.getvalue()
+        self.out_buffer.close()
+
+        return value
+
+    def raw_input(self, prompt=""):
+        output = self.unset_buffer()
+        # payload format: 'prompt' ? '\n' 'output'
+        self.ipc.send('\n'.join((prompt, output)))
+
+        cmd = self.ipc.recv()
+
+        self.set_buffer()
+
+        return cmd
+
+
+class ReversePythonShell(threading.Thread, pyrasite.PyrasiteIPC):
+    """A reverse Python shell that behaves like Python interactive interpreter.
+
+    """
+
+    host = 'localhost'
+    port = 9001
+    reliable = True
+
+    def __init__(self, host=None, port=None):
+        super(ReversePythonShell, self).__init__()
+
+    def run(self):
+        try:
+            for res in socket.getaddrinfo(self.host, self.port,
+                    socket.AF_UNSPEC, socket.SOCK_STREAM):
+                af, socktype, proto, canonname, sa = res
+                try:
+                    self.sock = socket.socket(af, socktype, proto)
+                    try:
+                        self.sock.connect(sa)
+                    except socket.error:
+                        self.sock.close()
+                        self.sock = None
+                        continue
+                except socket.error:
+                    self.sock = None
+                    continue
+                break
+
+            if not self.sock:
+                raise Exception('pyrasite cannot establish reverse ' +
+                        'connection to %s:%d' % (self.host, self.port))
+
+            DistantInteractiveConsole(self).interact()
+
+        except SystemExit:
+            pass
+        except:
+            traceback.print_exc(file=sys.__stderr__)
+        finally:
+            sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__
+            self.close()

pyrasite/tools/shell.py

@@ -26,18 +26,46 @@ def shell():
         print("Usage: pyrasite-shell <PID>")
         sys.exit(1)
 
-    ipc = pyrasite.PyrasiteIPC(int(sys.argv[1]))
+    ipc = pyrasite.PyrasiteIPC(int(sys.argv[1]), 'ReversePythonShell')
     ipc.connect()
 
     print("Pyrasite Shell %s" % pyrasite.__version__)
     print("Connected to '%s'" % ipc.title)
-    print(ipc.cmd('import sys; print("Python " + sys.version + ' +
-                  '" on " + sys.platform)').strip())
+
+    prompt, payload = ipc.recv().split('\n', 1)
+    print(payload)
+
+    try:
+        import readline
+    except ImportError:
+        pass
+
+    # py3k compat
+    try:
+        input_ = raw_input
+    except NameError:
+        input_ = input
 
     try:
         while True:
-            print(ipc.cmd(raw_input('>>> ')))
+            try:
+                input_line = input_(prompt)
+            except EOFError:
+                input_line = 'exit()'
+                print('')
+            except KeyboardInterrupt:
+                input_line = 'None'
+                print('')
+
+            ipc.send(input_line)
+            payload = ipc.recv()
+            if payload is None:
+                break
+            prompt, payload = payload.split('\n', 1)
+            if payload != '':
+                print(payload)
     except:
         print('')
+        raise
 
     ipc.close()