Some serious refactoring, documenting, and optimizing.

- Move the low level socket code into the PyrasiteIPC for re-use throughout
  our GUI and payloads
- Create an 'unreliable' data transfer mode that lets payloads, like
  the reverse shells, work properly with netcat on the other end.
  By default, the IPC will prepend a length header to the data, to
  ensure reliability.
- Pull our ReverseConnection and ReversePythonConnection out into a new
  pyrasite.reverse module

pyrasite/inject.py

@@ -34,6 +34,7 @@ import warnings
 from utils import run
 
 class CodeInjector(object):
+    """Injects code into a running Python process"""
 
     def __init__(self, pid, filename=None, verbose=False, gdb_prefix=""):
         self.pid = pid
@@ -46,6 +47,7 @@ class CodeInjector(object):
             self.filename = os.path.abspath(filename)
 
     def inject(self, filename=None):
+        """Inject a given file into `self.pid` using gdb"""
         if filename:
             self.filename = os.path.abspath(filename)
         gdb_cmds = [

pyrasite/inspect.py

@@ -24,6 +24,7 @@ class ObjectInspector(object):
         self.pid = pid
 
     def inspect(self, address):
+        """Return the value of an object at a given address"""
         cmd = ' '.join([
             'gdb --quiet -p %s -batch' % self.pid,
             '-eval-command="print (PyObject *)%s"' % address,

pyrasite/ipc.py

@@ -13,124 +13,149 @@
 # You should have received a copy of the GNU General Public License
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright (C) 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
+# Copyright (C) 2011, 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
 """
-:mod:`pyrasite.ipc` - Pyrasite Interprocess Communication
+:mod:`pyrasite.ipc` - Pyrasite Inter-Python Communication
 =========================================================
-
-This module contains :class:`PyrasiteIPC`, which handles injecting a reverse
-Python shell into a process, which connects back to us on a custom port that it
-is listening on. You can then execute commands in the process using
-:meth:`PyrastieIPC.cmd`, and get the stdout+stderr back.
-
 """
 
 import os
 import socket
 import struct
-import logging
 import tempfile
 import pyrasite
 
-log = logging.getLogger(__name__)
-
-REVERSE_SHELL = """\
-import sys, struct
-sys.path.insert(0, "%s/../payloads/")
-
-from StringIO import StringIO
-from _reverseconnection import ReverseConnection
-
-class ReversePythonShell(ReverseConnection):
-    host = 'localhost'
-    port = %d
-
-    def on_command(self, s, cmd):
-        buffer = StringIO()
-        sys.stdout = buffer
-        sys.stderr = buffer
-        output = ''
-        try:
-            exec(cmd)
-            output = buffer.getvalue()
-        except Exception, e:
-            output = str(e)
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            buffer.close()
-        header = struct.pack('<L', len(output))
-        s.sendall(header + output)
-        return True
-
-ReversePythonShell().start()
-"""
-
+from os.path import dirname, abspath, join
 
 class PyrasiteIPC(object):
+    """Pyrasite Inter-Python Communication.
+
+    This object is used in communicating to or from another Python process.
+
+    It can perform a variety of tasks:
+
+    - Injection of the :class:`pyrasite.ReversePythonConnection` payload via
+      :meth:`PyrasiteIPC.connect()`, which causes the process to connect back
+      to a port that we are listening on. The connection with the process is
+      then available via `self.sock`.
+
+    - Python code can then be executed in the process using
+      :meth:`PyrasiteIPC.cmd`. Both stdout and stderr are returned.
+
+    - Low-level communication with the process, both reliably (via a length
+      header) or unreliably (raw data, ideal for use with netcat) with a
+      :class:`pyrasite.ReversePythonConnection` payload, via
+      :meth:`PyrasiteIPC.send(data)` and :meth:`PyrasiteIPC.recv(data)`.
+
+    The :class:`PyrasiteIPC` is subclassed by
+    :class:`pyrasite.tools.gui.Process` as well as
+    :class:`pyrasite.reverse.ReverseConnection`.
+
     """
-    An object that listens for connections from the reverse python shell payload,
-    and then allows you to run commands in the other process.
-    """
+    # Allow subclasses to disable this and just send/receive raw data, as
+    # opposed to prepending a length header, to ensure reliability. The reason
+    # to enable 'unreliable' connections is so we can still use our reverse
+    # shell payloads with netcat.
+    reliable = True
+
     def __init__(self, pid):
         super(PyrasiteIPC, self).__init__()
         self.pid = pid
-        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.sock.settimeout(5)
-        self.sock.bind(('localhost', 0))
-        self.sock.listen(20)
-        self.port = self.sock.getsockname()[1]
-        self.client = None
+        self.sock = None
+
+    def connect(self):
+        """
+        Setup a communication socket with the process by injecting
+        a reverse subshell and having it connect back to us.
+        """
+        self.listen()
+        self.inject()
+        self.wait()
+
+    def listen(self):
+        """Listen on a random port"""
+        self.server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.server_sock.settimeout(5)
+        self.server_sock.bind(('localhost', 0))
+        self.server_sock.listen(1)
+        self.port = self.server_sock.getsockname()[1]
         self.running = True
 
-    def inject(self):
-        # Write out a reverse subprocess payload with a custom port
+    def create_payload(self):
+        """Write out a reverse python connection payload with a custom port"""
         (fd, filename) = tempfile.mkstemp()
-        self.filename = filename
         tmp = os.fdopen(fd, 'w')
-        tmp.write(REVERSE_SHELL % (
-            os.path.abspath(os.path.dirname(pyrasite.__file__)),
-            self.port))
+        path = dirname(abspath(join(pyrasite.__file__, '..')))
+        payload = file(join(path, 'pyrasite', 'reverse.py'))
+        tmp.write('import sys; sys.path.insert(0, "%s")\n' % path)
+
+        for line in payload.readlines():
+            if line.startswith('#'):
+                continue
+            line = line.replace('port = 9001', 'port = %d' % self.port)
+            line = line.replace('reliable = False', 'reliable = True')
+            tmp.write(line)
+
+        tmp.write('ReversePythonConnection().start()\n')
         tmp.close()
+        payload.close()
 
+        return filename
+
+    def inject(self):
+        """Inject the payload into the process."""
+        filename = self.create_payload()
         injector = pyrasite.CodeInjector(self.pid)
         injector.inject(filename)
+        os.unlink(filename)
 
-    def listen(self):
-        (clientsocket, address) = self.sock.accept()
-        self.client = clientsocket
-        self.client.settimeout(3)
+    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.address = address
 
     def cmd(self, cmd):
-        self.client.sendall(cmd + '\n')
-        try:
-            header_data = self._recv_bytes(4)
+        """Send a python command to exec in the process and return the output"""
+        self.send(cmd + '\n')
+        return self.recv()
+
+    def send(self, data):
+        """Send arbitrary data to the process via self.sock"""
+        header = ''
+        if self.reliable:
+            header = struct.pack('<L', len(data))
+        self.sock.sendall(header + data)
+
+    def recv(self):
+        """Receive a command from a given socket"""
+        if self.reliable:
+            header_data = self.recv_bytes(4)
             if len(header_data) == 4:
                 msg_len = struct.unpack('<L', header_data)[0]
-                data = self._recv_bytes(msg_len)
+                data = self.recv_bytes(msg_len)
                 if len(data) == msg_len:
                     return data
-                else:
-                    log.error("Response doesn't match header len (%s) : %r" % (
-                          msg_len, data))
-        except Exception, e:
-            log.exception(e)
-            self.close()
-
-    def _recv_bytes(self, n):
+        else:
+            return self.sock.recv(4096)
+
+    def recv_bytes(self, n):
+        """Receive n bytes from a socket"""
         data = ''
         while len(data) < n:
-            chunk = self.client.recv(n - len(data))
+            chunk = self.sock.recv(n - len(data))
             if chunk == '':
                 break
             data += chunk
         return data
 
     def close(self):
-        os.unlink(self.filename)
-        if self.client:
-            self.client.sendall('exit\n')
-            self.client.close()
+        if self.sock:
+            self.sock.close()
 
     def __repr__(self):
         return "<%s %s>" % (self.__class__.__name__, self.pid)
+
+    def __str__(self):
+        return self.title

pyrasite/payloads/reverse_python_shell.py

@@ -13,37 +13,16 @@
 # You should have received a copy of the GNU General Public License
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright (C) 2011 Red Hat, Inc.
+# Copyright (C) 2011, 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
 
 import sys
+import pyrasite
 
-from StringIO import StringIO
-
-from _reverseconnection import ReverseConnection
-
-class ReversePythonShell(ReverseConnection):
-
-    host = '127.0.0.1'
+class ReversePythonShell(pyrasite.ReversePythonConnection):
     port = 9001
+    reliable = False
 
-    def on_connect(self, s):
-        s.send("Python %s\nType 'quit' to exit\n>>> " % sys.version)
-
-    def on_command(self, s, cmd):
-        buffer = StringIO()
-        sys.stdout = buffer
-        sys.stderr = buffer
-        output = ''
-        try:
-            exec(cmd)
-            output = buffer.getvalue()
-        except Exception, e:
-            output = str(e)
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            buffer.close()
-        s.send(output + '\n>>> ')
-        return True
+    def on_connect(self):
+        self.send("Python %s\nType 'quit' to exit\n>>> " % sys.version)
 
 ReversePythonShell().start()

pyrasite/payloads/reverse_shell.py

@@ -13,31 +13,25 @@
 # You should have received a copy of the GNU General Public License
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright (C) 2011 Red Hat, Inc.
+# Copyright (C) 2011, 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
 
-import subprocess
+import pyrasite
 
-from _reverseconnection import ReverseConnection
+class ReverseShell(pyrasite.ReverseConnection):
 
-class ReverseShell(ReverseConnection):
+    reliable = False # This payload is designed to be used with netcat
+    host = '127.0.0.1'
+    port = 9001
 
-    host = '127.0.0.1'    # The remote host
-    port = 9001           # The same port as used by the server
+    def on_connect(self):
+        uname = pyrasite.utils.run('uname -a')[1]
+        self.send("%sType 'quit' to exit\n%% " % uname)
 
-    def on_connect(self, s):
-        uname = self._run('uname -a')[0]
-        s.send("%sType 'quit' to exit\n%% " % uname)
-
-    def on_command(self, s, cmd):
-        out, err = self._run(cmd)
+    def on_command(self, cmd):
+        p, out, err = pyrasite.utils.run(cmd)
         if err:
             out += err
-        s.send(out + '\n% ')
+        self.send(out + '\n% ')
         return True
 
-    def _run(self, cmd):
-        p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
-        out, err = p.communicate()
-        return out, err
-
 ReverseShell().start()

pyrasite/payloads/_reverseconnection.py → pyrasite/reverse.py

@@ -13,38 +13,71 @@
 # You should have received a copy of the GNU General Public License
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright (C) 2011 Red Hat, Inc.
+# Copyright (C) 2011, 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
+"""
+:mod:`pyrasite.reverse` - Pyrasite Reverse Connection Payload
+=============================================================
+"""
 
-import time, socket, threading
+import sys
+import socket
+import threading
 
-class ReverseConnection(threading.Thread):
+from StringIO import StringIO
+from pyrasite.ipc import PyrasiteIPC
 
-    host = '127.0.0.1'    # The remote host
-    port = 9001           # The same port as used by the server
+class ReverseConnection(threading.Thread, PyrasiteIPC):
+
+    host = '127.0.0.1'
+    port = 9001
+
+    def __init__(self, port=None):
+        super(ReverseConnection, self).__init__()
+        self.sock = None
+        if port:
+            self.port = port
+
+    def on_connect(self):
+        """Called when we successfuly connect to `self.host`"""
+
+    def on_command(self, cmd):
+        """Called when the host sends us a command"""
 
     def run(self):
         running = True
         while running:
             try:
-                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                s.connect((self.host, self.port))
-                self.on_connect(s)
+                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                self.sock.connect((self.host, self.port))
+                self.on_connect()
                 while running:
-                    data = s.recv(1024)
-                    if data == "quit\n" or len(data) == 0:
+                    cmd = self.recv()
+                    if cmd is None or cmd == "quit\n" or len(cmd) == 0:
                         running = False
                     else:
-                        try:
-                            running = self.on_command(s, data)
-                        except:
-                            running = False
-                s.close()
-            except socket.error, e:
+                        running = self.on_command(cmd)
+            except Exception, e:
                 print(str(e))
-            time.sleep(5)
+                running = False
+            if not running:
+                self.close()
+
 
-    def on_connect(self, s):
-        pass
+class ReversePythonConnection(ReverseConnection):
 
-    def on_command(self, s, cmd):
-        raise NotImplementedError("You must prove your own on_command method")
+    def on_command(self, cmd):
+        buffer = StringIO()
+        sys.stdout = buffer
+        sys.stderr = buffer
+        output = ''
+        try:
+            exec(cmd)
+            output = buffer.getvalue()
+        except Exception, e:
+            output = str(e)
+        finally:
+            sys.stdout = sys.__stdout__
+            sys.stderr = sys.__stderr__
+            buffer.close()
+        self.send(output)
+        return True

pyrasite/tools/gui.py

@@ -15,6 +15,10 @@
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
 # Copyright (C) 2012 Red Hat, Inc., Luke Macken <lmacken@redhat.com>
+#
+# This interface may contain some code from the gtk-demo, written
+# by John (J5) Palmieri, and licensed under the LGPLv2.1
+# http://git.gnome.org/browse/pygobject/tree/demos/gtk-demo/gtk-demo.py
 
 import os
 import sys
@@ -25,51 +29,28 @@ import keyword
 import tokenize
 
 from meliae import loader
-from gi.repository import GLib, GObject, Pango, GdkPixbuf, Gtk
+from gi.repository import GLib, GObject, Pango, Gtk
 
 import pyrasite
-from pyrasite.utils import run, setup_logger
-from pyrasite.ipc import PyrasiteIPC
+from pyrasite.utils import setup_logger, run
 
 log = logging.getLogger('pyrasite')
 
 
-class Process(GObject.GObject):
-
-    def __init__(self, pid):
-        super(Process, self).__init__()
-        self.pid = pid
-        self.title = run('ps --no-heading -o cmd= -p %d' % pid)[1]
-        self.command = self.title
-        self.title = self.title[:25]
-        self.ipc = None
-        self.filename = None
-
-    def connect(self):
-        """
-        Setup a communication socket with the process by injecting
-        a reverse subshell and having it connect back to us.
-        """
-        self.ipc = PyrasiteIPC(self.pid)
-        self.ipc.inject()
-        self.ipc.listen()
-
-    def cmd(self, cmd, *args, **kw):
-        return self.ipc.cmd(cmd, *args, **kw)
-
-    def close(self):
-        log.debug("Closing %r" % self)
-        self.ipc.close()
-
-    def __repr__(self):
-        return "<Process %d '%s'>" % (self.pid, self.title)
+class Process(pyrasite.PyrasiteIPC, GObject.GObject):
+    """
+    A :class:`GObject.GObject` subclass for use in the :class:`ProcessTreeStore`
+    """
 
-    def __str__(self):
-        return self.title
+    @property
+    def title(self):
+        if not getattr(self, '_title', None):
+            self._title = run('ps --no-heading -o cmd= -p %d' % self.pid)[1]
+        return self._title
 
 
 class ProcessTreeStore(Gtk.TreeStore):
-    """ This TreeStore finds all running python processes.  """
+    """This TreeStore finds all running python processes."""
 
     def __init__(self, *args):
         Gtk.TreeStore.__init__(self, str, Process, Pango.Style)
@@ -80,35 +61,14 @@ class ProcessTreeStore(Gtk.TreeStore):
                 try:
                     maps = open('/proc/%d/maps' % pid).read().strip()
                     if 'python' in maps:
-                        self.append(None, (proc.title, proc, Pango.Style.NORMAL))
+                        self.append(None, (proc.title.strip(), proc,
+                                           Pango.Style.NORMAL))
                 except IOError:
                     pass
             except ValueError:
                 pass
 
 
-class InputStream(object):
-    '''
-    Simple Wrapper for File-like objects. [c]StringIO doesn't provide
-    a readline function for use with generate_tokens.
-    Using a iterator-like interface doesn't succeed, because the readline
-    function isn't used in such a context. (see <python-lib>/tokenize.py)
-    '''
-    def __init__(self, data):
-        self.__data = [ '%s\n' % x for x in data.splitlines() ]
-        self.__lcount = 0
-
-    def readline(self):
-        try:
-            line = self.__data[self.__lcount]
-            self.__lcount += 1
-        except IndexError:
-            line = ''
-            self.__lcount = 0
-
-        return line
-
-
 class PyrasiteWindow(Gtk.Window):
 
     def __init__(self):
@@ -240,7 +200,8 @@ class PyrasiteWindow(Gtk.Window):
         self.shell_prompt.set_activates_default(True)
         self.shell_button.set_receives_default(True)
 
-        notebook.append_page(shell_hbox, Gtk.Label.new_with_mnemonic('_Shell'))
+        shell_label = Gtk.Label.new_with_mnemonic('_Shell')
+        notebook.append_page(shell_hbox, shell_label)
 
         # To try and grab focus of our text input
         notebook.connect('switch-page', self.switch_page)
@@ -400,8 +361,8 @@ class PyrasiteWindow(Gtk.Window):
 
         self.fontify()
         self.update_progress(1.0)
-        self.update_progress(0.0)
         self.progress.hide()
+        self.update_progress(0.0)
 
     def dump_objects(self, proc):
         cmd = ';'.join(["import os, shutil", "from meliae import scanner",
@@ -415,12 +376,6 @@ class PyrasiteWindow(Gtk.Window):
         self.obj_store.clear()
         self.update_progress(0.4, "Loading object dump")
 
-        obj_dump = '/tmp/%d.objects' % proc.pid
-        if not os.path.exists(obj_dump):
-            time.sleep(1)
-            if not os.path.exists(obj_dump):
-                time.sleep(2)
-
         objects = loader.load('/tmp/%d.objects' % proc.pid)
         objects.compute_referrers()
         self.update_progress(0.45)
@@ -440,10 +395,11 @@ class PyrasiteWindow(Gtk.Window):
                 obj = summary.summaries[i - 2]
                 self.obj_store.append([str(obj.max_address)] +
                                        map(intify, line.split()[1:]))
+
     def dump_stacks(self, proc):
         self.update_progress(0.55, "Dumping stacks")
         payloads = os.path.join(os.path.abspath(os.path.dirname(
-            pyrasite.__file__)), '..', 'payloads')
+            pyrasite.__file__)), 'payloads')
         dump_stacks = os.path.join(payloads, 'dump_stacks.py')
         code = proc.cmd(file(dump_stacks).read())
         self.update_progress(0.6)
@@ -456,7 +412,7 @@ class PyrasiteWindow(Gtk.Window):
         self.update_progress(0.7, "Tracing call stack")
         proc.cmd('import pycallgraph; pycallgraph.start_trace()')
         self.update_progress(0.8)
-        time.sleep(1)
+        time.sleep(1) # TODO: make this configurable in the UI
         self.update_progress(0.9, "Generating call stack graph")
         image = '/tmp/%d-callgraph.png' % proc.pid
         proc.cmd('import pycallgraph; pycallgraph.make_dot_graph("%s")' % image)
@@ -617,6 +573,28 @@ class PyrasiteWindow(Gtk.Window):
             process.close()
 
 
+class InputStream(object):
+    '''
+    Simple Wrapper for File-like objects. [c]StringIO doesn't provide
+    a readline function for use with generate_tokens.
+    Using a iterator-like interface doesn't succeed, because the readline
+    function isn't used in such a context. (see <python-lib>/tokenize.py)
+    '''
+    def __init__(self, data):
+        self.__data = [ '%s\n' % x for x in data.splitlines() ]
+        self.__lcount = 0
+
+    def readline(self):
+        try:
+            line = self.__data[self.__lcount]
+            self.__lcount += 1
+        except IndexError:
+            line = ''
+            self.__lcount = 0
+
+        return line
+
+
 def main():
     mainloop = GLib.MainLoop()
 

tests/test_code_injection.py

@@ -13,7 +13,7 @@
 # You should have received a copy of the GNU General Public License
 # along with pyrasite.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright (C) 2011 Red Hat, Inc.
+# Copyright (C) 2011, 2012 Red Hat, Inc.
 
 import unittest
 
@@ -27,7 +27,7 @@ class TestCodeInjection(unittest.TestCase):
         p = run(cmd, communicate=False)[0]
 
         ci = CodeInjector(p.pid, verbose=True)
-        ci.inject('payloads/helloworld.py')
+        ci.inject('pyrasite/payloads/helloworld.py')
 
         stdout, stderr = p.communicate()
         assert 'Hello World!' in stdout, "Code injection failed"
@@ -41,7 +41,7 @@ class TestCodeInjection(unittest.TestCase):
         p = run('python -c "%s"' % ';'.join(cmd), communicate=False)[0]
 
         ci = CodeInjector(p.pid, verbose=True)
-        ci.inject('payloads/helloworld.py')
+        ci.inject('pyrasite/payloads/helloworld.py')
 
         stdout, stderr = p.communicate()
         assert 'Hello World!' in stdout, "Multi-threaded code injection failed"