encrypter.py

一个多点下载且源码公开的P2P软件

        if ( s != self.Encoder.download_id
             or not self.Encoder.check_ip(ip=self.get_ip()) ):
            return None
        if not self.locally_initiated:
            if not self.encrypted:
                self.Encoder.connecter.external_connection_made += 1
            self.write(chr(len(protocol_name)) + protocol_name + 
                option_pattern + self.Encoder.download_id + self.Encoder.my_id)
        return 20, self.read_peer_id

    def read_peer_id(self, s):
        if not self.encrypted and self.Encoder.config['crypto_only']:
            return None     # allows older trackers to ping,
                            # but won't proceed w/ connections
        if not self.id:
            self.id = s
            self.readable_id = make_readable(s)
        else:
            if s != self.id:
                return None
        self.complete = self.Encoder.got_id(self)
        if not self.complete:
            return None
        if self.locally_initiated:
            self.write(self.Encoder.my_id)
            incompletecounter.decrement()
        self._switch_to_read2()
        c = self.Encoder.connecter.connection_made(self)
        self.keepalive = c.send_keepalive
        return 4, self.read_len

    def read_len(self, s):
        l = toint(s)
        if l > self.Encoder.max_len:
            return None
        return l, self.read_message

    def read_message(self, s):
        if s != '':
            self.connecter.got_message(self, s)
        return 4, self.read_len

    def read_dead(self, s):
        return None

    def _auto_close(self):
        if not self.complete:
            self.close()

    def close(self):
        if not self.closed:
            self.connection.close()
            self.sever()

    def sever(self):
        if self.log:
            self.log.write('closed\n')
            self.log.close()
        self.closed = True
        del self.Encoder.connections[self.connection]
        if self.complete:
            self.connecter.connection_lost(self)
        elif self.locally_initiated:
            incompletecounter.decrement()

    def send_message_raw(self, message):
        self.write(message)

    def _write(self, message):
        if not self.closed:
            self.connection.write(message)

    def data_came_in(self, connection, s):
        self.read(s)

    def _write_buffer(self, s):
        self.buffer = s+self.buffer

    def _read(self, s):
        if self.log:
            self.log.write('r:'+b2a_hex(s)+'\n')
        self.Encoder.measurefunc(len(s))
        self.buffer += s
        while True:
            if self.closed:
                return
            # self.next_len = # of characters function expects
            # or 0 = all characters in the buffer
            # or -1 = wait for next read, then all characters in the buffer
            # not compatible w/ keepalives, switch out after all negotiation complete
            if self.next_len <= 0:
                m = self.buffer
                self.buffer = ''
            elif len(self.buffer) >= self.next_len:
                m = self.buffer[:self.next_len]
                self.buffer = self.buffer[self.next_len:]
            else:
                return
            try:
                x = self.next_func(m)
            except:
                self.next_len, self.next_func = 1, self.read_dead
                raise
            if x is None:
                self.close()
                return
            self.next_len, self.next_func = x
            if self.next_len < 0:  # already checked buffer
                return             # wait for additional data
            if self.bufferlen is not None:
                self._read2('')
                return

    def _switch_to_read2(self):
        self._write_buffer = None
        if self.encrypter:
            self.encrypter.setrawaccess(self._read2,self._write)
        else:
            self.read = self._read2
        self.bufferlen = len(self.buffer)
        self.buffer = [self.buffer]

    def _read2(self, s):    # more efficient, requires buffer['',''] & bufferlen
        if self.log:
            self.log.write('r:'+b2a_hex(s)+'\n')
        self.Encoder.measurefunc(len(s))
        while True:
            if self.closed:
                return
            p = self.next_len-self.bufferlen
            if self.next_len == 0:
                m = ''
            elif s:
                if p > len(s):
                    self.buffer.append(s)
                    self.bufferlen += len(s)
                    return
                self.bufferlen = len(s)-p
                self.buffer.append(s[:p])
                m = ''.join(self.buffer)
                if p == len(s):
                    self.buffer = []
                else:
                    self.buffer=[s[p:]]
                s = ''
            elif p <= 0:
                # assert len(self.buffer) == 1
                s = self.buffer[0]
                self.bufferlen = len(s)-self.next_len
                m = s[:self.next_len]
                if p == 0:
                    self.buffer = []
                else:
                    self.buffer = [s[self.next_len:]]
                s = ''
            else:
                return
            try:
                x = self.next_func(m)
            except:
                self.next_len, self.next_func = 1, self.read_dead
                raise
            if x is None:
                self.close()
                return
            self.next_len, self.next_func = x
            if self.next_len < 0:  # already checked buffer
                return             # wait for additional data
            

    def connection_flushed(self, connection):
        if self.complete:
            self.connecter.connection_flushed(self)

    def connection_lost(self, connection):
        if self.Encoder.connections.has_key(connection):
            self.sever()


class _dummy_banlist:
    def includes(self, x):
        return False

class Encoder:
    def __init__(self, connecter, raw_server, my_id, max_len,
            schedulefunc, keepalive_delay, download_id, 
            measurefunc, config, bans=_dummy_banlist() ):
        self.raw_server = raw_server
        self.connecter = connecter
        self.my_id = my_id
        self.max_len = max_len
        self.schedulefunc = schedulefunc
        self.keepalive_delay = keepalive_delay
        self.download_id = download_id
        self.measurefunc = measurefunc
        self.config = config
        self.connections = {}
        self.banned = {}
        self.external_bans = bans
        self.to_connect = []
        self.paused = False
        if self.config['max_connections'] == 0:
            self.max_connections = 2 ** 30
        else:
            self.max_connections = self.config['max_connections']
        schedulefunc(self.send_keepalives, keepalive_delay)

    def send_keepalives(self):
        self.schedulefunc(self.send_keepalives, self.keepalive_delay)
        if self.paused:
            return
        for c in self.connections.values():
            c.keepalive()

    def start_connections(self, list):
        if not self.to_connect:
            self.raw_server.add_task(self._start_connection_from_queue)
        self.to_connect = list

    def _start_connection_from_queue(self):
        if self.connecter.external_connection_made:
            max_initiate = self.config['max_initiate']
        else:
            max_initiate = int(self.config['max_initiate']*1.5)
        cons = len(self.connections)
        if cons >= self.max_connections or cons >= max_initiate:
            delay = 60
        elif self.paused or incompletecounter.toomany():
            delay = 1
        else:
            delay = 0
            dns, id, encrypted = self.to_connect.pop(0)
            self.start_connection(dns, id, encrypted)
        if self.to_connect:
            self.raw_server.add_task(self._start_connection_from_queue, delay)

    def start_connection(self, dns, id, encrypted = None):
        if ( self.paused
             or len(self.connections) >= self.max_connections
             or id == self.my_id
             or not self.check_ip(ip=dns[0]) ):
            return True
        if self.config['crypto_only']:
            if encrypted is None or encrypted:  # fails on encrypted = 0
                encrypted = True
            else:
                return True
        for v in self.connections.values():
            if v is None:
                continue
            if id and v.id == id:
                return True
            ip = v.get_ip(True)
            if self.config['security'] and ip != 'unknown' and ip == dns[0]:
                return True
        try:
            c = self.raw_server.start_connection(dns)
            con = Connection(self, c, id, encrypted = encrypted)
            self.connections[c] = con
            c.set_handler(con)
        except socketerror:
            return False
        return True

    def _start_connection(self, dns, id, encrypted = None):
        def foo(self=self, dns=dns, id=id, encrypted=encrypted):
            self.start_connection(dns, id, encrypted)
        self.schedulefunc(foo, 0)

    def check_ip(self, connection=None, ip=None):
        if not ip:
            ip = connection.get_ip(True)
        if self.config['security'] and self.banned.has_key(ip):
            return False
        if self.external_bans.includes(ip):
            return False
        return True

    def got_id(self, connection):
        if connection.id == self.my_id:
            self.connecter.external_connection_made -= 1
            return False
        ip = connection.get_ip(True)
        for v in self.connections.values():
            if connection is not v:
                if connection.id == v.id:
                    if ip == v.get_ip(True):
                        v.close()
                    else:
                        return False
                if self.config['security'] and ip != 'unknown' and ip == v.get_ip(True):
                    v.close()
        return True

    def external_connection_made(self, connection):
        if self.paused or len(self.connections) >= self.max_connections:
            connection.close()
            return False
        con = Connection(self, connection, None)
        self.connections[connection] = con
        connection.set_handler(con)
        return True

    def externally_handshaked_connection_made(self, connection, options,
                                              already_read, encrypted = None):
        if ( self.paused
             or len(self.connections) >= self.max_connections
             or not self.check_ip(connection=connection) ):
            connection.close()
            return False
        con = Connection(self, connection, None,
                ext_handshake = True, encrypted = encrypted, options = options)
        self.connections[connection] = con
        connection.set_handler(con)
        if already_read:
            con.data_came_in(con, already_read)
        return True

    def close_all(self):
        for c in self.connections.values():
            c.close()
        self.connections = {}

    def ban(self, ip):
        self.banned[ip] = 1

    def pause(self, flag):
        self.paused = flag