# Written by Arno Bakker, Jie Yang
# see LICENSE.txt for license information
import socket
from binascii import b2a_hex
from struct import pack,unpack
from StringIO import StringIO
DEBUG=False
current_version = 2
lowest_version = 2
protocol_name = "BitTorrent protocol"
# Enable Swapper extensions:
# Left-most bit = Azureus Enhanced Messaging Protocol (AEMP)
# Left+42 bit = Swapper Simple Merkle Hashes extension
# Left+43 bit = Swapper Overlay swarm extension
# Right-most bit = BitTorrent DHT extension
option_pattern = '\x00\x00\x00\x00\x00\x30\x00\x00'
overlay_infohash = '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
def toint(s):
return long(b2a_hex(s), 16)
def tobinary(i):
return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
chr((i >> 8) & 0xFF) + chr(i & 0xFF))
class BTConnection:
def __init__(self,hostname,port,opensock=None):
self.hisport = port
self.buffer = StringIO()
self.myport = 481
self.myid = "".zfill(20)
self.myid = self.myid[:16] + pack('H', lowest_version) + pack('H', current_version)
self.myid = self.myid[:14] + pack('H', self.myport) + self.myid[16:]
self.hisid = None
if opensock:
self.s = opensock
else:
self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.s.connect((hostname, port))
handshake = chr(len(protocol_name))
handshake += protocol_name
handshake += option_pattern
handshake += overlay_infohash
handshake += self.myid
self.s.send(handshake)
def get_my_id(self):
return self.myid
def get_his_id(self):
return self.hisid
def read_handshake(self):
data = self._readn(68)
assert(data[0] == chr(len(protocol_name)))
assert(data[1:20] == protocol_name)
assert(data[20:28] == option_pattern)
assert(data[28:48] == overlay_infohash)
self.hisid = data[48:68]
hisport = unpack('H', self.hisid[14:16])[0]
assert(hisport == self.hisport)
low_ver = unpack('H', self.hisid[16:18])[0]
assert(low_ver == lowest_version)
cur_ver = unpack('H', self.hisid[18:20])[0]
assert(cur_ver == current_version)
def read_handshake_medium_rare(self):
data = self._readn(68)
assert(data[0] == chr(len(protocol_name)))
assert(data[1:20] == protocol_name)
assert(data[20:28] == option_pattern)
assert(data[28:48] == overlay_infohash)
self.hisid = data[48:68]
# don't check encoded fields
def close(self):
self.s.close()
def send(self,data):
""" send length-prefixed message """
self.s.send(tobinary(len(data)))
self.s.send(data)
def recv(self):
""" received length-prefixed message """
size_data = self._readn(4)
size = toint(size_data)
if DEBUG:
print "btconn: waiting for message size",size
return self._readn(size)
def _readn(self,n):
""" read n bytes from socket stream """
nwant = n
while True:
data = self.s.recv(nwant)
if DEBUG:
print "btconn: _readn got",len(data),"bytes"
if len(data) == 0:
raise Exception('arno says connection closed')
nwant -= len(data)
self.buffer.write(data)
if nwant == 0:
break
self.buffer.seek(0)
data = self.buffer.read(n)
self.buffer.seek(0)
return data
|