# Written by Bram Cohen# see LICENSE.txt for license informationfrom BitTornado.bitfield import Bitfieldfrom sha import shafrom BitTornado.clock import clockfrom traceback import print_excfrom random import randrangetry:    Trueexcept:    True = 1    False = 0try:    from bisect import insortexcept:    def insort(l, item):        l.append(item)        l.sort()DEBUG = FalseSTATS_INTERVAL = 0.2def dummy_status(fractionDone = None, activity = None):    passclass Olist:    def __init__(self, l = []):        self.d = {}        for i in l:            self.d[i] = 1    def __len__(self):        return len(self.d)    def includes(self, i):        return self.d.has_key(i)    def add(self, i):        self.d[i] = 1    def extend(self, l):        for i in l:            self.d[i] = 1    def pop(self, n=0):        # assert self.d        k = self.d.keys()        if n == 0:            i = min(k)        elif n == -1:            i = max(k)        else:            k.sort()            i = k[n]        del self.d[i]        return i    def remove(self, i):        if self.d.has_key(i):            del self.d[i]class fakeflag:    def __init__(self, state=False):        self.state = state    def wait(self):        pass    def isSet(self):        return self.stateclass StorageWrapper:    def __init__(self, storage, request_size, hashes,             piece_size, finished, failed,             statusfunc = dummy_status, flag = fakeflag(), check_hashes = True,            data_flunked = lambda x: None, backfunc = None,            config = {}, unpauseflag = fakeflag(True) ):        self.storage = storage        self.request_size = long(request_size)        self.hashes = hashes        self.piece_size = long(piece_size)        self.piece_length = long(piece_size)        self.finished = finished        self.failed = failed        self.statusfunc = statusfunc        self.flag = flag        self.check_hashes = check_hashes        self.data_flunked = data_flunked        self.backfunc = backfunc        self.config = config        self.unpauseflag = unpauseflag                self.alloc_type = config.get('alloc_type','normal')        self.double_check = config.get('double_check', 0)        self.triple_check = config.get('triple_check', 0)        if self.triple_check:            self.double_check = True        self.bgalloc_enabled = False        self.bgalloc_active = False        self.total_length = storage.get_total_length()        self.amount_left = self.total_length        if self.total_length <= self.piece_size * (len(hashes) - 1):            raise ValueError, 'bad data in responsefile - total too small'        if self.total_length > self.piece_size * len(hashes):            raise ValueError, 'bad data in responsefile - total too big'        self.numactive = [0] * len(hashes)        self.inactive_requests = [1] * len(hashes)        self.amount_inactive = self.total_length        self.amount_obtained = 0        self.amount_desired = self.total_length        self.have = Bitfield(len(hashes))        self.have_cloaked_data = None        self.blocked = [False] * len(hashes)        self.blocked_holes = []        self.blocked_movein = Olist()        self.blocked_moveout = Olist()        self.waschecked = [False] * len(hashes)        self.places = {}        self.holes = []        self.stat_active = {}        self.stat_new = {}        self.dirty = {}        self.stat_numflunked = 0        self.stat_numdownloaded = 0        self.stat_numfound = 0        self.download_history = {}        self.failed_pieces = {}        self.out_of_place = 0        self.write_buf_max = config['write_buffer_size']*1048576L        self.write_buf_size = 0L        self.write_buf = {}   # structure:  piece: [(start, data), ...]        self.write_buf_list = []        self.initialize_tasks = [            ['checking existing data', 0, self.init_hashcheck, self.hashcheckfunc],            ['moving data', 1, self.init_movedata, self.movedatafunc],            ['allocating disk space', 1, self.init_alloc, self.allocfunc] ]        self.backfunc(self._bgalloc,0.1)        self.backfunc(self._bgsync,max(self.config['auto_flush']*60,60))    def _bgsync(self):        if self.config['auto_flush']:            self.sync()        self.backfunc(self._bgsync,max(self.config['auto_flush']*60,60))    def old_style_init(self):        while self.initialize_tasks:            msg, done, init, next = self.initialize_tasks.pop(0)            if init():                self.statusfunc(activity = msg, fractionDone = done)                t = clock() + STATS_INTERVAL                x = 0                while x is not None:                    if t < clock():                        t = clock() + STATS_INTERVAL                        self.statusfunc(fractionDone = x)                    self.unpauseflag.wait()                    if self.flag.isSet():                        return False                    x = next()        self.statusfunc(fractionDone = 0)        return True    def initialize(self, donefunc, statusfunc = None):        self.initialize_done = donefunc        if statusfunc is None:            statusfunc = self.statusfunc        self.initialize_status = statusfunc        self.initialize_next = None                    self.backfunc(self._initialize)    def _initialize(self):        if not self.unpauseflag.isSet():            self.backfunc(self._initialize, 1)            return                if self.initialize_next:            x = self.initialize_next()            if x is None:                self.initialize_next = None            else:                self.initialize_status(fractionDone = x)        else:            if not self.initialize_tasks:                self.initialize_done()                return            msg, done, init, next = self.initialize_tasks.pop(0)            if init():                self.initialize_status(activity = msg, fractionDone = done)                self.initialize_next = next        self.backfunc(self._initialize)    def init_hashcheck(self):        if self.flag.isSet():            return False        self.check_list = []        if len(self.hashes) == 0 or self.amount_left == 0:            self.check_total = 0            self.finished()            return False        self.check_targets = {}        got = {}        for p,v in self.places.items():            assert not got.has_key(v)            got[v] = 1        for i in xrange(len(self.hashes)):            if self.places.has_key(i):  # restored from pickled                self.check_targets[self.hashes[i]] = []                if self.places[i] == i:                    continue                else:                    assert not got.has_key(i)                    self.out_of_place += 1            if got.has_key(i):                continue            if self._waspre(i):                if self.blocked[i]:                    self.places[i] = i                else:                    self.check_list.append(i)                continue            if not self.check_hashes:                self.failed('told file complete on start-up, but data is missing')                return False            self.holes.append(i)            if self.blocked[i] or self.check_targets.has_key(self.hashes[i]):                self.check_targets[self.hashes[i]] = [] # in case of a hash collision, discard            else:                self.check_targets[self.hashes[i]] = [i]        self.check_total = len(self.check_list)        self.check_numchecked = 0.0        self.lastlen = self._piecelen(len(self.hashes) - 1)        self.numchecked = 0.0        return self.check_total > 0    def _markgot(self, piece, pos):        if DEBUG:            print str(piece)+' at '+str(pos)        self.places[piece] = pos        self.have[piece] = True        len = self._piecelen(piece)        self.amount_obtained += len        self.amount_left -= len        self.amount_inactive -= len        self.inactive_requests[piece] = None        self.waschecked[piece] = self.check_hashes        self.stat_numfound += 1    def hashcheckfunc(self):        if self.flag.isSet():            return None        if not self.check_list:            return None                i = self.check_list.pop(0)        if not self.check_hashes:            self._markgot(i, i)        else:            d1 = self.read_raw(i,0,self.lastlen)            if d1 is None:                return None            sh = sha(d1[:])            d1.release()            sp = sh.digest()            d2 = self.read_raw(i,self.lastlen,self._piecelen(i)-self.lastlen)            if d2 is None:                return None            sh.update(d2[:])            d2.release()            s = sh.digest()            if s == self.hashes[i]:                self._markgot(i, i)            elif ( self.check_targets.get(s)                   and self._piecelen(i) == self._piecelen(self.check_targets[s][-1]) ):                self._markgot(self.check_targets[s].pop(), i)                self.out_of_place += 1            elif ( not self.have[-1] and sp == self.hashes[-1]                   and (i == len(self.hashes) - 1                        or not self._waspre(len(self.hashes) - 1)) ):                self._markgot(len(self.hashes) - 1, i)                self.out_of_place += 1            else:                self.places[i] = i        self.numchecked += 1        if self.amount_left == 0:            self.finished()        return (self.numchecked / self.check_total)    def init_movedata(self):        if self.flag.isSet():            return False        if self.alloc_type != 'sparse':            return False        self.storage.top_off()  # sets file lengths to their final size        self.movelist = []        if self.out_of_place == 0:            for i in self.holes:                self.places[i] = i            self.holes = []            return False        self.tomove = float(self.out_of_place)        for i in xrange(len(self.hashes)):            if not self.places.has_key(i):                self.places[i] = i            elif self.places[i] != i:                self.movelist.append(i)        self.holes = []        return True    def movedatafunc(self):        if self.flag.isSet():            return None        if not self.movelist:            return None        i = self.movelist.pop(0)        old = self.read_raw(self.places[i], 0, self._piecelen(i))        if old is None:            return None        if not self.write_raw(i, 0, old):            return None        if self.double_check and self.have[i]:            if self.triple_check:                old.release()                old = self.read_raw( i, 0, self._piecelen(i),                                            flush_first = True )                if old is None:                    return None            if sha(old[:]).digest() != self.hashes[i]:                self.failed('download corrupted; please restart and resume')                return None        old.release()        self.places[i] = i        self.tomove -= 1        return (self.tomove / self.out_of_place)            def init_alloc(self):        if self.flag.isSet():            return False        if not self.holes:            return False        self.numholes = float(len(self.holes))        self.alloc_buf = chr(0xFF) * self.piece_size        if self.alloc_type == 'pre-allocate':