Tiny .torrent file creator.

torr.py

#!/usr/bin/env python2
'''
A tiny .torrent file creator with no external dependencies.

Easy to hack!
'''

import argparse
import hashlib
import os
from cStringIO import StringIO


def bencode(v):
    # https://en.wikipedia.org/wiki/Bencode
    buf = StringIO()
    def emit(v):
        if isinstance(v, (int, long, bool)):
            buf.write('i%de' % v)
        elif isinstance(v, str):
            buf.write('%d:%s' % (len(v), v))
        elif isinstance(v, dict):
            buf.write('d')
            for k, v in sorted(v.items()):
                emit(k)
                emit(v)
            buf.write('e')
        elif isinstance(v, list):
            buf.write('l')
            for e in v:
                emit(e)
            buf.write('e')
        else:
            raise TypeError('unhandled type: %s in %r' % (type(v), v))
    emit(v)
    return buf.getvalue()


class Torrent(object):
    def __init__(self, base_dir=None, piece_length=2**(8*10)):
        self.base_dir = base_dir
        self.piece_length = piece_length
        self.pieces = ''
        self.hasher = hashlib.sha1()
        self.hasher_fill = 0
        self.files = []
        self.file_set = set()

    def add_file(self, path, torrent_path=None):
        if torrent_path is None:
            torrent_path = path
            if self.base_dir:
                torrent_path = os.path.relpath(path, self.base_dir)
        assert torrent_path not in self.file_set, 'Duplicate file: %s' % torrent_path
        self.file_set.add(torrent_path)
        self.files.append({'path': torrent_path.split('/'), 'length': os.stat(path).st_size})
        with open(path, 'rb') as f:
            while True:
                chunk = f.read(min(self.piece_length - self.hasher_fill, self.piece_length))
                if not chunk:
                    break
                self.hasher.update(chunk)
                self.hasher_fill += len(chunk)
                assert self.hasher_fill <= self.piece_length
                if self.hasher_fill == self.piece_length:
                    self.pieces += self.hasher.digest()
                    self.hasher = hashlib.sha1()
                    self.hasher_fill = 0

    def get_encoded(self, name, **kwargs):
        if self.hasher_fill:
            self.pieces += self.hasher.digest()
            self.hasher = None
        d = {
            'info': {
                'name': name,
                'piece length': self.piece_length,
                'pieces': self.pieces,
            }
        }
        if len(self.files) > 1:
            d['info']['files'] = self.files
        elif self.files:
            d['info']['length'] = self.files[0]['length']
        else:
            raise ValueError('cannot create empty torrent file')
        d.update((k,v) for k,v in kwargs.iteritems() if v)
        return bencode(d)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-a', '--announce', help='Tracker announce URL.')
    parser.add_argument('-p', '--private', action='store_true', help='Disable peer exchange and DHT.')
    parser.add_argument('-s', '--piece_length', type=int, default=18, choices=range(10, 31),
                        help='Power of 2 chunk size (default: 18=256KiB).')
    parser.add_argument('-b', '--base', help='Base directory to make paths relative to.')
    parser.add_argument('-n', '--name', help='Suggested name of the torrent. Defaults to base directory name or first file directory name.')
    parser.add_argument('-o', '--output', metavar='OUT.TORRENT', help='Output .torrent name.')
    parser.add_argument('files', nargs='+', help='Files to include in torrent.')
    options = parser.parse_args()

    piece_length = 2 ** options.piece_length
    if piece_length < 2**20:
        print 'Using %d KiB chunks.' % (piece_length / 1024)
    else:
        print 'Using %d MiB chunks.' % (piece_length / 1024 ** 2)

    torr = Torrent(base_dir=options.base, piece_length=piece_length)

    for fname in options.files:
        assert os.path.isfile(fname)
        torr.add_file(fname)
        print 'Added', fname

    name = options.files[0]
    if options.name:
        name = options.name
    elif options.base:
        name = os.path.basename(options.base)
    elif len(options.files) > 1:
        name = os.path.basename(os.path.dirname(os.path.abspath(options.files[0])))

    encoded = torr.get_encoded(name=name, private=options.private, announce=options.announce)

    with open(options.output, 'w') as f:
        f.write(encoded)

    print 'Done.'

if __name__ == '__main__':
    main()