Traffic.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366
  1. #!/usr/bin/env python
  2. #
  3. # Copyright 2013 The Tor Project
  4. #
  5. # You may do anything with this work that copyright law would normally
  6. # restrict, so long as you retain the above notice(s) and this license
  7. # in all redistributed copies and derived works. There is no warranty.
  8. # Do select/read/write for binding to a port, connecting to it and
  9. # write, read what's written and verify it. You can connect over a
  10. # SOCKS proxy (like Tor).
  11. #
  12. # You can create a TrafficTester and give it an IP address/host and
  13. # port to bind to. If a Source is created and added to the
  14. # TrafficTester, it will connect to the address/port it was given at
  15. # instantiation and send its data. A Source can be configured to
  16. # connect over a SOCKS proxy. When everything is set up, you can
  17. # invoke TrafficTester.run() to start running. The TrafficTester will
  18. # accept the incoming connection and read from it, verifying the data.
  19. #
  20. # For example code, see main() below.
  21. from __future__ import print_function
  22. import sys
  23. import socket
  24. import struct
  25. import errno
  26. import time
  27. import os
  28. import asyncore
  29. import asynchat
  30. from chutney.Debug import debug_flag, debug
  31. def addr_to_family(addr):
  32. for family in [socket.AF_INET, socket.AF_INET6]:
  33. try:
  34. socket.inet_pton(family, addr)
  35. return family
  36. except (socket.error, OSError):
  37. pass
  38. return socket.AF_INET
  39. def socks_cmd(addr_port):
  40. """
  41. Return a SOCKS command for connecting to addr_port.
  42. SOCKSv4: https://en.wikipedia.org/wiki/SOCKS#Protocol
  43. SOCKSv5: RFC1928, RFC1929
  44. """
  45. ver = 4 # Only SOCKSv4 for now.
  46. cmd = 1 # Stream connection.
  47. user = b'\x00'
  48. dnsname = ''
  49. host, port = addr_port
  50. try:
  51. addr = socket.inet_aton(host)
  52. except socket.error:
  53. addr = b'\x00\x00\x00\x01'
  54. dnsname = '%s\x00' % host
  55. debug("Socks 4a request to %s:%d" % (host, port))
  56. if type(dnsname) != type(b""):
  57. dnsname = dnsname.encode("ascii")
  58. return struct.pack('!BBH', ver, cmd, port) + addr + user + dnsname
  59. class TestSuite(object):
  60. """Keep a tab on how many tests are pending, how many have failed
  61. and how many have succeeded."""
  62. def __init__(self):
  63. self.tests = {}
  64. self.not_done = 0
  65. self.successes = 0
  66. self.failures = 0
  67. def add(self, name):
  68. if name not in self.tests:
  69. debug("Registering %s"%name)
  70. self.not_done += 1
  71. self.tests[name] = 'not done'
  72. def success(self, name):
  73. if self.tests[name] == 'not done':
  74. debug("Succeeded %s"%name)
  75. self.tests[name] = 'success'
  76. self.not_done -= 1
  77. self.successes += 1
  78. def failure(self, name):
  79. if self.tests[name] == 'not done':
  80. debug("Failed %s"%name)
  81. self.tests[name] = 'failure'
  82. self.not_done -= 1
  83. self.failures += 1
  84. def failure_count(self):
  85. return self.failures
  86. def all_done(self):
  87. return self.not_done == 0
  88. def status(self):
  89. return('%s: %d/%d/%d' % (self.tests, self.not_done, self.successes,
  90. self.failures))
  91. class Listener(asyncore.dispatcher):
  92. "A TCP listener, binding, listening and accepting new connections."
  93. def __init__(self, tt, endpoint):
  94. asyncore.dispatcher.__init__(self)
  95. self.create_socket(addr_to_family(endpoint[0]), socket.SOCK_STREAM)
  96. self.set_reuse_addr()
  97. self.bind(endpoint)
  98. self.listen(0)
  99. self.tt = tt
  100. def handle_accept(self):
  101. # deprecated in python 3.2
  102. pair = self.accept()
  103. if pair is not None:
  104. newsock, endpoint = pair
  105. debug("new client from %s:%s (fd=%d)" %
  106. (endpoint[0], endpoint[1], newsock.fileno()))
  107. handler = Sink(newsock, self.tt)
  108. self.tt.add(handler)
  109. def fileno(self):
  110. return self.socket.fileno()
  111. class DataSource(object):
  112. """A data source generates some number of bytes of data, and then
  113. returns None.
  114. For convenience, it conforms to the 'producer' api.
  115. """
  116. def __init__(self, data, repetitions=1):
  117. self.data = data
  118. self.repetitions = repetitions
  119. self.sent_any = False
  120. def copy(self):
  121. assert not self.sent_any
  122. return DataSource(self.data, self.repetitions)
  123. def more(self):
  124. self.sent_any = True
  125. if self.repetitions > 0:
  126. self.repetitions -= 1
  127. return self.data
  128. return None
  129. class DataChecker(object):
  130. """A data checker verifies its input against bytes in a stream."""
  131. def __init__(self, source):
  132. self.source = source
  133. self.pending = b""
  134. self.succeeded = False
  135. self.failed = False
  136. def consume(self, inp):
  137. if self.failed:
  138. return
  139. if self.succeeded and len(inp):
  140. self.succeeded = False
  141. self.failed = True
  142. return
  143. while len(inp):
  144. n = min(len(inp), len(self.pending))
  145. if inp[:n] != self.pending[:n]:
  146. self.failed = True
  147. return
  148. inp = inp[n:]
  149. self.pending = self.pending[n:]
  150. if not self.pending:
  151. self.pending = self.source.more()
  152. if self.pending is None:
  153. if len(inp):
  154. self.failed = True
  155. else:
  156. self.succeeded = True
  157. return
  158. class Sink(asynchat.async_chat):
  159. "A data sink, reading from its peer and verifying the data."
  160. def __init__(self, sock, tt):
  161. asynchat.async_chat.__init__(self, sock)
  162. self.set_terminator(None)
  163. self.tt = tt
  164. self.data_checker = DataChecker(tt.data_source.copy())
  165. self.testname = "recv-data%s"%id(self)
  166. def get_test_names(self):
  167. return [ self.testname ]
  168. def collect_incoming_data(self, inp):
  169. # shortcut read when we don't ever expect any data
  170. debug("successfully received (bytes=%d)" % len(inp))
  171. self.data_checker.consume(inp)
  172. if self.data_checker.succeeded:
  173. debug("successful verification")
  174. self.close()
  175. self.tt.success(self.testname)
  176. elif self.data_checker.failed:
  177. debug("receive comparison failed")
  178. self.tt.failure(self.testname)
  179. self.close()
  180. def fileno(self):
  181. return self.socket.fileno()
  182. class CloseSourceProducer:
  183. """Helper: when this producer is returned, a source is successful."""
  184. def __init__(self, source):
  185. self.source = source
  186. def more(self):
  187. self.source.sent_ok()
  188. return b""
  189. class Source(asynchat.async_chat):
  190. """A data source, connecting to a TCP server, optionally over a
  191. SOCKS proxy, sending data."""
  192. CONNECTING = 1
  193. CONNECTING_THROUGH_PROXY = 2
  194. CONNECTED = 5
  195. def __init__(self, tt, server, proxy=None):
  196. asynchat.async_chat.__init__(self)
  197. self.data_source = tt.data_source.copy()
  198. self.inbuf = b''
  199. self.proxy = proxy
  200. self.server = server
  201. self.tt = tt
  202. self.testname = "send-data%s"%id(self)
  203. self.set_terminator(None)
  204. dest = (self.proxy or self.server)
  205. self.create_socket(addr_to_family(dest[0]), socket.SOCK_STREAM)
  206. debug("socket %d connecting to %r..."%(self.fileno(),dest))
  207. self.state = self.CONNECTING
  208. self.connect(dest)
  209. def get_test_names(self):
  210. return [ self.testname ]
  211. def sent_ok(self):
  212. self.tt.success(self.testname)
  213. def handle_connect(self):
  214. if self.proxy:
  215. self.state = self.CONNECTING_THROUGH_PROXY
  216. self.push(socks_cmd(self.server))
  217. else:
  218. self.state = self.CONNECTED
  219. self.push_output()
  220. def collect_incoming_data(self, data):
  221. self.inbuf += data
  222. if self.state == self.CONNECTING_THROUGH_PROXY:
  223. if len(self.inbuf) >= 8:
  224. if self.inbuf[:2] == b'\x00\x5a':
  225. debug("proxy handshake successful (fd=%d)" % self.fileno())
  226. self.state = self.CONNECTED
  227. debug("successfully connected (fd=%d)" % self.fileno())
  228. self.inbuf = self.inbuf[8:]
  229. self.push_output()
  230. else:
  231. debug("proxy handshake failed (0x%x)! (fd=%d)" %
  232. (ord(self.inbuf[1]), self.fileno()))
  233. self.state = self.NOT_CONNECTED
  234. self.close()
  235. def push_output(self):
  236. self.push_with_producer(self.data_source)
  237. self.push_with_producer(CloseSourceProducer(self))
  238. self.close_when_done()
  239. def fileno(self):
  240. return self.socket.fileno()
  241. class TrafficTester(object):
  242. """
  243. Hang on select.select() and dispatch to Sources and Sinks.
  244. Time out after self.timeout seconds.
  245. Keep track of successful and failed data verification using a
  246. TestSuite.
  247. Return True if all tests succeed, else False.
  248. """
  249. def __init__(self,
  250. endpoint,
  251. data=b"",
  252. timeout=3,
  253. repetitions=1,
  254. dot_repetitions=0):
  255. self.listener = Listener(self, endpoint)
  256. self.pending_close = []
  257. self.timeout = timeout
  258. self.tests = TestSuite()
  259. self.data_source = DataSource(data, repetitions)
  260. # sanity checks
  261. self.dot_repetitions = dot_repetitions
  262. debug("listener fd=%d" % self.listener.fileno())
  263. def add(self, item):
  264. """Register a single item."""
  265. # We used to hold on to these items for their fds, but now
  266. # asyncore manages them for us.
  267. if hasattr(item, "get_test_names"):
  268. for name in item.get_test_names():
  269. self.tests.add(name)
  270. def success(self, name):
  271. """Declare that a single test has passed."""
  272. self.tests.success(name)
  273. def failure(self, name):
  274. """Declare that a single test has failed."""
  275. self.tests.failure(name)
  276. def run(self):
  277. start = now = time.time()
  278. end = time.time() + self.timeout
  279. while now < end and not self.tests.all_done():
  280. # run only one iteration at a time, with a nice short timeout, so we
  281. # can actually detect completion and timeouts.
  282. asyncore.loop(0.2, False, None, 1)
  283. now = time.time()
  284. debug("Test status: %s"%self.tests.status())
  285. if not debug_flag:
  286. sys.stdout.write('\n')
  287. sys.stdout.flush()
  288. debug("Done with run(); all_done == %s and failure_count == %s"
  289. %(self.tests.all_done(), self.tests.failure_count()))
  290. self.listener.close()
  291. return self.tests.all_done() and self.tests.failure_count() == 0
  292. def main():
  293. """Test the TrafficTester by sending and receiving some data."""
  294. DATA = b"a foo is a bar" * 1000
  295. bind_to = ('localhost', int(sys.argv[1]))
  296. tt = TrafficTester(bind_to, DATA)
  297. # Don't use a proxy for self-testing, so that we avoid tor entirely
  298. tt.add(Source(tt, bind_to))
  299. success = tt.run()
  300. if success:
  301. return 0
  302. return 255
  303. if __name__ == '__main__':
  304. sys.exit(main())