instrumentation/next-share/BaseLib/Test/btconn.py

   1 # Written by Arno Bakker, Jie Yang
   2 # see LICENSE.txt for license information
   3
   4 import socket
   5 import sys
   6 from binascii import b2a_hex
   7 from struct import pack,unpack
   8 from StringIO import StringIO
   9
  10 from BaseLib.Core.Overlay.SecureOverlay import OLPROTO_VER_CURRENT
  11
  12 DEBUG=False
  13
  14 current_version = OLPROTO_VER_CURRENT
  15 lowest_version = 2
  16
  17 protocol_name = "BitTorrent protocol"
  18 # Enable Tribler extensions:
  19 # Left-most bit = Azureus Enhanced Messaging Protocol (AEMP)
  20 # Left+42 bit = Tribler Simple Merkle Hashes extension v0. Outdated, but still sent for compatibility.
  21 # Left+43 bit = Tribler Overlay swarm extension
  22 # Right-most bit = BitTorrent DHT extension
  23 tribler_option_pattern = '\x00\x00\x00\x00\x00\x30\x00\x00'
  24 overlay_infohash = '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
  25
  26 def toint(s):
  27     return long(b2a_hex(s), 16)
  28
  29 def tobinary(i):
  30     return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
  31         chr((i >> 8) & 0xFF) + chr(i & 0xFF))
  32
  33 class BTConnection:
  34     def __init__(self,hostname,port,opensock=None,user_option_pattern=None,user_infohash=None,myid=None,mylistenport=None,myoversion=None):
  35         assert user_option_pattern is None or isinstance(user_option_pattern, str)
  36         assert user_option_pattern is None or len(user_option_pattern) == 8
  37         assert user_infohash is None or isinstance(user_infohash, str)
  38         assert user_infohash is None or len(user_infohash) == 20
  39         assert myid is None or isinstance(myid, str)
  40         assert myid is None or len(myid) == 20
  41         self.hisport = port
  42         self.buffer = StringIO()
  43         if mylistenport is None:
  44             self.myport = 481
  45         else:
  46             self.myport = mylistenport
  47         if myid is None:
  48             self.myid = "".zfill(20)
  49             if myoversion is None:
  50                 myoversion = current_version
  51             self.myid = self.myid[:16] + pack('<H', lowest_version) + pack('<H', myoversion)
  52             self.myid = self.myid[:14] + pack('<H', self.myport) + self.myid[16:]
  53         else:
  54             self.myid = myid
  55         self.hisid = None
  56
  57         if opensock:
  58             self.s = opensock
  59         else:
  60             self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  61             self.s.connect((hostname, port))
  62         handshake = chr(len(protocol_name))
  63         handshake += protocol_name
  64         if user_option_pattern is None:
  65             handshake += tribler_option_pattern
  66         else:
  67             handshake += user_option_pattern
  68
  69         if user_infohash is None:
  70             self.expected_infohash = overlay_infohash
  71         else:
  72             self.expected_infohash = user_infohash
  73         handshake += self.expected_infohash
  74         handshake += self.myid
  75         if DEBUG:
  76             print >>sys.stderr,"btconn: Sending handshake len",len(handshake)
  77         self.s.send(handshake)
  78
  79     def get_my_id(self):
  80         return self.myid
  81
  82     def get_his_id(self):
  83         return self.hisid
  84
  85     def get_my_fake_listen_port(self):
  86         return self.myport
  87
  88     def read_handshake(self):
  89         data = self._readn(68)
  90         assert(data[0] == chr(len(protocol_name)))
  91         assert(data[1:20] == protocol_name)
  92         assert(data[20:28] == tribler_option_pattern)
  93         assert(data[28:48] == self.expected_infohash)
  94         self.hisid = data[48:68]
  95         hisport = unpack('<H', self.hisid[14:16])[0]
  96         assert(hisport == self.hisport)
  97         low_ver = unpack('<H', self.hisid[16:18])[0]
  98         assert(low_ver == lowest_version)
  99         cur_ver = unpack('<H', self.hisid[18:20])[0]
 100         #if DEBUG:
 101         #    print >> sys.stderr, "btconn: his cur_ver: ", cur_ver
 102         #    print >> sys.stderr, "btconn: my curr_ver: ", current_version
 103         assert(cur_ver == current_version)
 104
 105     def read_handshake_medium_rare(self,close_ok = False):
 106         data = self._readn(68)
 107         if len(data) == 0:
 108             if close_ok:
 109                 return
 110             else:
 111                 assert(len(data) > 0)
 112         assert(data[0] == chr(len(protocol_name)))
 113         assert(data[1:20] == protocol_name)
 114         assert(data[20:28] == tribler_option_pattern)
 115         assert(data[28:48] == self.expected_infohash)
 116         self.hisid = data[48:68]
 117         # don't check encoded fields
 118
 119     def close(self):
 120         self.s.close()
 121
 122     def send(self,data):
 123         """ send length-prefixed message """
 124         self.s.send(tobinary(len(data)))
 125         self.s.send(data)
 126
 127     def recv(self):
 128         """ received length-prefixed message """
 129         size_data = self._readn(4)
 130         if len(size_data) == 0:
 131             return size_data
 132         size = toint(size_data)
 133         if DEBUG and size > 10000:
 134             print >> sys.stderr,"btconn: waiting for message size",size
 135         if size == 0:
 136             # BT keep alive message, don't report upwards
 137             return self.recv()
 138         else:
 139             return self._readn(size)
 140
 141     def _readn(self,n):
 142         """ read n bytes from socket stream """
 143         nwant = n
 144         while True:
 145             try:
 146                 data = self.s.recv(nwant)
 147             except socket.error, e:
 148                 if e[0] == 10035:
 149                     # WSAEWOULDBLOCK on Windows
 150                     continue
 151                 elif e[0] == 10054:
 152                     # WSAECONNRESET on Windows
 153                     print >>sys.stderr,"btconn:",e,"converted to EOF"
 154                     return '' # convert to EOF
 155                 else:
 156                     raise e
 157             if DEBUG:
 158                 print >> sys.stderr,"btconn: _readn got",len(data),"bytes"
 159             if len(data) == 0:
 160                 #raise socket.error(ECONNRESET,'arno says connection closed')
 161                 return data
 162             nwant -= len(data)
 163             self.buffer.write(data)
 164             if nwant == 0:
 165                 break
 166         self.buffer.seek(0)
 167         data = self.buffer.read(n)
 168         self.buffer.seek(0)
 169         return data