diff options
| -rw-r--r-- | lib/chutney/Traffic.py | 49 |
1 files changed, 47 insertions, 2 deletions
diff --git a/lib/chutney/Traffic.py b/lib/chutney/Traffic.py index 9b0a1f7..ab9a578 100644 --- a/lib/chutney/Traffic.py +++ b/lib/chutney/Traffic.py @@ -6,19 +6,38 @@ # restrict, so long as you retain the above notice(s) and this license # in all redistributed copies and derived works. There is no warranty. +# Do select/read/write for binding to a port, connecting to it and +# write, read what's written and verify it. You can connect over a +# SOCKS proxy (like Tor). +# +# You can create a TrafficTester and give it an IP address/host and +# port to bind to. If a Source is created and added to the +# TrafficTester, it will connect to the address/port it was given at +# instantiation and send its data. A Source can be configured to +# connect over a SOCKS proxy. When everything is set up, you can +# invoke TrafficTester.run() to start running. The TrafficTester will +# accept the incoming connection and read from it, verifying the data. +# +# For example code, see main() below. + import socket import select import struct import errno +# Set debug_flag=True in order to debug this program or to get hints +# about what's going wrong in your system. debug_flag = False def debug(s): + "Print a debug message on stdout if debug_flag is True." if debug_flag: print("DEBUG: %s" % s) def socks_cmd(addr_port): """ + Return a SOCKS command for connecting to addr_port. + SOCKSv4: https://en.wikipedia.org/wiki/SOCKS#Protocol SOCKSv5: RFC1928, RFC1929 """ @@ -63,6 +82,7 @@ class TestSuite(object): return('%d/%d/%d' % (self.not_done, self.successes, self.failures)) class Peer(object): + "Base class for Listener, Source and Sink." LISTENER = 1 SOURCE = 2 SINK = 3 @@ -84,6 +104,7 @@ class Peer(object): return self.type == self.SINK class Listener(Peer): + "A TCP listener, binding, listening and accepting new connections." def __init__(self, tt, endpoint): super(Listener, self).__init__(Peer.LISTENER, tt) self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) @@ -97,11 +118,17 @@ class Listener(Peer): self.tt.add(Sink(self.tt, newsock)) class Sink(Peer): + "A data sink, reading from its peer and verifying the data." def __init__(self, tt, s): super(Sink, self).__init__(Peer.SINK, tt, s) self.inbuf = '' def on_readable(self): + """Invoked when the socket becomes readable. + Return 0 on finished, successful verification. + -1 on failed verification + >0 if more data needs to be read + """ return self.verify(self.tt.data) def verify(self, data): @@ -115,6 +142,8 @@ class Sink(Peer): return len(data) - len(self.inbuf) class Source(Peer): + """A data source, connecting to a TCP server, optionally over a + SOCKS proxy, sending data.""" NOT_CONNECTED = 0 CONNECTING = 1 CONNECTING_THROUGH_PROXY = 2 @@ -143,6 +172,10 @@ class Source(Peer): raise def on_readable(self): + """Invoked when the socket becomes readable. + Return -1 on failure + >0 if more data needs to be read or written + """ if self.state == self.CONNECTING_THROUGH_PROXY: self.inbuf += self.s.recv(8 - len(self.inbuf)) if len(self.inbuf) == 8: @@ -168,6 +201,11 @@ class Source(Peer): return False def on_writable(self): + """Invoked when the socket becomes writable. + Return 0 when done writing + -1 on failure (like connection refused) + >0 if more data needs to be written + """ if self.state == self.CONNECTING: if self.proxy is None: self.state = self.CONNECTED @@ -189,6 +227,13 @@ class Source(Peer): return len(self.outbuf) # When 0, we're being removed. class TrafficTester(): + """ + Hang on select.select() and dispatch to Sources and Sinks. + Time out after self.timeout seconds. + Keep track of successful and failed data verification using a + TestSuite. + Return True if all tests succeed, else False. + """ def __init__(self, endpoint, data={}, timeout=3): self.listener = Listener(self, endpoint) self.pending_close = [] @@ -246,8 +291,8 @@ class TrafficTester(): self.remove(p) for fd in sets[1]: # writable fd's - p = self.peers.get(fd) # Might have been removed above. - if p is not None: + p = self.peers.get(fd) + if p is not None: # Might have been removed above. n = p.on_writable() if n == 0: self.remove(p) |