encrypter.py

BitTorrentABC-Linux-V.2.4.3源码

# Written by Bram Cohen
# see LICENSE.txt for license information

from cStringIO import StringIO
from binascii import b2a_hex
from socket import error as socketerror
from traceback import print_exc
true = 1
false = 0

protocol_name = 'BitTorrent protocol'

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))

# header, reserved, download id, my id, [length, message]

class Connection:
    def __init__(self, Encoder, connection, id):
        self.Encoder = Encoder
        self.connection = connection
        self.id = id
        self.locally_initiated = (id != None)
        self.complete = false
        self.closed = false
        self.buffer = StringIO()
        self.next_len = 1
        self.next_func = self.read_header_len
        if self.locally_initiated:
            self.send_handshake()

    def send_handshake(self):
        self.connection.write(chr(len(protocol_name)) + protocol_name + 
            (chr(0) * 8) + self.Encoder.download_id + self.Encoder.my_id)

    def get_ip(self):
        return self.connection.get_ip()

    def get_id(self):
        return self.id

    def is_locally_initiated(self):
        return self.locally_initiated

    def is_flushed(self):
        return self.connection.is_flushed()

    def read_header_len(self, s):
        if ord(s) != len(protocol_name):
            return None
        return len(protocol_name), self.read_header

    def read_header(self, s):
        if s != protocol_name:
            return None
        return 8, self.read_reserved

    def read_reserved(self, s):
        return 20, self.read_download_id

    def read_download_id(self, s):
        if s != self.Encoder.download_id:
            return None
        if not self.locally_initiated:
            self.send_handshake()
        return 20, self.read_peer_id

    def read_peer_id(self, s):
        if self.id is None:
            self.id = s
        else:
            if s != self.id:
                return None
        self.complete = self.Encoder.got_id(self)
        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.Encoder.connecter.got_message(self, s)
        return 4, self.read_len

    def read_dead(self, s):
        return None

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

    def sever(self):
        self.closed = true
        del self.Encoder.connections[self.connection]
        if self.complete:
            self.Encoder.connecter.connection_lost(self)

    def send_message(self, message):
        self.connection.write(tobinary(len(message)) + message)

    def data_came_in(self, s):
        while true:
            if self.closed:
                return
            i = self.next_len - self.buffer.tell()
            if i > len(s):
                self.buffer.write(s)
                return
            self.buffer.write(s[:i])
            s = s[i:]
            m = self.buffer.getvalue()
            self.buffer.reset()
            self.buffer.truncate()
            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

class Encoder:
    def __init__(self, connecter, raw_server, my_id, max_len,
            schedulefunc, keepalive_delay, download_id, 
            config):
        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.config = config
        self.connections = {}
        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)
        for c in self.connections.values():
            if c is not None and c.complete:
                c.send_message('')

    def start_connection(self, dns, id):
        if len(self.connections) > self.max_connections:
            return true
        if len(self.connections) >= self.config['max_initiate']:
            return true
        if id == self.my_id:
            return true
        for v in self.connections.values():
            if v is None:
                continue
            if v.id == id:
                return true
            ip = v.get_ip()
            if self.config['security'] and ip != 'unknown' and ip == dns[0]:
                return true
        try:
            c = self.raw_server.start_connection(dns)
            self.connections[c] = Connection(self, c, id)
        except socketerror:
            return false
        return true
    
    def _start_connection(self, dns, id):
        def foo(self=self, dns=dns, id=id):
            self.start_connection(dns, id)
        
        self.schedulefunc(foo, 0)
        
    def got_id(self, connection):
        ip = connection.get_ip()
        for v in self.connections.values():
            if v is None:
                continue
            if connection is not v:
                if connection.id == v.id:
                    connection.close()
                    return false
                if self.config['security'] and ip != 'unknown' and ip == v.get_ip():
                    v.close()
        self.connecter.connection_made(connection)
        return true

    def external_connection_made(self, connection):
        self.connecter.external_connection_made = true
        if len(self.connections) > self.max_connections:
            self.connections[connection] = None
            connection.close()
            return false
        self.connections[connection] = Connection(self, 
            connection, None)
        return true

    def connection_flushed(self, connection):
        c = self.connections[connection]
        if c.complete:
            self.connecter.connection_flushed(c)

    def connection_lost(self, connection):
        if self.connections[connection] is None:
            del self.connections[connection]
            return
        self.connections[connection].sever()
        
    def data_came_in(self, connection, data):
        self.connections[connection].data_came_in(data)

# everything below is for testing

class DummyConnecter:
    def __init__(self):
        self.log = []
        self.close_next = false
    
    def connection_made(self, connection):
        self.log.append(('made', connection))
        
    def connection_lost(self, connection):
        self.log.append(('lost', connection))

    def connection_flushed(self, connection):
        self.log.append(('flushed', connection))

    def got_message(self, connection, message):
        self.log.append(('got', connection, message))
        if self.close_next:
            connection.close()

class DummyRawServer:
    def __init__(self):
        self.connects = []
    
    def start_connection(self, dns):
        c = DummyRawConnection()
        self.connects.append((dns, c))
        return c

class DummyRawConnection:
    def __init__(self):
        self.closed = false
        self.data = []
        self.flushed = true

    def get_ip(self):
        return 'fake.ip'

    def is_flushed(self):
        return self.flushed

    def write(self, data):
        assert not self.closed
        self.data.append(data)
        
    def close(self):
        assert not self.closed
        self.closed = true

    def pop(self):
        r = ''.join(self.data)
        del self.data[:]
        return r

def dummyschedule(a, b):
    pass

def test_messages_in_and_out():
    c = DummyConnecter()
    rs = DummyRawServer()
    e = Encoder(c, rs, 'a' * 20, 500, dummyschedule, 30, 'd' * 20)
    c1 = DummyRawConnection()
    e.external_connection_made(c1)
    assert c1.pop() == chr(len(protocol_name)) + protocol_name + \
        chr(0) * 8 + 'd' * 20 + 'a' * 20
    assert c.log == []
    assert rs.connects == []
    assert not c1.closed

    e.data_came_in(c1, chr(len(protocol_name)) + protocol_name + \
        chr(0) * 8 + 'd' * 20)
    assert c1.pop() == ''
    assert c.log == []
    assert rs.connects == []
    assert not c1.closed

    e.data_came_in(c1, 'b' * 20)
    assert c1.pop() == ''
    assert len(c.log) == 1
    assert c.log[0][0] == 'made'
    ch = c.log[0][1]
    del c.log[:]
    assert rs.connects == []
    assert not c1.closed
    assert ch.get_ip() == 'fake.ip'
    
    ch.send_message('abc')
    assert c1.pop() == chr(0) * 3 + chr(3) + 'abc'
    assert c.log == []
    assert rs.connects == []
    assert not c1.closed
    
    e.data_came_in(c1, chr(0) * 3 + chr(3) + 'def')
    assert c1.pop() == ''
    assert c.log == [('got', ch, 'def')]
    del c.log[:]
    assert rs.connects == []
    assert not c1.closed

def test_flushed():