# Written by Bram Cohen# see LICENSE.txt for license informationfrom cStringIO import StringIOfrom binascii import b2a_hexfrom socket import error as socketerrorfrom urllib import quotefrom traceback import print_excfrom BitTornado.BTcrypto import Cryptotry:    Trueexcept:    True = 1    False = 0    bool = lambda x: not not xDEBUG = FalseMAX_INCOMPLETE = 8protocol_name = 'BitTorrent protocol'option_pattern = chr(0)*8def 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):