def finishHandshake(self): self.btm = self.factory.btm self.bitfield = Bitfield(self.btm.metainfo.pieces_size) self.upload = BTUpload(self) self.download = BTDownload(self) self.upload.start() self.__uploadMonitor = self.upload._uploadMonitor self.download.start() self.__downloadMonitor = self.download._downloadMonitor self.send_bitfield(self.btm.pieceManager.bitfield) self.send_keep_alive() if self.btm.connectionManager.isAlreadyConnected(self.peer_id) : # Already connected, dropping the connection reactor.callLater(0, self.transport.loseConnection) else: self.factory.addActiveConnection(self.peer_id, self) self.status = 'started'
class BTProtocol(protocol.Protocol): msg_choke = '\x00' msg_unchoke = '\x01' msg_interested = '\x02' msg_not_interested = '\x03' msg_have = '\x04' msg_bitfield = '\x05' msg_request = '\x06' msg_piece = '\x07' msg_cancel = '\x08' msg_port = '\x09' msg_type = {'\x00' : 'choke', '\x01' : 'unchoke', '\x02' : 'interested', '\x03' : 'not_interested', '\x04' : 'have', '\x05' : 'bitfield', '\x06' : 'request', '\x07' : 'piece', '\x08' : 'cancel', '\x09' : 'port'} def __init__(self): self.peer_id = None self.status = None def connectionMade(self): self.status = 'handshake' self.data = '' self._handle_data = self.handle_data() self._next_data_len = self._handle_data.next() self.preHandshake() def finishHandshake(self): self.btm = self.factory.btm self.bitfield = Bitfield(self.btm.metainfo.pieces_size) self.upload = BTUpload(self) self.download = BTDownload(self) self.upload.start() self.__uploadMonitor = self.upload._uploadMonitor self.download.start() self.__downloadMonitor = self.download._downloadMonitor self.send_bitfield(self.btm.pieceManager.bitfield) self.send_keep_alive() if self.btm.connectionManager.isAlreadyConnected(self.peer_id) : # Already connected, dropping the connection reactor.callLater(0, self.transport.loseConnection) else: self.factory.addActiveConnection(self.peer_id, self) self.status = 'started' def connectionLost(self, reason=None): if self.status == 'started': self.upload.stop() self.download.stop() del self.__uploadMonitor del self.__downloadMonitor del self.upload del self.download del self.btm self.factory.removeActiveConnection(self) self.status = 'stopped' def stopConnection(self): if self.connected: self.transport.loseConnection() def send_data(self, data): if not self.connected: return prefix = struct.pack('!I', len(data)) self.transport.write(prefix + data) def send_message(self, _type, data): self.send_data(_type + data) self.__uploadMonitor(_type, data) def __uploadMonitor(self, _type, data): pass def send_handshake(self): info_hash = self.factory.btm.metainfo.info_hash my_id = self.factory.btm.my_peer_id reserved = '\x00'*7 + '\x01' data = '\x13' + 'BitTorrent protocol' + reserved + info_hash + my_id self.transport.write(data) @defer.inlineCallbacks def send_keep_alive(self): yield sleep(60.0) while self.connected: self.send_data('') yield sleep(60.0) def send_choke(self): self.am_choke = True self.send_data(self.msg_choke) def send_unchoke(self): self.am_choke = False self.send_data(self.msg_unchoke) def send_interested(self): self.am_interested = True self.send_data(self.msg_interested) def send_not_interested(self): self.am_interested = False self.send_data(self.msg_not_interested) def send_have(self, index): data = struct.pack('!I', index) self.send_message(self.msg_have, data) def send_bitfield(self, bitfield): if type(bitfield) is str : data = bitfield elif type(bitfield) is Bitfield : data = bitfield.tostring() else : raise TypeError('bitfield should be str or Bitfield') self.send_message(self.msg_bitfield, data) def send_request(self, index, begin, length): data = struct.pack('!III', index, begin, length) self.send_message(self.msg_request, data) def send_piece(self, index, begin, piece): data = struct.pack('!II', index, begin) + piece self.send_message(self.msg_piece, data) def send_cancel(self, idx, begin, length): data = struct.pack('!III', idx, begin, length) self.send_message(self.msg_cancel, data) def send_port(self, port): data = struct.pack('!I', port) self.send_message(self.msg_port, data) def __downloadMonitor(self, data): pass def dataReceived(self, data): self.__downloadMonitor(data) data = self.data + data nd_len = self._next_data_len while len(data) >= nd_len: data_send, data = data[:nd_len], data[nd_len:] nd_len = self._handle_data.send(data_send) self.data = data self._next_data_len = nd_len def handle_data(self): protocol = yield ord((yield 1)) reserved = yield 8 info_hash = yield 20 peer_id = yield 20 self.handle_handshake(protocol, reserved, info_hash, peer_id) self.postHandshake() self.finishHandshake() while True: size, = struct.unpack('!I', (yield 4)) if size == 0 : self.handle_keep_alive() else: _type = yield 1 self.cur_msg_type = _type data = yield (size - 1) method_name = 'handle_'+self.msg_type[_type] method = getattr(self, method_name, None) if method: method(data) else: raise NotImplementedError(method_name) def handle_handshake(self, protocol, reserved, info_hash, peer_id): log.msg('Connected to client ID: {0} v{1}'.format(*identify_client(peer_id))) self.peer_protocol = protocol self.peer_reserved = reserved self.peer_info_hash = info_hash self.peer_id = peer_id def handle_keep_alive(self): pass def handle_choke(self, data): self.download._choke(True) def handle_unchoke(self, data): self.download._choke(False) def handle_interested(self, data): self.upload._interested(True) def handle_not_interested(self, data): self.upload._interested(False) def handle_have(self, data): assert len(data) == 4 index, = struct.unpack('!I', data) self.download._have(index) # print '-->> have index=%d' % index def handle_bitfield(self, data): self.download._bitfield(data) def handle_request(self, data): index, begin, length = struct.unpack('!III', data) self.upload._request(index, begin, length) def handle_piece(self, data): index, begin = struct.unpack('!II', data[:8]) piece = data[8:] self.download._piece(index, begin, piece) def handle_cancel(self, data): index, begin, length = struct.unpack('!III', data) self.upload._cancel(index, begin, length) def handle_port(self, data): if self.btm.app.enable_DHT: dht_fp = open('dht.txt', 'wb') port, = struct.unpack('!H', data) self.dht_port = port addr = self.transport.getPeer().host self.dht_fp.write('{0}\t{1}\n'.format(addr, port)) self.dht_fp.flush() self.btm.connectionManager.handle_port(addr, port)