Improved proxy.py

diff --git a/scripts/proxy.py b/scripts/proxy.py
index ec32f415d898b3f645e751ecebcfe017008a8282..48ab1d85bef0f4314bfa3bf33f78f9a2198f72fc 100755 (executable)
--- a/scripts/proxy.py
+++ b/scripts/proxy.py
@@ -1,6 +1,19 @@
 #!/usr/bin/env python
 
+# Sample config file:
+
+# [proxy]
+# listen_port = 8000
+# host = google.com
+# port = 80
+# 
+# [allowed]
+# host1 = delx.cjb.net
+# host2 = 333.333.333.333
+
+
 import asyncore
+import ConfigParser
 import os
 import socket
 import sys
@@ -9,7 +22,7 @@ class Proxy(asyncore.dispatcher):
        def __init__(self, sock):
                asyncore.dispatcher.__init__(self, sock)
                self.other = None
-               self.towrite = ""
+               self.buffer = ""
        
        def meet(self, other):
                self.other = other
@@ -21,11 +34,11 @@ class Proxy(asyncore.dispatcher):
 
        def handle_read(self):
                data = self.recv(8192)
-               self.other.towrite += data
+               self.other.buffer += data
        
        def handle_write(self):
-               sent = self.send(self.towrite)
-               self.towrite = self.towrite[sent:]
+               sent = self.send(self.buffer)
+               self.buffer = self.buffer[sent:]
        
        def handle_close(self):
                if not self.other:
@@ -35,63 +48,79 @@ class Proxy(asyncore.dispatcher):
                self.other = None
 
 class Forwarder(asyncore.dispatcher):
-       def __init__(self, host, port):
+       def __init__(self, listen_port, host, port, allowed):
                asyncore.dispatcher.__init__(self)
                self.host = host
                self.port = port
+               self.allowed = allowed
                self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
-               self.bind(("", port))
+               self.bind(("", listen_port))
                self.listen(5)
        
        def handle_error(self):
                print >>sys.stderr, "Connection error!"
        
        def handle_accept(self):
-               # Get sockets
-               clientConnection, addr = self.accept()
+               client_connection, (addr, port) = self.accept()
+               if addr not in map(socket.gethostbyname, self.allowed):
+                       print >>sys.stderr, "Rejected connection from", addr
+                       client_connection.close()
+                       return
+
                print >>sys.stderr, "Accepted connection from", addr
-               serverConnection = socket.socket()
-               serverConnection.connect((self.host, self.port))
+               server_connection = socket.socket()
+               server_connection.connect((self.host, self.port))
 
-               # Hook them up to the event loop
-               client = Proxy(clientConnection)
-               server = Proxy(serverConnection)
+               # Hook the sockets up to the event loop
+               client = Proxy(client_connection)
+               server = Proxy(server_connection)
                server.meet(client)
 
-def main(host, port):
-       proxy = Forwarder(host, port)
+def main(listen_port, host, port, allowed):
+       proxy = Forwarder(listen_port, host, port, allowed)
        asyncore.loop()
 
 if __name__ == "__main__":
        try:
                if sys.argv[1] == "-d":
                        daemon = True
-                       host, port = sys.argv[2:]
+                       config = sys.argv[2]
                else:
-                       host, port = sys.argv[1:]
                        daemon = False
-               port = int(port)
+                       config = sys.argv[1]
        except ValueError:
-               print >> sys.stderr, "Usage: %s [-d] host port" % sys.argv[0]
+               print >>sys.stderr, "Usage: %s [-d] config" % sys.argv[0]
+               sys.exit(1)
+
+       try:
+               c = ConfigParser.RawConfigParser()
+               c.read(config)
+               listen_port = c.getint("proxy", "listen_port")
+               host = c.get("proxy", "host")
+               port = c.getint("proxy", "port")
+               allowed = c.items("allowed")
+               allowed = [host for _,host in c.items("allowed")]
+       except:
+               print >>sys.stderr, "Error parsing config!"
                sys.exit(1)
 
        if not daemon:
                try:
-                       main(host, port)
+                       main(listen_port, host, port, allowed)
                except KeyboardInterrupt:
                        print
        else:
-               sys.stdin.close()
-               sys.stdout.close()
-               sys.stderr.close()
-               sys.stdin = open("/dev/null")
-               sys.stdout = open("/dev/null")
-               sys.stderr = open("/dev/null")
+               os.close(0)
+               os.close(1)
+               os.close(2)
+               os.open("/dev/null", os.O_RDONLY)
+               os.open("/dev/null", os.O_RDWR)
+               os.dup(1)
 
                if os.fork() == 0:
                        # We are the child
                        try:
-                               main(host, port)
+                               main(listen_port, host, port, allowed)
                        except KeyboardInterrupt:
                                print
                        sys.exit(0)