proxy: add drop next packet option

1 files changed, 34 insertions, 15 deletions

diff --git a/proxy.py b/proxy.py
index a8478d6..03af6d6 100755
--- a/proxy.py
+++ b/proxy.py
@@ -52,7 +52,21 @@ def connect(addr):
     s.setblocking(False)
     return s
 
-def proxy(local_port, remote_addr):
+class Proxy(object):
+    def __init__(self, local_port, remote_addr):
+        self._drop_next = False
+        self._proxy = _proxy(self, local_port, remote_addr)
+    def drop_next(self):
+        self._drop_next = True
+    def check_drop_next(self):
+        dn = self._drop_next
+        self._drop_next = False
+        return dn
+    def __iter__(self):
+        for x in self._proxy:
+            yield x
+
+def _proxy(proxy, local_port, remote_addr):
     print "proxying from %s to %s" % (local_port, remote_addr)
     accepter = make_accepter(local_port)
     open_socks = twowaydict()
@@ -67,14 +81,15 @@ def proxy(local_port, remote_addr):
                 open_socks[s] = connect(remote_addr)
             else:
                 other = open_socks[s]
-                src_dst = [s, other]
+                src_dst_socks = [s, other]
+                src_dst = [None, None]
                 dont_recv = False
                 for i in xrange(len(src_dst)):
                     try:
-                        src_dst[i] = src_dst[i].getpeername()[1]
+                        src_dst[i] = src_dst_socks[i].getpeername()[1]
                     except Exception, e:
                         if e.errno in close_errnos:
-                            src_dst[1-i].close()
+                            src_dst_socks[1-i].close()
                             if s in open_socks:
                                 del open_socks[s]
                             dont_recv = True
@@ -86,17 +101,21 @@ def proxy(local_port, remote_addr):
                     del open_socks[s]
                     continue
                 yield src, dst, data
-                try:
-                    other.send(data)
-                except Exception, e:
-                    if e.errno in close_errnos:
-                        n = len(open_socks)
-                        s.close()
-                        open_socks[s].close()
-                        del open_socks[s]
-                        assert(len(open_socks) == n - 1)
-                    else:
-                        import pdb; pdb.set_trace()
+                if not proxy.check_drop_next():
+                    try:
+                        other.send(data)
+                    except Exception, e:
+                        if e.errno in close_errnos:
+                            n = len(open_socks)
+                            s.close()
+                            open_socks[s].close()
+                            del open_socks[s]
+                            assert(len(open_socks) == n - 1)
+                        else:
+                            import pdb; pdb.set_trace()
+
+def proxy(local_port, remote_addr):
+    return Proxy(local_port, remote_addr)
 
 def closeallsockets():
     for s in Socket.sockets: