fix

1 files changed, 23 insertions, 8 deletions

diff --git a/proxy.py b/proxy.py
index 8c43a25..1d7327c 100755
--- a/proxy.py
+++ b/proxy.py
@@ -1,10 +1,11 @@
 #!/usr/bin/python
-from socket import socket, AF_INET, SOCK_STREAM
+import socket
+from socket import AF_INET, SOCK_STREAM
 from select import select
 
 MAX_PACKET_SIZE=65536
 
-class Socket(socket):
+class Socket(socket.socket):
     sockets = []
     def __init__(self, *args, **kw):
         super(Socket, self).__init__(*args, **kw)
@@ -48,7 +49,10 @@ def make_accepter(port, host='127.0.0.1'):
 
 def connect(addr):
     s = Socket(AF_INET, SOCK_STREAM)
-    s.connect(addr)
+    try:
+        s.connect(addr)
+    except socket.error:
+        return None
     s.setblocking(False)
     return s
 
@@ -58,7 +62,7 @@ class Proxy(object):
     def __init__(self, local_port, remote_addr, host='127.0.0.1'):
         self._drop_next = False
         handle_input, select_based_iterator, get_fds = make_proxy(
-            self, local_port, remote_addr, host)
+            local_port, remote_addr, host, check_drop_next=self.check_drop_next)
         self.select_based_iterator = select_based_iterator()
 
     def drop_next(self):
@@ -75,13 +79,16 @@ class Proxy(object):
 
 
 def make_proxy(local_port, remote_addr, host = '127.0.0.1',
-               check_drop_next=lambda: True):
+               check_drop_next=lambda: False, debug=False):
     print "proxying from %s to %s" % (local_port, remote_addr)
     accepter = make_accepter(local_port, host)
     open_socks = twowaydict()
     close_errnos = set([BROKEN_PIPE_ERRNO, TRANSPORT_NOT_CONNECTED_ERRNO])
     def get_fds():
-        return [accepter] + open_socks.allkeys()
+        fds = [accepter] + open_socks.allkeys()
+        if debug:
+            print repr(fds)
+        return fds
     def iterate_packets():
         """
         An iterator with inputs, designed to work with some external
@@ -109,7 +116,11 @@ def make_proxy(local_port, remote_addr, host = '127.0.0.1',
         """
         if s is accepter:
             s, _addr = accepter.accept()
-            open_socks[s] = connect(remote_addr)
+            s_or_none = connect(remote_addr)
+            if s_or_none:
+                open_socks[s] = s_or_none
+            else:
+                print "connection to remote %s failed" % repr(remote_addr)
         else:
             other = open_socks[s]
             src_dst_socks = [s, other]
@@ -180,10 +191,14 @@ def tests():
     import sys
     p = argparse.ArgumentParser(description="proxy multiple socket connections")
     p.add_argument('--listen-port', default=11000)
+    p.add_argument('--listen-host', default='127.0.0.1')
     p.add_argument('--remote-port', default=12000)
     p.add_argument('--remote-host', default='127.0.0.1')
     args = p.parse_args(sys.argv[1:])
-    for x in proxy(args.listen_port, (args.remote_host, args.remote_port)):
+    handle_input, select_based_iterator, get_fds = make_proxy(
+        args.listen_port, (args.remote_host, args.remote_port),
+        args.listen_host, debug=True)
+    for x in select_based_iterator():
         print x
 
 if __name__ == '__main__':