medashare/ui/medashare/btv/__init__.py

from . import bencode
from functools import reduce
from hashlib import sha1
import importlib.resources
import itertools
import os
import pathlib
import shutil
import sys
import tempfile
import unittest

_encoding = 'utf-8'

__all__ = [ 'validate', 'validate_file' ]

class Storage:
	'''A class to help read pieces of a torrent.
	'''

	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._files = files
		self._piecelen = piecelen

		self._buildindex()

	def _filepaths(self):
		'''Iterates over all the files in the torrent.

		Each item is a tuple of:
			array of file path components (undecoded)
			a pathlib.PurePath for the file
			a pathlib.Path for file on disk
		'''

		for curfile in self._files:
			fname = pathlib.PurePath(
			    *(x.decode(_encoding) for x in
			    curfile['path']))
			curfilepath = self._rootpath / fname

			yield curfile, fname, curfilepath

	def allfiles(self):
		'''Iterator that returns each on disk path name for
		each file.'''

		for x, y, curfilepath in self._filepaths():
			yield curfilepath

	def _buildindex(self):
		'''Internal function to build the needed indexes for
		pieces and files.'''

		self._pieceindex = []
		self._fileindex = {}
		files = self._filepaths()
		left = 0
		curfile = None

		while True:
			if curfile is None or curfileoff == curfile['length']:
				# next file
				try:
					curfile, fname, curfilepath = next(files)
				except StopIteration:
					break
				curfileoff = 0

			if left == 0:
				current = []
				self._fileindex.setdefault(fname,
				    []).append(len(self._pieceindex))
				self._pieceindex.append(current)
				left = self._piecelen

			sz = min(curfile['length'] - curfileoff, left)

			current.append(dict(file=curfilepath, fname=fname,
			    offset=curfileoff, size=sz))

			curfileoff += sz
			left -= sz

	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()

	def filesforpiece(self, idx):
		'''Return a list of files that are covered by piece idx.'''

		for x in self._pieceindex[idx]:
			yield x['file']

	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]:
			with open(i['file'], 'rb') as fp:
				fp.seek(i['offset'])
				fun(fp.read(i['size']))

def validate_file(fname):
	fname = pathlib.Path(fname)

	with open(fname, 'rb') as fp:
		torrent = bencode.bdecode(fp.read())

	dirname = list(fname.parent.rglob(torrent['info']['name'].decode(_encoding)))[0]

	tordir = dirname.parent

	return validate(torrent, tordir)

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']

	basedir = pathlib.Path(basedir)

	torrentdir = basedir / info['name'].decode(_encoding)

	stor = Storage(torrentdir, info['files'], info['piece length'])

	pieces = info['pieces']
	piecescnt = len(pieces) // 20
	valid = [ None ] * piecescnt
	for num, i in enumerate(pieces[x:x+20] for x in range(0, len(pieces),
	    20)):
		hash = sha1()

		stor.apply_piece(num, hash.update)

		if hash.digest() == i:
			valid[num] = True
		else:
			valid[num] = False

	# if any piece of a file is bad, it's bad
	allfiles = set(stor.allfiles())

	badfiles = { torrentdir / x for x, y in stor.filepieces() if
	    not all(valid[i] for i in y) }

	return allfiles - badfiles, badfiles

class _TestCases(unittest.TestCase):
	dirname = 'somedir'

	# file contents for somedir.torrent
	origfiledata = {
		'filea.txt':	b'foo\n',
		'fileb.txt':	b'bar\n',
		'filec.txt':	b'bleha\n',
		'filed.txt':	b'somehow\n',
		'filee.txt':	b'nowab\n',
		'filef/filef.txt':	b'\n',
	}

	# some munging to make some files bad
	badfiles = {
		'filea.txt': b'',
		'filec.txt': b'\x00\x00\x00\x00a\n',
		'filee.txt': b'no',
	}

	def setUp(self):
		d = pathlib.Path(tempfile.mkdtemp()).resolve()

		self.basetempdir = d

		tor = importlib.resources.files(__name__)
		tor = tor / 'fixtures' / 'somedir.torrent'
		with tor.open('rb') as fp:
			self.torrent = bencode.bdecode(fp.read())

		self.oldcwd = os.getcwd()

		os.chdir(d)

	def tearDown(self):
		shutil.rmtree(self.basetempdir)

		os.chdir(self.oldcwd)

	@staticmethod
	def make_files(dname, fdict):
		dname = pathlib.Path(dname)
		for k, v in fdict.items():
			k = dname / pathlib.PurePosixPath(k)
			k.parent.mkdir(parents=True, exist_ok=True)
			with open(k, 'wb') as fp:
				fp.write(v)

	def test_completeverif(self):
		tf = self.basetempdir / 'a.torrent'
		with open(tf, 'wb') as fp:
			fp.write(bencode.bencode(self.torrent))

		sd = self.basetempdir / 'anotherdir' / self.dirname
		sd.parent.mkdir()
		sd.mkdir()

		self.make_files(sd, self.origfiledata)

		good, bad = validate_file(tf)

		self.assertFalse(bad)

		# that utf-8 encoded names work

		sd = self.basetempdir / 'thai'
		sd.mkdir()

		self.make_files(sd, { 'thai - สวัสดี.txt': b'hello\n'
		})

		tor = importlib.resources.files(__name__)
		tor = tor / 'fixtures' / 'thai.torrent'
		with tor.open('rb') as fp:
			torrent = bencode.bdecode(fp.read())

		good, bad = validate(torrent, self.basetempdir)

		self.assertFalse(bad)

	def test_verification(self):
		# Testing for "missing" files
		# piece size 2 (aka 4 bytes)
		# empty file of 4 bytes		'foo\n'
		# complete file of 4 bytes	'bar\n'
		# partial missing file, 6 bytes, last two correct	'bleha\n'
		# complete file of 8 bytes (multiple pieces)		'somehow\n'
		# partial missing file, starting w/ 2 bytes, length 6	'nowab\n'
		# complete file (length 1)				'\n'

		missingfiles = self.origfiledata.copy()

		missingfiles.update(self.badfiles)

		sd = self.basetempdir / self.dirname
		sd.mkdir()

		self.make_files(sd, missingfiles)

		val, inval = validate(self.torrent, self.basetempdir)

		self.assertEqual(set(val), { sd / x for x in
		    missingfiles.keys() if x not in self.badfiles })
		self.assertEqual(set(inval), { sd / x for x in
		    self.badfiles.keys() })