encrypter.py

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

# 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 urllib import quote
from traceback import print_exc
from BitTornado.BTcrypto import Crypto

try:
    True
except:
    True = 1
    False = 0
    bool = lambda x: not not x

DEBUG = False

MAX_INCOMPLETE = 8

protocol_name = 'BitTorrent protocol'
option_pattern = chr(0)*8

def toint(s):
    return long(b2a_hex(s), 16)

def tobinary16(i):
    return chr((i >> 8) & 0xFF) + chr(i & 0xFF)

hexchars = '0123456789ABCDEF'
hexmap = []
for i in xrange(256):
    hexmap.append(hexchars[(i&0xF0)/16]+hexchars[i&0x0F])

def tohex(s):
    r = []
    for c in s:
        r.append(hexmap[ord(c)])
    return ''.join(r)

def make_readable(s):
    if not s:
        return ''
    if quote(s).find('%') >= 0:
        return tohex(s)
    return '"'+s+'"'
   

class IncompleteCounter:
    def __init__(self):
        self.c = 0
    def increment(self):
        self.c += 1
    def decrement(self):
        self.c -= 1
    def toomany(self):
        return self.c >= MAX_INCOMPLETE
    
incompletecounter = IncompleteCounter()


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

class Connection:
    def __init__(self, Encoder, connection, id,
                 ext_handshake=False, encrypted = None, options = None):
        self.Encoder = Encoder
        self.connection = connection
        self.connecter = Encoder.connecter
        self.id = id
        self.locally_initiated = (id != None)
        self.readable_id = make_readable(id)
        self.complete = False
        self.keepalive = lambda: None
        self.closed = False
        self.buffer = ''
        self.bufferlen = None
        self.log = None
        self.read = self._read
        self.write = self._write
        self.cryptmode = 0
        self.encrypter = None
        if self.locally_initiated:
            incompletecounter.increment()
            if encrypted:
                self.encrypted = True
                self.encrypter = Crypto(True)
                self.write(self.encrypter.pubkey+self.encrypter.padding())
            else:
                self.encrypted = False
                self.write(chr(len(protocol_name)) + protocol_name + 
                    option_pattern + self.Encoder.download_id )
            self.next_len, self.next_func = 1+len(protocol_name), self.read_header
        elif ext_handshake:
            self.Encoder.connecter.external_connection_made += 1
            if encrypted:   # passed an already running encrypter
                self.encrypter = encrypted
                self.encrypted = True
                self._start_crypto()
                self.next_len, self.next_func = 14, self.read_crypto_block3c
            else:
                self.encrypted = False
                self.options = options
                self.write(self.Encoder.my_id)
                self.next_len, self.next_func = 20, self.read_peer_id
        else:
            self.encrypted = None       # don't know yet
            self.next_len, self.next_func = 1+len(protocol_name), self.read_header
        self.Encoder.raw_server.add_task(self._auto_close, 30)


    def _log_start(self):   # only called with DEBUG = True
        self.log = open('peerlog.'+self.get_ip()+'.txt','a')
        self.log.write('connected - ')
        if self.locally_initiated:
            self.log.write('outgoing\n')
        else:
            self.log.write('incoming\n')
        self._logwritefunc = self.write
        self.write = self._log_write

    def _log_write(self, s):
        self.log.write('w:'+b2a_hex(s)+'\n')
        self._logwritefunc(s)
        

    def get_ip(self, real=False):
        return self.connection.get_ip(real)

    def get_id(self):
        return self.id

    def get_readable_id(self):
        return self.readable_id

    def is_locally_initiated(self):
        return self.locally_initiated

    def is_encrypted(self):
        return bool(self.encrypted)

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

    def _read_header(self, s):
        if s == chr(len(protocol_name))+protocol_name:
            return 8, self.read_options
        return None

    def read_header(self, s):
        if self._read_header(s):
            if self.encrypted or self.Encoder.config['crypto_stealth']:
                return None
            return 8, self.read_options
        if self.locally_initiated and not self.encrypted:
            return None
        elif not self.Encoder.config['crypto_allowed']:
            return None
        if not self.encrypted:
            self.encrypted = True
            self.encrypter = Crypto(self.locally_initiated)
        self._write_buffer(s)
        return self.encrypter.keylength, self.read_crypto_header

    ################## ENCRYPTION SUPPORT ######################

    def _start_crypto(self):
        self.encrypter.setrawaccess(self._read,self._write)
        self.write = self.encrypter.write
        self.read = self.encrypter.read
        if self.buffer:
            self.buffer = self.encrypter.decrypt(self.buffer)

    def _end_crypto(self):
        self.read = self._read
        self.write = self._write
        self.encrypter = None

    def read_crypto_header(self, s):
        self.encrypter.received_key(s)
        self.encrypter.set_skey(self.Encoder.download_id)
        if self.locally_initiated:
            if self.Encoder.config['crypto_only']:
                cryptmode = '\x00\x00\x00\x02'    # full stream encryption
            else:
                cryptmode = '\x00\x00\x00\x03'    # header or full stream
            padc = self.encrypter.padding()
            self.write( self.encrypter.block3a
                      + self.encrypter.block3b
                      + self.encrypter.encrypt(
                            ('\x00'*8)            # VC
                          + cryptmode             # acceptable crypto modes
                          + tobinary16(len(padc))
                          + padc                  # PadC
                          + '\x00\x00' ) )        # no initial payload data
            self._max_search = 520
            return 1, self.read_crypto_block4a
        self.write(self.encrypter.pubkey+self.encrypter.padding())
        self._max_search = 520
        return 0, self.read_crypto_block3a

    def _search_for_pattern(self, s, pat):
        p = s.find(pat)
        if p < 0:
            if len(s) >= len(pat):
                self._max_search -= len(s)+1-len(pat)
            if self._max_search < 0:
                self.close()
                return False
            self._write_buffer(s[1-len(pat):])
            return False
        self._write_buffer(s[p+len(pat):])
        return True

    ### INCOMING CONNECTION ###

    def read_crypto_block3a(self, s):
        if not self._search_for_pattern(s,self.encrypter.block3a):
            return -1, self.read_crypto_block3a     # wait for more data
        return len(self.encrypter.block3b), self.read_crypto_block3b

    def read_crypto_block3b(self, s):
        if s != self.encrypter.block3b:
            return None
        self.Encoder.connecter.external_connection_made += 1
        self._start_crypto()
        return 14, self.read_crypto_block3c

    def read_crypto_block3c(self, s):
        if s[:8] != ('\x00'*8):             # check VC
            return None
        self.cryptmode = toint(s[8:12]) % 4
        if self.cryptmode == 0:
            return None                     # no encryption selected
        if ( self.cryptmode == 1            # only header encryption
             and self.Encoder.config['crypto_only'] ):
            return None
        padlen = (ord(s[12])<<8)+ord(s[13])
        if padlen > 512:
            return None
        return padlen+2, self.read_crypto_pad3

    def read_crypto_pad3(self, s):
        s = s[-2:]
        ialen = (ord(s[0])<<8)+ord(s[1])
        if ialen > 65535:
            return None
        if self.cryptmode == 1:
            cryptmode = '\x00\x00\x00\x01'    # header only encryption
        else:
            cryptmode = '\x00\x00\x00\x02'    # full stream encryption
        padd = self.encrypter.padding()
        self.write( ('\x00'*8)            # VC
                  + cryptmode             # encryption mode
                  + tobinary16(len(padd))
                  + padd )                # PadD
        if ialen:
            return ialen, self.read_crypto_ia
        return self.read_crypto_block3done()

    def read_crypto_ia(self, s):
        if DEBUG:
            self._log_start()
            self.log.write('r:'+b2a_hex(s)+'(ia)\n')
            if self.buffer:
                self.log.write('r:'+b2a_hex(self.buffer)+'(buffer)\n')
        return self.read_crypto_block3done(s)

    def read_crypto_block3done(self, ia=''):
        if DEBUG:
            if not self.log:
                self._log_start()
        if self.cryptmode == 1:     # only handshake encryption
            assert not self.buffer  # oops; check for exceptions to this
            self._end_crypto()
        if ia:
            self._write_buffer(ia)
        return 1+len(protocol_name), self.read_encrypted_header

    ### OUTGOING CONNECTION ###

    def read_crypto_block4a(self, s):
        if not self._search_for_pattern(s,self.encrypter.VC_pattern()):
            return -1, self.read_crypto_block4a     # wait for more data
        self._start_crypto()
        return 6, self.read_crypto_block4b

    def read_crypto_block4b(self, s):
        self.cryptmode = toint(s[:4]) % 4
        if self.cryptmode == 1:             # only header encryption
            if self.Encoder.config['crypto_only']:
                return None
        elif self.cryptmode != 2:
            return None                     # unknown encryption
        padlen = (ord(s[4])<<8)+ord(s[5])
        if padlen > 512:
            return None
        if padlen:
            return padlen, self.read_crypto_pad4
        return self.read_crypto_block4done()

    def read_crypto_pad4(self, s):
        # discard data
        return self.read_crypto_block4done()

    def read_crypto_block4done(self):
        if DEBUG:
            self._log_start()
        if self.cryptmode == 1:     # only handshake encryption
            if not self.buffer:  # oops; check for exceptions to this
                return None
            self._end_crypto()
        self.write(chr(len(protocol_name)) + protocol_name + 
            option_pattern + self.Encoder.download_id)
        return 1+len(protocol_name), self.read_encrypted_header

    ### START PROTOCOL OVER ENCRYPTED CONNECTION ###

    def read_encrypted_header(self, s):
        return self._read_header(s)

    ################################################

    def read_options(self, s):
        self.options = s
        return 20, self.read_download_id

    def read_download_id(self, s):