]> code.delx.au - monosys/blobdiff - scripts/proxy.py
Renamed underscores to hyphens
[monosys] / scripts / proxy.py
index 0690611a11f0ecc7bef247ab3c7442766116ac97..66220ef12ae243bea888a8915454ed91fbd8e1f5 100755 (executable)
@@ -84,6 +84,7 @@ import os
 import socket
 import struct
 import sys
+import traceback
 
 
 if sys.platform == "linux2":
@@ -122,8 +123,17 @@ elif sys.platform == "darwin":
 
 
 class Proxy(asyncore.dispatcher):
-       def __init__(self, sock):
-               asyncore.dispatcher.__init__(self, sock)
+       def __init__(self, arg):
+               if isinstance(arg, tuple):
+                       asyncore.dispatcher.__init__(self)
+                       self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+                       self.connect(arg)
+               else:
+                       asyncore.dispatcher.__init__(self, arg)
+               self.init()
+
+       def init(self):
+               self.end = False
                self.other = None
                self.buffer = ""
 
@@ -132,8 +142,8 @@ class Proxy(asyncore.dispatcher):
                other.other = self
 
        def handle_error(self):
-               print >>sys.stderr, "Proxy error:", sys.exc_info()
-               self.handle_close()
+               print >>sys.stderr, "Proxy error:", traceback.format_exc()
+               self.close()
 
        def handle_read(self):
                data = self.recv(8192)
@@ -143,6 +153,8 @@ class Proxy(asyncore.dispatcher):
        def handle_write(self):
                sent = self.send(self.buffer)
                self.buffer = self.buffer[sent:]
+               if len(self.buffer) == 0 and self.end:
+                       self.close()
 
        def writable(self):
                return len(self.buffer) > 0
@@ -151,8 +163,10 @@ class Proxy(asyncore.dispatcher):
                if not self.other:
                        return
                print >>sys.stderr, "Proxy closed"
-               self.other.close()
                self.close()
+               if len(self.other.buffer) == 0:
+                       self.other.close()
+               self.other.end = True
                self.other = None
 
 class ConnectProxy(asyncore.dispatcher):
@@ -161,8 +175,8 @@ class ConnectProxy(asyncore.dispatcher):
                self.buffer = ""
 
        def handle_error(self):
-               print >>sys.stderr, "ConnectProxy error:", sys.exc_info()
-               self.handle_close()
+               print >>sys.stderr, "ConnectProxy error:", traceback.format_exc()
+               self.close()
 
        def handle_read(self):
                self.buffer += self.recv(8192)
@@ -191,15 +205,13 @@ class ConnectProxy(asyncore.dispatcher):
                print >>sys.stderr, "Forwarding connection", host, port
 
                # Create server proxy
-               server_connection = socket.socket()
-               server_connection.connect((host, port))
-               server = Proxy(server_connection)
+               server = Proxy((host, port))
+               server.buffer = self.buffer
 
                # Morph and connect
                self.__class__ = Proxy
+               self.init()
                server.meet(self)
-               server.buffer = self.buffer
-               self.buffer = ""
 
 
 class BasicForwarder(asyncore.dispatcher):
@@ -212,9 +224,10 @@ class BasicForwarder(asyncore.dispatcher):
                self.set_reuse_addr()
                self.bind(("", listen_port))
                self.listen(50)
+               print >>sys.stderr, "BasicForwarder bound on", listen_port
 
        def handle_error(self):
-               print >>sys.stderr, "BasicForwarder error:", sys.exc_info()
+               print >>sys.stderr, "BasicForwarder error:", traceback.format_exc()
 
        def handle_accept(self):
                client_connection, source_addr = self.accept()
@@ -224,12 +237,10 @@ class BasicForwarder(asyncore.dispatcher):
                        return
 
                print >>sys.stderr, "Accepted connection from", source_addr
-               server_connection = socket.socket()
-               server_connection.connect((self.host, self.port))
 
                # Hook the sockets up to the event loop
                client = Proxy(client_connection)
-               server = Proxy(server_connection)
+               server = Proxy((self.host, self.port))
                server.meet(client)
 
 class Forwarder(asyncore.dispatcher):
@@ -240,9 +251,10 @@ class Forwarder(asyncore.dispatcher):
                self.set_reuse_addr()
                self.bind(("", listen_port))
                self.listen(50)
+               print >>sys.stderr, "Forwarder bound on", listen_port
 
        def handle_error(self):
-               print >>sys.stderr, "Forwarder error:", sys.exc_info()
+               print >>sys.stderr, "Forwarder error:", traceback.format_exc()
 
        def handle_accept(self):
                client_connection, source_addr = self.accept()
@@ -263,21 +275,20 @@ class Interceptor(asyncore.dispatcher):
                self.set_reuse_addr()
                self.bind(("0.0.0.0", listen_port))
                self.listen(50)
+               print >>sys.stderr, "Interceptor bound on", listen_port
 
        def handle_error(self):
-               print >>sys.stderr, "Interceptor error!", sys.exc_info()
+               print >>sys.stderr, "Interceptor error!", traceback.format_exc()
 
        def handle_accept(self):
                # Get sockets
                client_connection, source_addr = self.accept()
                dest = get_original_dest(client_connection)
                print >>sys.stderr, "Accepted connection from", source_addr
-               server_connection = socket.socket()
-               server_connection.connect((self.host, self.port))
 
                # Hook them up to the event loop
                client = Proxy(client_connection)
-               server = Proxy(server_connection)
+               server = Proxy((self.host, self.port))
                server.buffer += "%s\n%d\n" % dest
                server.meet(client)