connector.py

bittorrent source by python. please enjoy

                self.connection.write(chr(len(protocol_name)) + protocol_name +
                                      FLAGS + self.parent.download_id)
                yield padlen
            else:
                dhstr = self._message
                yield DH_BYTES - len(dhstr)
                dhstr += self._message
                privkey = bytetonum(urandom(20))
                pub = numtobyte(pow(2, privkey, dh_prime))
                self.connection.write(pub + urandom(randrange(PAD_MAX)))
                pub = bytetonum(dhstr)
                S = numtobyte(pow(pub, privkey, dh_prime))
                dhstr = pub = privkey = None
                streamid = sha('req1' + S).digest()
                x = ''
                while 1:
                    yield 1
                    x += self._message
                    i = (x + self._rest).find(streamid)
                    if i >= 0:
                        break
                    yield len(self._rest)
                    x += self._message
                    if len(x) >= 532:
                        self.protocol_violation('incoming VC not found',
                                           self.connection)
                        return
                yield i + 20 + 20 + 8 + 4 + 2 - len(x)
                self._message = (x + self._message)[-34:]
                streamid = self._message[0:20]
                x = sha('req3' + S).digest()
                streamid = ''.join([chr(ord(streamid[i]) ^ ord(x[i]))
                                    for i in range(20)])
                self.parent.select_torrent_obfuscated(self, streamid)
                if self.parent.download_id is None:
                    self.protocol_violation('download id unknown/rejected',
                                       self.connection)
                    return
                self.logger = logging.getLogger(
                    self.log_prefix + '.' + repr(self.parent.download_id) +
                    '.peer_id_not_yet' )
                SKEY = self.parent.download_id
                decrypt = ARC4.new(sha('keyA' + S + SKEY).digest()).decrypt
                decrypt('x'*1024)
                s = decrypt(self._message[20:34])
                if s[0:8] != '\x00' * 8:
                    self.protocol_violation('BAD VC', self.connection)
                    return
                crypto_provide = toint(s[8:12])
                padlen = (ord(s[12]) << 8) + ord(s[13])
                if padlen > 512:
                    self.protocol_violation('BAD padlen, too long', self.connection)
                    return
                self._decrypt = decrypt
                yield padlen + 2
                s = self._message
                encrypt = ARC4.new(sha('keyB' + S + SKEY).digest()).encrypt
                encrypt('x'*1024)
                self.connection.encrypt = encrypt
                if not crypto_provide & 2:
                    self.protocol_violation("peer doesn't support crypto mode 2", self.connection)
                    return
                padlen = randrange(PAD_MAX)
                s = '\x00' * 11 + '\x02\x00' + chr(padlen) + urandom(padlen)
                self.connection.write(s)
            S = SKEY = s = x = streamid = VC = padlen = None
            yield 1 + len(protocol_name)
            if self._message != chr(len(protocol_name)) + protocol_name:
                self.protocol_violation('classic handshake fails', self.connection)
                return

        yield 8  # reserved
        # dht is on last reserved byte
        if ord(self._message[7]) & DHT:
            self.uses_dht = True
        if ord(self._message[7]) & FAST_EXTENSION:
            if disable_fast_extension:
                self.uses_fast_extension = False
            else:
                if noisy: log( "Implements FAST_EXTENSION")
                self.uses_fast_extension = True
        
        yield 20 # download id (i.e., infohash)
        if self.parent.download_id is None:  # incoming connection
            # modifies self.parent if successful
            self.parent.select_torrent(self, self._message)
            if self.parent.download_id is None:
                self.protocol_violation("no download_id from parent (peer from a torrent you're not running)", self.connection)
                return
        elif self._message != self.parent.download_id:
            self.protocol_violation("incorrect download_id from parent",
                               self.connection)
            return
        if not self.locally_initiated:
            self.connection.write(chr(len(protocol_name)) + protocol_name +
                FLAGS + self.parent.download_id + self.parent.my_id)

        yield 20  # peer id
        if not self.id:
            self.id = self._message
            ns = (self.log_prefix + '.' + repr(self.parent.download_id) +
                  '.' + self._message.encode('hex'))
            self.logger = logging.getLogger(ns)

            if self.id == self.parent.my_id:
                self.protocol_violation("talking to self", self.connection)
                return
            for v in self.parent.connections.itervalues():
                if v is not self:
                    if v.id == self.id:
                        self.protocol_violation(
                            "duplicate connection (id collision)",
                            self.connection)
                        return
                    if (self.parent.config['one_connection_per_ip'] and
                        v.ip == self.ip):
                        self.protocol_violation(
                            "duplicate connection (ip collision)",
                            self.connection)
                        return
            if self.locally_initiated:
                self.connection.write(self.parent.my_id)
            else:
                self.parent.everinc = True
        else:
            if self._message != self.id:
                self.protocol_violation("incorrect id")
                return
        self.complete = True
        self.parent.connection_completed(self)

        while True:
            yield 4   # message length
            l = toint(self._message)
            if l > self.parent.config['max_message_length']:
                self.protocol_violation("message length exceeds max (%s %s)" %
                    (l, self.parent.config['max_message_length']),
                    self.connection)
                return
            if l > 0:
                yield l
                self._got_message(self._message)

    def _got_message(self, message):
        t = message[0]
        #if noisy: log( "GOT %s" % message_dict[t] )
        if t in [BITFIELD, HAVE_ALL, HAVE_NONE] and self.got_anything:
            self.close()
            return
        self.got_anything = True
        if (t in [CHOKE, UNCHOKE, INTERESTED, NOT_INTERESTED] and
                len(message) != 1):
            self.close()
            return
        if t == CHOKE:
            if noisy: log( "GOT %s" % message_dict[t] )
            self.download.got_choke()
        elif t == UNCHOKE:
            if noisy: log( "GOT %s" % message_dict[t] )
            self.download.got_unchoke()
        elif t == INTERESTED:
            if noisy: log( "GOT %s" % message_dict[t] )
            self.upload.got_interested()
        elif t == NOT_INTERESTED:
            if noisy: log( "GOT %s" % message_dict[t] )
            self.upload.got_not_interested()
        elif t == HAVE:
            i = unpack("!xi", message)[0]
            if noisy: log( "GOT HAVE %d" % i )
            if i >= self.parent.numpieces:
                self.close()
                return
            self.download.got_have(i)
        elif t == BITFIELD:
            try:
                b = Bitfield(self.parent.numpieces, message[1:])
            except ValueError:
                self.close()
                return
            self.download.got_have_bitfield(b)
        elif t == REQUEST:
            if len(message) != 13:
                self.close()
                return
            i, a, b = unpack("!xiii", message)
            if noisy: log( "GOT REQUEST %d %d %d" % (i, a, b) )
            if i >= self.parent.numpieces:
                self.protocol_violation(
                     "Requested piece index %d > numpieces which is %d" %
                     (i,self.parent.numpieces), self.connection )
                self.close()
                return
            if self.download.have[i]:
                self.protocol_violation(
                     "Requested piece index %d which the peer already has" %
                     (i,), self.connection )
                self.close()
                return
            self.upload.got_request(i, a, b)
        elif t == CANCEL:
            if len(message) != 13:
                self.protocol_violation(
                    "Invalid cancel message length %d" % len(message),
                    self.connection )
                self.close()
                return
            i, a, b = unpack("!xiii", message)
            if noisy: log( "GOT CANCEL %d %d %d" % (i, a, b) )
            if i >= self.parent.numpieces:
                self.protocol_violation(
                     "Cancelled piece index %d > numpieces which is %d" %
                     (i,self.parent.numpieces), self.connection )
                self.close()
                return
            self.upload.got_cancel(i, a, b)
        elif t == PIECE:
            if len(message) <= 9:
                self.close()
                return
            n = len(message) - 9
            i, a, b = unpack("!xii%ss" % n, message)
            if noisy: log( "GOT PIECE %d %d" % (i, a) )
            if i >= self.parent.numpieces:
                self.close()
                return
            self.download.got_piece(i, a, b)
        elif t == PORT:
            if len(message) != 3:
                self.close()
                return
            self.dht_port = unpack('!H', message[1:3])[0]
            self.parent.got_port(self)
        elif t == SUGGEST_PIECE:
            if not self.uses_fast_extension:
                self.protocol_violation(
                    "Received 'SUGGEST_PIECE' when fast extension disabled.",
                    self.connection )
                self.close()
                return
            i = unpack("!xi", message)[0]
            if noisy: log( "GOT SUGGEST_PIECE %d" % i )
            if i >= self.parent.numpieces:
                self.protocol_violation(
                    "Received 'SUGGEST_PIECE' with piece id %d > numpieces." %
                    self.parent.numpieces, self.connection )
                self.close()
                return
            self.download.got_suggest_piece(i)
        elif t == HAVE_ALL:
            if noisy: log( "GOT %s" % message_dict[t] )
            if not self.uses_fast_extension:
                self.protocol_violation(
                    "Received 'HAVE_ALL' when fast extension disabled.",
                    self.connection )
                self.close()
                return
            self.download.got_have_all()
        elif t == HAVE_NONE:
            if noisy: log( "GOT %s" % message_dict[t] )
            if not self.uses_fast_extension:
                self.protocol_violation(
                    "Received 'HAVE_NONE' when fast extension disabled.",
                    self.connection )
                self.close()
                return
            self.download.got_have_none()
        elif t == REJECT_REQUEST:
            if not self.uses_fast_extension:
                self.protocol_violation(
                    "Received 'REJECT_REQUEST' when fast extension disabled.",
                    self.connection )
                self.close()
                return
            if len(message) != 13:
                self.protocol_violation(
                    "Received 'REJECT_REQUEST' with length %d != 13." %
                    len(message), self.connection )
                self.close()
                return
            i, a, b = unpack("!xiii", message)
            if noisy: log( "GOT REJECT_REQUEST %d %d" % (i,a) )
            if i >= self.parent.numpieces:
                self.close()
                return
            self.download.got_reject_request(i, a, b)
        elif t == ALLOWED_FAST:
            if not self.uses_fast_extension:
                self.protocol_violation(
                    "Received 'ALLOWED_FAST' when fast extension disabled.",
                    self.connection )
                self.close()
                return
            i = unpack("!xi", message)[0]
            if noisy: log( "GOT ALLOWED_FAST %d" % i )
            self.download.got_allowed_fast(i)
        else:
            self.close()

    def _send_message(self, message):
        if self.closed:
            return
        s = tobinary(len(message)) + message
        if self._partial_message is not None:
            self._outqueue.append(s)
        else:
            self.connection.write(s)

    def data_came_in(self, conn, s):
        self.received_data = True
        if not self.download:
            # this is really annoying.
            self.sloppy_pre_connection_counter += len(s)
        else:
            l = self.sloppy_pre_connection_counter + len(s)
            self.sloppy_pre_connection_counter = 0
            
        while True:
            if self.closed:
                return
            i = self._next_len - self._buffer_len
            if i > len(s):
                self._buffer.append(s)
                self._buffer_len += len(s)
                return
            m = s[:i]
            if self._buffer_len > 0:
                self._buffer.append(m)
                m = ''.join(self._buffer)
                self._buffer = []
                self._buffer_len = 0
            s = s[i:]
            if self._decrypt is not None:
                m = self._decrypt(m)
            self._message = m
            self._rest = s
            try:
                self._next_len = self._reader.next()
            except StopIteration:
                self.close()
                return

    def _optional_restart(self):            
        if self.locally_initiated and not self.obfuscate_outgoing and not self.received_data:
            dns = (self.connection.ip, self.connection.port)
            self.parent.start_connection(dns, id=None, encrypt=True)

    def connection_lost(self, conn):
        assert conn is self.connection
        self.closed = True
        self._reader = None
        del self.parent.connections[self.connection]
        # ARG. Thanks Uoti
        if hasattr(self.parent, 'ratelimiter'):
            self.parent.ratelimiter.dequeue(self)
        self._optional_restart()
        self.connection = None
        self.parent.replace_connection()
        if self.complete:
            self.parent.complete_connections.remove(self)
            if self.download is not None:
                self.download.disconnected()
            self.parent.choker.connection_lost(self)
            self.upload = None
            self.download = None

    def connection_flushed(self, connection):
        if self.complete and self.next_upload is None and (self._partial_message is not None
                                             or (self.upload and self.upload.buffer)):
            if self.lan:
                # bypass upload rate limiter
                self.send_partial(self.parent.ratelimiter.unitsize)
            else:
                self.parent.ratelimiter.queue(self)