Merge pull request #10 from gregjurman/feature/localhost-fix

IPv6 fix for Issue #8

docs/Payloads.rst

@@ -18,7 +18,7 @@ This lets you easily introspect or alter any objects in your running process.
 ::
 
     $ pyrasite <PID> pyrasite/payloads/reverse_python_shell.py
-    $ nc -l localhost 9001
+    $ nc -l 9001
     Python 2.7.1 (r271:86832, Apr 12 2011, 16:15:16)
     [GCC 4.6.0 20110331 (Red Hat 4.6.0-2)]
     Type 'quit' to exit.
@@ -62,7 +62,7 @@ Reverse Shell
 ::
 
     $ pyrasite <PID> pyrasite/payloads/reverse_shell.py
-    $ nc -l localhost 9001
+    $ nc -l 9001
     Linux tomservo 2.6.40.3-0.fc15.x86_64 #1 SMP Tue Aug 16 04:10:59 UTC 2011 x86_64 x86_64 x86_64 GNU/Linux
     Type 'quit' to exit.
     % ls

pyrasite/ipc.py

@@ -74,11 +74,23 @@ class PyrasiteIPC(object):
 
     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))
+        for res in socket.getaddrinfo('localhost', None, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, 0):
+            af, socktype, proto, canonname, sa = res
+            try:
+                self.server_sock = socket.socket(af, socktype, proto)
+            except socket.error, msg:
+                self.server_sock = None
+                continue
+            try:
+                self.server_sock.bind(sa)
                 self.server_sock.listen(1)
-        self.port = self.server_sock.getsockname()[1]
+            except socket.error, msg:
+                self.server_sock.close()
+                self.server_sock = None
+                continue
+            break
+
+        self.hostname, self.port = self.server_sock.getsockname()[0:2]
         self.running = True
 
     def create_payload(self):

pyrasite/payloads/reverse_shell.py

@@ -20,7 +20,6 @@ import pyrasite
 class ReverseShell(pyrasite.ReverseConnection):
 
     reliable = False # This payload is designed to be used with netcat
-    host = '127.0.0.1'
     port = 9001
 
     def on_connect(self):

pyrasite/reverse.py

@@ -29,7 +29,7 @@ from pyrasite.ipc import PyrasiteIPC
 class ReverseConnection(threading.Thread, PyrasiteIPC):
     """A payload that connects to a given host:port and receives commands"""
 
-    host = '127.0.0.1'
+    host = 'localhost'
     port = 9001
 
     def __init__(self, host=None, port=None):
@@ -50,8 +50,21 @@ class ReverseConnection(threading.Thread, PyrasiteIPC):
         running = True
         while running:
             try:
-                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                self.sock.connect((self.host, self.port))
+                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)
+                    except socket.error, msg:
+                        self.sock = None
+                        continue
+                    try:
+                        self.sock.connect(sa)
+                    except socket.error, msg:
+                        self.sock.close()
+                        self.sock = None
+                        continue
+                    break
+
                 self.on_connect()
                 while running:
                     cmd = self.recv()