Exemple #1
0
    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'
Exemple #2
0
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)