1
0

StreamCheck.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3. from cStringIO import StringIO
  4. from binascii import b2a_hex
  5. from socket import error as socketerror
  6. from urllib import quote
  7. from traceback import print_exc
  8. import Connecter
  9. try:
  10. True
  11. except:
  12. True = 1
  13. False = 0
  14. DEBUG = False
  15. protocol_name = 'BitTorrent protocol'
  16. option_pattern = chr(0)*8
  17. def toint(s):
  18. return long(b2a_hex(s), 16)
  19. def tobinary(i):
  20. return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
  21. chr((i >> 8) & 0xFF) + chr(i & 0xFF))
  22. hexchars = '0123456789ABCDEF'
  23. hexmap = []
  24. for i in xrange(256):
  25. hexmap.append(hexchars[(i&0xF0)/16]+hexchars[i&0x0F])
  26. def tohex(s):
  27. r = []
  28. for c in s:
  29. r.append(hexmap[ord(c)])
  30. return ''.join(r)
  31. def make_readable(s):
  32. if not s:
  33. return ''
  34. if quote(s).find('%') >= 0:
  35. return tohex(s)
  36. return '"'+s+'"'
  37. def toint(s):
  38. return long(b2a_hex(s), 16)
  39. # header, reserved, download id, my id, [length, message]
  40. streamno = 0
  41. class StreamCheck:
  42. def __init__(self):
  43. global streamno
  44. self.no = streamno
  45. streamno += 1
  46. self.buffer = StringIO()
  47. self.next_len, self.next_func = 1, self.read_header_len
  48. def read_header_len(self, s):
  49. if ord(s) != len(protocol_name):
  50. print self.no, 'BAD HEADER LENGTH'
  51. return len(protocol_name), self.read_header
  52. def read_header(self, s):
  53. if s != protocol_name:
  54. print self.no, 'BAD HEADER'
  55. return 8, self.read_reserved
  56. def read_reserved(self, s):
  57. return 20, self.read_download_id
  58. def read_download_id(self, s):
  59. if DEBUG:
  60. print self.no, 'download ID ' + tohex(s)
  61. return 20, self.read_peer_id
  62. def read_peer_id(self, s):
  63. if DEBUG:
  64. print self.no, 'peer ID' + make_readable(s)
  65. return 4, self.read_len
  66. def read_len(self, s):
  67. l = toint(s)
  68. if l > 2 ** 23:
  69. print self.no, 'BAD LENGTH: '+str(l)+' ('+s+')'
  70. return l, self.read_message
  71. def read_message(self, s):
  72. if not s:
  73. return 4, self.read_len
  74. m = s[0]
  75. if ord(m) > 8:
  76. print self.no, 'BAD MESSAGE: '+str(ord(m))
  77. if m == Connecter.REQUEST:
  78. if len(s) != 13:
  79. print self.no, 'BAD REQUEST SIZE: '+str(len(s))
  80. return 4, self.read_len
  81. index = toint(s[1:5])
  82. begin = toint(s[5:9])
  83. length = toint(s[9:])
  84. print self.no, 'Request: '+str(index)+': '+str(begin)+'-'+str(begin)+'+'+str(length)
  85. elif m == Connecter.CANCEL:
  86. if len(s) != 13:
  87. print self.no, 'BAD CANCEL SIZE: '+str(len(s))
  88. return 4, self.read_len
  89. index = toint(s[1:5])
  90. begin = toint(s[5:9])
  91. length = toint(s[9:])
  92. print self.no, 'Cancel: '+str(index)+': '+str(begin)+'-'+str(begin)+'+'+str(length)
  93. elif m == Connecter.PIECE:
  94. index = toint(s[1:5])
  95. begin = toint(s[5:9])
  96. length = len(s)-9
  97. print self.no, 'Piece: '+str(index)+': '+str(begin)+'-'+str(begin)+'+'+str(length)
  98. else:
  99. print self.no, 'Message '+str(ord(m))+' (length '+str(len(s))+')'
  100. return 4, self.read_len
  101. def write(self, s):
  102. while True:
  103. i = self.next_len - self.buffer.tell()
  104. if i > len(s):
  105. self.buffer.write(s)
  106. return
  107. self.buffer.write(s[:i])
  108. s = s[i:]
  109. m = self.buffer.getvalue()
  110. self.buffer.reset()
  111. self.buffer.truncate()
  112. x = self.next_func(m)
  113. self.next_len, self.next_func = x