maketorrent.py

# -*- coding: utf-8 -*-
#
# Copyright (C) 2009 Andrew Resch <andrewresch@gmail.com>
#
# This file is part of Deluge and is licensed under GNU General Public License 3.0, or later, with
# the additional special exception to link portions of this program with the OpenSSL library.
# See LICENSE for more details.
#

from __future__ import division, unicode_literals, print_function

import os
import sys
from hashlib import sha1 as sha
import logging

from bencode import bencode

def removeit(p):
	while p[0] == "\\" or p[0] == "/":
		p = p[1:]
	return p

def get_path_size(path):
	"""
	Gets the size in bytes of 'path'

	:param path: the path to check for size
	:type path: string
	:returns: the size in bytes of the path or -1 if the path does not exist
	:rtype: int

	"""
	if not os.path.exists(path):
		return -1

	if os.path.isfile(path):
		return os.path.getsize(path)

	dir_size = 0
	for (p, dummy_dirs, files) in os.walk(path):
		for _file in files:
			filename = os.path.join(p, _file)
			dir_size += os.path.getsize(filename)
	return dir_size

def utf8_encode_structure(data):
	"""Recursively convert all unicode keys and values in a data structure to utf8.

	e.g. converting keys and values for a dict with nested dicts and lists etc.

	Args:
		data (any): This can be any structure, dict, list or tuple.

	Returns:
		input type: The data with unicode keys and values converted to utf8.

	"""
	if isinstance(data, (list, tuple)):
		return type(data)([utf8_encode_structure(d) for d in data])
	elif isinstance(data, dict):
		return dict([utf8_encode_structure(d) for d in data.items()])
	elif not isinstance(data, bytes):
		try:
			return data.encode('utf8')
		except AttributeError:
			pass
	return data


class InvalidPath(Exception):
	"""Raised when an invalid path is supplied."""
	pass


class InvalidPieceSize(Exception):
	"""Raised when an invalid piece size is set.

	Note:
		Piece sizes must be multiples of 16KiB.
	"""
	pass


class TorrentMetadata(object):
	"""This class is used to create .torrent files.

	Examples:

		>>> t = TorrentMetadata()
		>>> t.data_path = '/tmp/torrent'
		>>> t.comment = 'My Test Torrent'
		>>> t.trackers = [['http://tracker.openbittorent.com']]
		>>> t.save('/tmp/test.torrent')

	"""
	def __init__(self):
		self.__data_path = None
		self.__piece_size = 0
		self.__comment = ''
		self.__private = False
		self.__trackers = []
		self.__webseeds = []
		self.__pad_files = False
		self.__source = ''
		self.skip = [
			"Thumbs.db",
			"edition.txt",
			"tracklist.txt",
			"formatted_tracklist.txt",
			"credits.txt",
		]

	def save(self, torrent_path, progress=None):
		"""Creates and saves the torrent file to `path`.

		Args:
			torrent_path (str): Location to save the torrent file.
			progress(func, optional): The function to be called when a piece is hashed. The
				provided function should be in the format `func(num_completed, num_pieces)`.

		Raises:
			InvalidPath: If the data_path has not been set.

		"""
		if not self.data_path:
			raise InvalidPath('Need to set a data_path!')

		torrent = {
			'info': {}
		}

		if self.source:
			torrent['info']['source'] = self.source

		if self.comment:
			torrent['comment'] = self.comment

		if self.private:
			torrent['info']['private'] = True

		if self.trackers:
			torrent['announce'] = self.trackers[0][0]
			#torrent['announce-list'] = self.trackers
		else:
			torrent['announce'] = ''

		if self.webseeds:
			httpseeds = []
			webseeds = []
			for w in self.webseeds:
				if w.endswith('.php'):
					httpseeds.append(w)
				else:
					webseeds.append(w)

			if httpseeds:
				torrent['httpseeds'] = httpseeds
			if webseeds:
				torrent['url-list'] = webseeds

		datasize = get_path_size(self.data_path)

		if self.piece_size:
			piece_size = self.piece_size * 1024
		else:
			# We need to calculate a piece size
			piece_size = 16384
			while (datasize // piece_size) > 1024 and piece_size < (8192 * 1024):
				piece_size *= 2

		# Calculate the number of pieces we will require for the data
		num_pieces = datasize // piece_size
		if datasize % piece_size:
			num_pieces += 1

		torrent['info']['piece length'] = piece_size
		torrent['info']['name'] = os.path.split(self.data_path)[1]

		# Create the info
		if os.path.isdir(self.data_path):
			files = []
			padding_count = 0
			# Collect a list of file paths and add padding files if necessary
			for (dirpath, dirnames, filenames) in os.walk(self.data_path):
				for index, filename in enumerate(filenames):
					if filename in self.skip or filename.endswith(".torrent") or filename.endswith(".json") or filename.endswith(".ini") or filename.endswith(".ico"):
						logging.info("Skipped " + filename)
					else:
						size = get_path_size(os.path.join(self.data_path, dirpath, filename))
						p = dirpath[len(self.data_path):]
						p = p.lstrip('/')
						p = p.split('/')
						if p[0]:
							p += [filename]
						else:
							p = [filename]
						files.append((size, p))
						# Add a padding file if necessary
						if self.pad_files and (index + 1) < len(filenames):
							left = size % piece_size
							if left:
								p = list(p)
								p[-1] = '_____padding_file_' + str(padding_count)
								files.append((piece_size - left, p))
								padding_count += 1

			# Run the progress function with 0 completed pieces
			if progress:
				progress(0, num_pieces)

			fs = []
			pieces = []
			# Create the piece hashes
			buf = b''
			for size, path in files:
				logging.info(path)
				path = [removeit(s) for s in path]
				#path = [removeit(s.decode(sys.getfilesystemencoding()).encode('UTF-8')) for s in path]
				fs.append({'length': size, 'path': path})
				#if path[-1].startswith('_____padding_file_'):
				#	buf += '\0' * size
				#	pieces.append(sha(buf).digest())
				#	buf = b''
				#	fs[-1]['attr'] = 'p'
				#else:
				#print (self.data_path, path, "==>", os.path.join(self.data_path, *path))
				with open(os.path.join(self.data_path, *path), 'rb') as _file:
					r = _file.read(piece_size - len(buf))
					while r:
						buf = buf + r
						if len(buf) == piece_size:
							pieces.append(sha(buf).digest())
							# Run the progress function if necessary
							if progress:
								progress(len(pieces), num_pieces)
							buf = b''
						else:
							break
						r = _file.read(piece_size - len(buf))
			torrent['info']['files'] = fs
			if buf:
				pieces.append(sha(buf).digest())
				if progress:
					progress(len(pieces), num_pieces)
				buf = ''

		elif os.path.isfile(self.data_path):
			torrent['info']['length'] = get_path_size(self.data_path)
			pieces = []

			with open(self.data_path, 'rb') as _file:
				r = _file.read(piece_size)
				while r:
					pieces.append(sha(r).digest())
					if progress:
						progress(len(pieces), num_pieces)

					r = _file.read(piece_size)

		torrent['info']['pieces'] = b''.join(pieces)

		# Write out the torrent file
		with open(torrent_path, 'wb') as _file:
			_file.write(bencode(utf8_encode_structure(torrent)))

		#f = open(os.path.abspath(torrent_path), 'wb')
		#f.write(bencode(utf8_encode_structure(torrent)))
		#f.close()

	def get_data_path(self):
		"""Get the path to the files that the torrent will contain.

		Note:
			It can be either a file or a folder.

		Returns:
			str: The torrent data path, either a file or a folder.

		"""
		return self.__data_path

	def set_data_path(self, path):
		"""Set the path to the files (data) that the torrent will contain.

		Note:
			This property needs to be set before the torrent file can be created and saved.

		Args:
			path (str): The path to the torrent data and can be either a file or a folder.

		Raises:
			InvalidPath: If the path is not found.

		"""
		if os.path.exists(path) and (os.path.isdir(path) or os.path.isfile(path)):
			self.__data_path = os.path.abspath(path)
		else:
			raise InvalidPath('No such file or directory: %s' % path)

	def get_piece_size(self):
		"""The size of the pieces.

		Returns:
			int: The piece size in multiples of 16 KiBs.
		"""
		return self.__piece_size

	def set_piece_size(self, size):
		"""Set piece size.

		Note:
			If no piece size is set, one will be automatically selected to
			produce a torrent with less than 1024 pieces or the smallest possible
			with a 8192KiB piece size.

		Args:
			size (int): The desired piece size in multiples of 16 KiBs.

		Raises:
			InvalidPieceSize: If the piece size is not a valid multiple of 16 KiB.

		"""
		if size % 16 and size:
			raise InvalidPieceSize('Piece size must be a multiple of 16 KiB')
		self.__piece_size = size

	def get_comment(self):
		"""Get the torrent comment.

		Returns:
			str: An informational string about the torrent.

		"""
		return self.__comment

	def set_comment(self, comment):
		"""Set the comment for the torrent.

		Args:
			comment (str): An informational string about the torrent.

		"""
		self.__comment = comment

	def get_source(self):
		"""Get the torrent source.

		Returns:
			str: The source of the torrent.

		"""
		return self.__source

	def set_source(self, source):
		"""Set the source for the torrent.

		Args:
			source (str): The source of the torrent.

		"""
		self.__source = source

	def get_private(self):
		"""Get the private flag of the torrent.

		Returns:
			bool: True if private flag has been set, else False.

		"""
		return self.__private

	def set_private(self, private):
		"""Set the torrent private flag.

		Note:
			Private torrents only announce to trackers and will not use DHT or
			Peer Exchange. See http://bittorrent.org/beps/bep_0027.html

		Args:
			private (bool): True if the torrent is to be private.

		"""
		self.__private = private

	def get_trackers(self):
		"""Get the announce trackers.

		Note:
			See http://bittorrent.org/beps/bep_0012.html

		Returns:
			list of lists: A list containing a list of trackers.

		"""
		return self.__trackers

	def set_trackers(self, trackers):
		"""Set the announce trackers.

		Args:
			private (list of lists): A list containing lists of trackers as strings, each list is a tier.

		"""
		self.__trackers = trackers

	def get_webseeds(self):
		"""Get the webseeds.

		Note:
			The web seeds can either be:
				Hoffman-style: http://bittorrent.org/beps/bep_0017.html
				GetRight-style: http://bittorrent.org/beps/bep_0019.html

			If the url ends in '.php' then it will be considered Hoffman-style, if
			not it will be considered GetRight-style.

		Returns:
			list: The webseeds.

		"""
		return self.__webseeds

	def set_webseeds(self, webseeds):
		"""Set webseeds.

		Note:
			The web seeds can either be:
				Hoffman-style: http://bittorrent.org/beps/bep_0017.html
				GetRight-style: http://bittorrent.org/beps/bep_0019.html

			If the url ends in '.php' then it will be considered Hoffman-style, if
			not it will be considered GetRight-style.

		Args:
			private (list): The webseeds URLs which can be either Hoffman or GetRight style.

		"""
		self.__webseeds = webseeds

	def get_pad_files(self):
		"""Get status of padding files for the torrent.

		Returns:
			bool: True if padding files have been enabled to align files on piece boundaries.

		"""
		return self.__pad_files

	def set_pad_files(self, pad):
		"""Enable padding files for the torrent.

		Args:
			private (bool): True adds padding files to align files on piece boundaries.

		"""
		self.__pad_files = pad

	data_path = property(get_data_path, set_data_path)
	piece_size = property(get_piece_size, set_piece_size)
	comment = property(get_comment, set_comment)
	source = property(get_source, set_source)
	private = property(get_private, set_private)
	trackers = property(get_trackers, set_trackers)
	webseeds = property(get_webseeds, set_webseeds)
	pad_files = property(get_pad_files, set_pad_files)