storagewrapper.py

BitTorrent（简称BT

# The contents of this file are subject to the BitTorrent Open Source License
# Version 1.0 (the License).  You may not copy or use this file, in either
# source code or executable form, except in compliance with the License.  You
# may obtain a copy of the License at http://www.bittorrent.com/license/.
#
# Software distributed under the License is distributed on an AS IS basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
# for the specific language governing rights and limitations under the
# License.

# Written by Bram Cohen

from __future__ import division

from sha import sha
from array import array
from binascii import b2a_hex

from BitTorrent.bitfield import Bitfield
from BitTorrent import BTFailure, INFO, WARNING, ERROR, CRITICAL

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

def tobinary(i):
    return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
        chr((i >> 8) & 0xFF) + chr(i & 0xFF))

NO_PLACE = -1

ALLOCATED = -1
UNALLOCATED = -2
FASTRESUME_PARTIAL = -3

class StorageWrapper(object):

    def __init__(self, storage, config, hashes, piece_size, finished,
            statusfunc, flag, data_flunked, infohash, errorfunc, resumefile):
        self.numpieces = len(hashes)
        self.storage = storage
        self.config = config
        check_hashes = config['check_hashes']
        self.hashes = hashes
        self.piece_size = piece_size
        self.data_flunked = data_flunked
        self.errorfunc = errorfunc
        self.total_length = storage.get_total_length()
        self.amount_left = self.total_length
        self.partial_mark = "BitTorrent - this part has not been "+\
                            "downloaded yet."+infohash+\
                            tobinary(config['download_slice_size'])
        if self.total_length <= piece_size * (self.numpieces - 1):
            raise BTFailure, 'bad data in responsefile - total too small'
        if self.total_length > piece_size * self.numpieces:
            raise BTFailure, 'bad data in responsefile - total too big'
        self.finished = finished
        self.numactive = array('H', [0] * self.numpieces)
        self.inactive_requests = [1] * self.numpieces
        self.amount_inactive = self.total_length
        self.endgame = False
        self.have = Bitfield(self.numpieces)
        self.waschecked = Bitfield(self.numpieces)
        if self.numpieces < 32768:
            typecode = 'h'
        else:
            typecode = 'l'
        self.places = array(typecode, [NO_PLACE] * self.numpieces)
        if not check_hashes:
            self.rplaces = array(typecode, range(self.numpieces))
            fastresume = True
        else:
            self.rplaces = self._load_fastresume(resumefile, typecode)
            if self.rplaces is not None:
                fastresume = True
            else:
                self.rplaces = array(typecode, [UNALLOCATED] * self.numpieces)
                fastresume = False
        self.holepos = 0
        self.stat_numfound = 0
        self.stat_numflunked = 0
        self.stat_numdownloaded = 0
        self.stat_active = {}
        self.stat_new = {}
        self.stat_dirty = {}
        self.download_history = {}
        self.failed_pieces = {}

        if self.numpieces == 0:
            return
        targets = {}
        total = 0
        if not fastresume:
            for i in xrange(self.numpieces):
                if self._waspre(i):
                    self.rplaces[i] = ALLOCATED
                    total += 1
                else:
                    targets[hashes[i]] = i
        if total and check_hashes:
            statusfunc('checking existing file', 0)
        def markgot(piece, pos):
            if self.have[piece]:
                if piece != pos:
                    return
                self.rplaces[self.places[pos]] = ALLOCATED
                self.places[pos] = self.rplaces[pos] = pos
                return
            self.places[piece] = pos
            self.rplaces[pos] = piece
            self.have[piece] = True
            self.amount_left -= self._piecelen(piece)
            self.amount_inactive -= self._piecelen(piece)
            self.inactive_requests[piece] = None
            if not fastresume:
                self.waschecked[piece] = True
            self.stat_numfound += 1
        lastlen = self._piecelen(self.numpieces - 1)
        partials = {}
        for i in xrange(self.numpieces):
            if not self._waspre(i):
                if self.rplaces[i] != UNALLOCATED:
                    raise BTFailure("--check_hashes 0 or fastresume info "
                                    "doesn't match file state (missing data)")
                continue
            elif fastresume:
                t = self.rplaces[i]
                if t >= 0:
                    markgot(t, i)
                    continue
                if t == UNALLOCATED:
                    raise BTFailure("Bad fastresume info (files contain more "
                                    "data)")
                if t == ALLOCATED:
                    continue
                if t!= FASTRESUME_PARTIAL:
                    raise BTFailure("Bad fastresume info (illegal value)")
                data = self.storage.read(self.piece_size * i,
                                         self._piecelen(i))
                self._check_partial(i, partials, data)
                self.rplaces[i] = ALLOCATED
            else:
                data = self.storage.read(piece_size * i, self._piecelen(i))
                sh = sha(buffer(data, 0, lastlen))
                sp = sh.digest()
                sh.update(buffer(data, lastlen))
                s = sh.digest()
                if s == hashes[i]:
                    markgot(i, i)
                elif s in targets and self._piecelen(i) == self._piecelen(targets[s]):
                    markgot(targets[s], i)
                elif not self.have[self.numpieces - 1] and sp == hashes[-1] and (i == self.numpieces - 1 or not self._waspre(self.numpieces - 1)):
                    markgot(self.numpieces - 1, i)
                else:
                    self._check_partial(i, partials, data)
                statusfunc(fractionDone = 1 - self.amount_left /
                           self.total_length)
            if flag.isSet():
                return
        self.amount_left_with_partials = self.amount_left
        for piece in partials:
            if self.places[piece] < 0:
                pos = partials[piece][0]
                self.places[piece] = pos
                self.rplaces[pos] = piece
                self._make_partial(piece, partials[piece][1])
        for i in xrange(self.numpieces):
            if self.rplaces[i] != UNALLOCATED:
                self.storage.allocated(piece_size * i, self._piecelen(i))
            if self.have[i]:
                self.storage.downloaded(piece_size * i, self._piecelen(i))

    def _waspre(self, piece):
        return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))

    def _piecelen(self, piece):
        if piece < self.numpieces - 1:
            return self.piece_size
        else:
            return self.total_length - piece * self.piece_size

    def _check_partial(self, pos, partials, data):
        index = None
        missing = False
        marklen = len(self.partial_mark)+4
        for i in xrange(0, len(data) - marklen,
                        self.config['download_slice_size']):
            if data[i:i+marklen-4] == self.partial_mark:
                ind = toint(data[i+marklen-4:i+marklen])
                if index is None:
                    index = ind
                    parts = []
                if ind >= self.numpieces or ind != index:
                    return
                parts.append(i)
            else:
                missing = True
        if index is not None and missing:
            i += self.config['download_slice_size']
            if i < len(data):
                parts.append(i)
            partials[index] = (pos, parts)

    def _make_partial(self, index, parts):
        length = self._piecelen(index)
        l = []