Traffic.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476
  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. # Future imports for Python 2.7, mandatory in 3.0
  22. from __future__ import division
  23. from __future__ import print_function
  24. from __future__ import unicode_literals
  25. import sys
  26. import socket
  27. import struct
  28. import errno
  29. import time
  30. import os
  31. import asyncore
  32. import asynchat
  33. from chutney.Debug import debug_flag, debug
  34. def note(s):
  35. sys.stderr.write("NOTE: %s\n"%s)
  36. def warn(s):
  37. sys.stderr.write("WARN: %s\n"%s)
  38. UNIQ_CTR = 0
  39. def uniq(s):
  40. global UNIQ_CTR
  41. UNIQ_CTR += 1
  42. return "%s-%s"%(s,UNIQ_CTR)
  43. if sys.version_info[0] >= 3:
  44. def byte_to_int(b):
  45. return b
  46. else:
  47. def byte_to_int(b):
  48. return ord(b)
  49. def addr_to_family(addr):
  50. for family in [socket.AF_INET, socket.AF_INET6]:
  51. try:
  52. socket.inet_pton(family, addr)
  53. return family
  54. except (socket.error, OSError):
  55. pass
  56. return socket.AF_INET
  57. def socks_cmd(addr_port):
  58. """
  59. Return a SOCKS command for connecting to addr_port.
  60. SOCKSv4: https://en.wikipedia.org/wiki/SOCKS#Protocol
  61. SOCKSv5: RFC1928, RFC1929
  62. """
  63. ver = 4 # Only SOCKSv4 for now.
  64. cmd = 1 # Stream connection.
  65. user = b'\x00'
  66. dnsname = ''
  67. host, port = addr_port
  68. try:
  69. addr = socket.inet_aton(host)
  70. except socket.error:
  71. addr = b'\x00\x00\x00\x01'
  72. dnsname = '%s\x00' % host
  73. debug("Socks 4a request to %s:%d" % (host, port))
  74. if type(dnsname) != type(b""):
  75. dnsname = dnsname.encode("ascii")
  76. return struct.pack('!BBH', ver, cmd, port) + addr + user + dnsname
  77. class TestSuite(object):
  78. """Keep a tab on how many tests are pending, how many have failed
  79. and how many have succeeded."""
  80. def __init__(self):
  81. self.tests = {}
  82. self.not_done = 0
  83. self.successes = 0
  84. self.failures = 0
  85. self.teststatus = {}
  86. def note(self, testname, status):
  87. self.teststatus[testname] = status
  88. def add(self, name):
  89. note("Registering %s"%name)
  90. if name not in self.tests:
  91. debug("Registering %s"%name)
  92. self.not_done += 1
  93. self.tests[name] = 'not done'
  94. else:
  95. warn("... already registered!")
  96. def success(self, name):
  97. note("Success for %s"%name)
  98. if self.tests[name] == 'not done':
  99. debug("Succeeded %s"%name)
  100. self.tests[name] = 'success'
  101. self.not_done -= 1
  102. self.successes += 1
  103. else:
  104. warn("... status was %s"%self.tests.get(name))
  105. def failure(self, name):
  106. note("Failure for %s"%name)
  107. if self.tests[name] == 'not done':
  108. debug("Failed %s"%name)
  109. self.tests[name] = 'failure'
  110. self.not_done -= 1
  111. self.failures += 1
  112. else:
  113. warn("... status was %s"%self.tests.get(name))
  114. def failure_count(self):
  115. return self.failures
  116. def all_done(self):
  117. return self.not_done == 0
  118. def status(self):
  119. return('%s: %d/%d/%d' % (self.tests, self.not_done, self.successes,
  120. self.failures))
  121. class Listener(asyncore.dispatcher):
  122. "A TCP listener, binding, listening and accepting new connections."
  123. def __init__(self, tt, endpoint):
  124. asyncore.dispatcher.__init__(self, map=tt.socket_map)
  125. self.create_socket(addr_to_family(endpoint[0]), socket.SOCK_STREAM)
  126. self.set_reuse_addr()
  127. self.bind(endpoint)
  128. self.listen(0)
  129. self.tt = tt
  130. def writable(self):
  131. return False
  132. def handle_accept(self):
  133. # deprecated in python 3.2
  134. pair = self.accept()
  135. if pair is not None:
  136. newsock, endpoint = pair
  137. debug("new client from %s:%s (fd=%d)" %
  138. (endpoint[0], endpoint[1], newsock.fileno()))
  139. self.tt.add_responder(newsock)
  140. def fileno(self):
  141. return self.socket.fileno()
  142. class DataSource(object):
  143. """A data source generates some number of bytes of data, and then
  144. returns None.
  145. For convenience, it conforms to the 'producer' api.
  146. """
  147. def __init__(self, data, repetitions=1):
  148. self.data = data
  149. self.repetitions = repetitions
  150. self.sent_any = False
  151. def copy(self):
  152. assert not self.sent_any
  153. return DataSource(self.data, self.repetitions)
  154. def more(self):
  155. self.sent_any = True
  156. if self.repetitions > 0:
  157. self.repetitions -= 1
  158. return self.data
  159. return None
  160. class DataChecker(object):
  161. """A data checker verifies its input against bytes in a stream."""
  162. def __init__(self, source):
  163. self.source = source
  164. self.pending = b""
  165. self.succeeded = False
  166. self.failed = False
  167. def consume(self, inp):
  168. if self.failed:
  169. return
  170. if self.succeeded and len(inp):
  171. self.succeeded = False
  172. self.failed = True
  173. return
  174. while len(inp):
  175. n = min(len(inp), len(self.pending))
  176. if inp[:n] != self.pending[:n]:
  177. self.failed = True
  178. return
  179. inp = inp[n:]
  180. self.pending = self.pending[n:]
  181. if not self.pending:
  182. self.pending = self.source.more()
  183. if self.pending is None:
  184. if len(inp):
  185. self.failed = True
  186. else:
  187. self.succeeded = True
  188. return
  189. class Sink(asynchat.async_chat):
  190. "A data sink, reading from its peer and verifying the data."
  191. def __init__(self, sock, tt):
  192. asynchat.async_chat.__init__(self, sock, map=tt.socket_map)
  193. self.set_terminator(None)
  194. self.tt = tt
  195. self.data_checker = DataChecker(tt.data_source.copy())
  196. self.testname = uniq("recv-data")
  197. def get_test_names(self):
  198. return [ self.testname ]
  199. def collect_incoming_data(self, inp):
  200. # shortcut read when we don't ever expect any data
  201. debug("successfully received (bytes=%d)" % len(inp))
  202. self.data_checker.consume(inp)
  203. if self.data_checker.succeeded:
  204. debug("successful verification")
  205. self.close()
  206. self.tt.success(self.testname)
  207. elif self.data_checker.failed:
  208. debug("receive comparison failed")
  209. self.tt.failure(self.testname)
  210. self.close()
  211. def fileno(self):
  212. return self.socket.fileno()
  213. class CloseSourceProducer:
  214. """Helper: when this producer is returned, a source is successful."""
  215. def __init__(self, source):
  216. self.source = source
  217. def more(self):
  218. self.source.note("Flushed")
  219. self.source.sent_ok()
  220. return b""
  221. class Source(asynchat.async_chat):
  222. """A data source, connecting to a TCP server, optionally over a
  223. SOCKS proxy, sending data."""
  224. NOT_CONNECTED = 0
  225. CONNECTING = 1
  226. CONNECTING_THROUGH_PROXY = 2
  227. CONNECTED = 5
  228. def __init__(self, tt, server, proxy=None):
  229. asynchat.async_chat.__init__(self, map=tt.socket_map)
  230. self.data_source = tt.data_source.copy()
  231. self.inbuf = b''
  232. self.proxy = proxy
  233. self.server = server
  234. self.tt = tt
  235. self.testname = uniq("send-data")
  236. self.set_terminator(None)
  237. dest = (self.proxy or self.server)
  238. self.create_socket(addr_to_family(dest[0]), socket.SOCK_STREAM)
  239. debug("socket %d connecting to %r..."%(self.fileno(),dest))
  240. self.state = self.CONNECTING
  241. self.connect(dest)
  242. def get_test_names(self):
  243. return [ self.testname ]
  244. def sent_ok(self):
  245. self.tt.success(self.testname)
  246. def note(self, s):
  247. self.tt.tests.note(self.testname, s)
  248. def handle_connect(self):
  249. if self.proxy:
  250. self.state = self.CONNECTING_THROUGH_PROXY
  251. self.note("connected, sending socks handshake")
  252. self.push(socks_cmd(self.server))
  253. else:
  254. self.state = self.CONNECTED
  255. self.push_output()
  256. def collect_incoming_data(self, data):
  257. self.inbuf += data
  258. if self.state == self.CONNECTING_THROUGH_PROXY:
  259. if len(self.inbuf) >= 8:
  260. if self.inbuf[:2] == b'\x00\x5a':
  261. self.note("proxy handshake successful")
  262. self.state = self.CONNECTED
  263. debug("successfully connected (fd=%d)" % self.fileno())
  264. self.inbuf = self.inbuf[8:]
  265. self.push_output()
  266. else:
  267. debug("proxy handshake failed (0x%x)! (fd=%d)" %
  268. (byte_to_int(self.inbuf[1]), self.fileno()))
  269. self.state = self.NOT_CONNECTED
  270. self.close()
  271. def push_output(self):
  272. self.note("pushed output")
  273. self.push_with_producer(self.data_source)
  274. self.push_with_producer(CloseSourceProducer(self))
  275. def fileno(self):
  276. return self.socket.fileno()
  277. class EchoServer(asynchat.async_chat):
  278. def __init__(self, sock, tt):
  279. asynchat.async_chat.__init__(self, sock, map=tt.socket_map)
  280. self.set_terminator(None)
  281. self.tt = tt
  282. self.am_closing = False
  283. def collect_incoming_data(self, data):
  284. self.push(data)
  285. class EchoClient(Source):
  286. def __init__(self, tt, server, proxy=None):
  287. Source.__init__(self, tt, server, proxy)
  288. self.data_checker = DataChecker(tt.data_source.copy())
  289. self.testname_check = uniq("check")
  290. self.am_closing = False
  291. def enote(self, s):
  292. self.tt.tests.note(self.testname_check, s)
  293. def get_test_names(self):
  294. return [ self.testname, self.testname_check ]
  295. def collect_incoming_data(self, data):
  296. if self.state == self.CONNECTING_THROUGH_PROXY:
  297. Source.collect_incoming_data(self, data)
  298. if self.state == self.CONNECTING_THROUGH_PROXY:
  299. return
  300. data = self.inbuf
  301. self.inbuf = b""
  302. self.data_checker.consume(data)
  303. self.enote("consumed some")
  304. if self.data_checker.succeeded:
  305. self.enote("successful verification")
  306. debug("successful verification")
  307. self.close()
  308. self.tt.success(self.testname_check)
  309. elif self.data_checker.failed:
  310. debug("receive comparison failed")
  311. self.tt.failure(self.testname_check)
  312. self.close()
  313. class TrafficTester(object):
  314. """
  315. Hang on select.select() and dispatch to Sources and Sinks.
  316. Time out after self.timeout seconds.
  317. Keep track of successful and failed data verification using a
  318. TestSuite.
  319. Return True if all tests succeed, else False.
  320. """
  321. def __init__(self,
  322. endpoint,
  323. data=b"",
  324. timeout=3,
  325. repetitions=1,
  326. dot_repetitions=0,
  327. chat_type="Echo"):
  328. if chat_type == "Echo":
  329. self.client_class = EchoClient
  330. self.responder_class = EchoServer
  331. else:
  332. self.client_class = Source
  333. self.responder_class = Sink
  334. self.socket_map = {}
  335. self.listener = Listener(self, endpoint)
  336. self.pending_close = []
  337. self.timeout = timeout
  338. self.tests = TestSuite()
  339. self.data_source = DataSource(data, repetitions)
  340. # sanity checks
  341. self.dot_repetitions = dot_repetitions
  342. debug("listener fd=%d" % self.listener.fileno())
  343. def add(self, item):
  344. """Register a single item."""
  345. # We used to hold on to these items for their fds, but now
  346. # asyncore manages them for us.
  347. if hasattr(item, "get_test_names"):
  348. for name in item.get_test_names():
  349. self.tests.add(name)
  350. def add_client(self, server, proxy=None):
  351. source = self.client_class(self, server, proxy)
  352. self.add(source)
  353. def add_responder(self, socket):
  354. sink = self.responder_class(socket, self)
  355. self.add(sink)
  356. def success(self, name):
  357. """Declare that a single test has passed."""
  358. self.tests.success(name)
  359. def failure(self, name):
  360. """Declare that a single test has failed."""
  361. self.tests.failure(name)
  362. def run(self):
  363. start = now = time.time()
  364. end = time.time() + self.timeout
  365. DUMP_TEST_STATUS_INTERVAL=0.5
  366. dump_at = start+DUMP_TEST_STATUS_INTERVAL
  367. while now < end and not self.tests.all_done():
  368. # run only one iteration at a time, with a nice short timeout, so we
  369. # can actually detect completion and timeouts.
  370. asyncore.loop(0.2, False, self.socket_map, 1)
  371. now = time.time()
  372. if now > dump_at:
  373. debug("Test status: %s"%self.tests.status())
  374. dump_at += DUMP_TEST_STATUS_INTERVAL
  375. if not debug_flag:
  376. sys.stdout.write('\n')
  377. sys.stdout.flush()
  378. debug("Done with run(); all_done == %s and failure_count == %s"
  379. %(self.tests.all_done(), self.tests.failure_count()))
  380. note("Status:\n%s"%self.tests.teststatus)
  381. self.listener.close()
  382. return self.tests.all_done() and self.tests.failure_count() == 0
  383. def main():
  384. """Test the TrafficTester by sending and receiving some data."""
  385. DATA = b"a foo is a bar" * 1000
  386. bind_to = ('localhost', int(sys.argv[1]))
  387. tt = TrafficTester(bind_to, DATA)
  388. # Don't use a proxy for self-testing, so that we avoid tor entirely
  389. tt.add_client(bind_to)
  390. success = tt.run()
  391. if success:
  392. return 0
  393. return 255
  394. if __name__ == '__main__':
  395. sys.exit(main())