| @@ -13,9 +13,19 @@ import unittest | |||||
| _encoding = 'utf-8' | _encoding = 'utf-8' | ||||
| __all__ = [ 'validate' ] | |||||
| class Storage: | class Storage: | ||||
| '''A class to help read pieces of a torrent. | |||||
| ''' | |||||
| def __init__(self, rootpath, files, piecelen): | def __init__(self, rootpath, files, piecelen): | ||||
| ''' | |||||
| rootpath - path to the dir of torrent files are in | |||||
| files - the files dictionary from the torrent info key | |||||
| piecelen - piece length from the torren info key | |||||
| ''' | |||||
| self._rootpath = pathlib.Path(rootpath) | self._rootpath = pathlib.Path(rootpath) | ||||
| self._files = files | self._files = files | ||||
| self._piecelen = piecelen | self._piecelen = piecelen | ||||
| @@ -40,10 +50,16 @@ class Storage: | |||||
| yield curfile, fname, curfilepath | yield curfile, fname, curfilepath | ||||
| def allfiles(self): | def allfiles(self): | ||||
| '''Iterator that returns each on disk path name for | |||||
| each file.''' | |||||
| for x, y, curfilepath in self._filepaths(): | for x, y, curfilepath in self._filepaths(): | ||||
| yield curfilepath | yield curfilepath | ||||
| def _buildindex(self): | def _buildindex(self): | ||||
| '''Internal function to build the needed indexes for | |||||
| pieces and files.''' | |||||
| self._pieceindex = [] | self._pieceindex = [] | ||||
| self._fileindex = {} | self._fileindex = {} | ||||
| files = self._filepaths() | files = self._filepaths() | ||||
| @@ -75,19 +91,37 @@ class Storage: | |||||
| left -= sz | left -= sz | ||||
| def filepieces(self): | def filepieces(self): | ||||
| '''Iterator that returns a pair, first item is the subpath | |||||
| to a file (that is relative to the torrent dir), and the | |||||
| pieces that cover the file.''' | |||||
| return self._fileindex.items() | return self._fileindex.items() | ||||
| def filesforpiece(self, idx): | def filesforpiece(self, idx): | ||||
| '''Return a list of files that are covered by piece idx.''' | |||||
| for x in self._pieceindex[idx]: | for x in self._pieceindex[idx]: | ||||
| yield x['file'] | yield x['file'] | ||||
| def apply_piece(self, idx, fun): | def apply_piece(self, idx, fun): | ||||
| '''Read the parts of piece idx, and call fun w/ each part. | |||||
| This is to hash the parts, e.g. | |||||
| hash = sha1() | |||||
| stor.apply_piece(num, hash.update) | |||||
| hash now contains the digest for the piece.''' | |||||
| for i in self._pieceindex[idx]: | for i in self._pieceindex[idx]: | ||||
| with open(i['file'], 'rb') as fp: | with open(i['file'], 'rb') as fp: | ||||
| fp.seek(i['offset']) | fp.seek(i['offset']) | ||||
| fun(fp.read(i['size'])) | fun(fp.read(i['size'])) | ||||
| def validate(torrent, basedir): | def validate(torrent, basedir): | ||||
| '''Take a decode torrent file, where it was stored in basedir, | |||||
| verify the torrent. Returns a pair of set, the first is all the | |||||
| files that are valid, the second are all the invalid files.''' | |||||
| info = torrent['info'] | info = torrent['info'] | ||||
| basedir = pathlib.Path(basedir) | basedir = pathlib.Path(basedir) | ||||