NatCheck.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. # Written by Bram Cohen
  11. from cStringIO import StringIO
  12. from socket import error as socketerror
  13. from BitTorrent.RawServer_twisted import Handler
  14. protocol_name = 'BitTorrent protocol'
  15. # header, reserved, download id, my id, [length, message]
  16. class NatCheck(Handler):
  17. def __init__(self, resultfunc, downloadid, peerid, ip, port, rawserver):
  18. self.resultfunc = resultfunc
  19. self.downloadid = downloadid
  20. self.peerid = peerid
  21. self.ip = ip
  22. self.port = port
  23. self.closed = False
  24. self.buffer = StringIO()
  25. self.next_len = 1
  26. self.next_func = self.read_header_len
  27. rawserver.start_connection((ip, port), self)
  28. def connection_made(self, s):
  29. self.connection = s
  30. self.connection.write(chr(len(protocol_name)) + protocol_name +
  31. (chr(0) * 8) + self.downloadid)
  32. def connection_failed(self, s, exception):
  33. self.answer(False)
  34. def answer(self, result):
  35. self.closed = True
  36. try:
  37. self.connection.close()
  38. except AttributeError:
  39. pass
  40. self.resultfunc(result, self.downloadid, self.peerid, self.ip, self.port)
  41. def read_header_len(self, s):
  42. if ord(s) != len(protocol_name):
  43. return None
  44. return len(protocol_name), self.read_header
  45. def read_header(self, s):
  46. if s != protocol_name:
  47. return None
  48. return 8, self.read_reserved
  49. def read_reserved(self, s):
  50. return 20, self.read_download_id
  51. def read_download_id(self, s):
  52. if s != self.downloadid:
  53. return None
  54. return 20, self.read_peer_id
  55. def read_peer_id(self, s):
  56. if s != self.peerid:
  57. return None
  58. self.answer(True)
  59. return None
  60. def data_came_in(self, connection, s):
  61. while True:
  62. if self.closed:
  63. return
  64. i = self.next_len - self.buffer.tell()
  65. if i > len(s):
  66. self.buffer.write(s)
  67. return
  68. self.buffer.write(s[:i])
  69. s = s[i:]
  70. m = self.buffer.getvalue()
  71. self.buffer.reset()
  72. self.buffer.truncate()
  73. x = self.next_func(m)
  74. if x is None:
  75. if not self.closed:
  76. self.answer(False)
  77. return
  78. self.next_len, self.next_func = x
  79. def connection_lost(self, connection):
  80. if not self.closed:
  81. self.closed = True
  82. self.resultfunc(False, self.downloadid, self.peerid, self.ip, self.port)
  83. def connection_flushed(self, connection):
  84. pass