lostfilm download script

lostfilm.py

#!/usr/bin/env python2
# -*- encoding: utf-8 -*-

import re
import os
import sys
import json
import argparse
import yaml
from logging import getLogger
from requests import session
from lxml import etree
from itertools import ifilter, imap, starmap
from urlparse import urljoin

logger = getLogger(__name__)

def arguments(*argspecs):
    def decorator(func):
        setattr(func, '__argspecs__', argspecs)
        return func
    return decorator

def get_command_args(command_func, args):
    argspecs = getattr(command_func, '__argspecs__', None)

    if argspecs:
        parser = argparse.ArgumentParser()
        for a, kw in argspecs:
            parser.add_argument(*a, **kw)
        return vars(parser.parse_args(args))
    else:
        return {}


def title_filter(*substrings):
    return lambda t: any(map(t.__contains__, substrings))

def load_config(filename='/etc/lostfilm.yml'):
    try:
        with open(filename, 'r') as f:
            return yaml.safe_load(f)
    except IOError:
        return {}

torrents_dir = '~/torrents'
transmission_url = 'http://localhost:9091/transmission/rpc'

BASE_URL = 'http://www.lostfilm.tv/'
LOGIN_URL = 'http://login1.bogi.ru/login.php'

def tap(o):
    #print(o, file=sys.stderr)
    print >> sys.stderr, unicode(o).encode('utf-8')
    #logger.info(o)
    return o


def login(email, password):

    s = session()
    s.headers.update({
        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/33.0.1750.152 Safari/537.36',
        })

    login_form = etree.HTML(
            s.post(LOGIN_URL,
                params={'referer': BASE_URL},
                data={
                    'login': email,
                    'password': password,
                    'module': '1',
                    'target': BASE_URL,
                    'repage': 'user',
                    'act': 'login'},
                headers={'Referer': BASE_URL}
                ).content
            ).find('body/form')

    action = login_form.get('action')
    inputs = dict(
            (i.get('name'), i.get('value'))
            for i in login_form.xpath('input')
            if i.get('value') is not None)

    response = s.post(action, data=inputs)

    return s


def get_torrent_urls(s, include=lambda t: True, exclude=lambda t: False):
    cookie_re = re.compile(r"setCookie\('(\w+)','([a-f0-9]+)'\)")
    params_re = re.compile(r"ShowAllReleases\('(\d+)','([\d.]+)','(\d+)'\)")
    torrent_link_re = re.compile(r'href="(http://tracktor\.in/td\.php\?s=[^"]+)"')

    charset = 'cp1251'

    return starmap(lambda c, p: s.cookies.set('%s_2' % c[0], c[1]) and
            next(torrent_link_re.finditer(s.get(urljoin(BASE_URL, 'nrdr.php?c=%s&s=%s&e=%s' % p)).content.decode(charset))).group(1),
                imap(lambda a: (next(cookie_re.finditer(a.get('onmouseover'))).groups(), next(params_re.finditer(a.get('onclick'))).groups()),
                    imap(lambda url: etree.HTML(s.get(url).content).xpath('//div[@class="content_body"]/div/a[@class="a_download"]')[0].attrib,
                        starmap(lambda title, url: tap(title) and url.replace('download.', 'details.').rsplit('&', 1)[0],
                            ifilter(lambda p: include(p[0]) and not exclude(p[0]),
                                imap(lambda item: (item.find('title').text, item.find('link').text),
                                    etree.XML(s.get(urljoin(BASE_URL, 'rssdd.xml')).content).xpath('channel/item')))))))

@arguments(
    (['--torrents-dir'], dict(help='directory to download torrents to', default='~/torrents', dest='torrents_dir'))
    )
def download_torrents(torrent_urls, torrents_dir=torrents_dir):
    '''
    Download torrent files
    '''
    torrents_dir = os.path.expanduser(torrents_dir)    

    try:
        os.makedirs(torrents_dir)
    except os.error as e:
        if e.errno != os.errno.EEXIST:
            raise

    s = session()

    for filename, content in imap(lambda resp: (os.path.join(torrents_dir, resp.headers['content-disposition'].rsplit('=', 1)[1].strip('"')), resp.content),
            imap(lambda url: s.get(url), torrent_urls)):
        with open(filename, 'wb') as f:
            f.write(content)

@arguments(
        (['--transmission-url'], dict(help='transmission RPC URL',
            default='http://localhost:9091/transmission/rpc', dest='transmission_url')),
        (['--download-dir'], dict(help='directory for Transmission to download to', dest='download_dir', nargs='?')),
        (['--paused'], dict(help='pause torrent after adding', action='store_true'))
        )
def add_torrents_to_transmission(torrent_urls, transmission_url=transmission_url, download_dir=None, paused=False):
    '''
    Enqueue torrents to Transmission
    '''
    s = session()
    token = ''
    args = {
            'paused': paused,
            'download-dir': download_dir
            }

    for tag, url in enumerate(torrent_urls):
        for _ in xrange(2):
            args['filename'] = url
            response = s.post(transmission_url, data=json.dumps({'tag': tag, 'method': 'torrent-add', 'arguments': args}),
                headers={'X-Transmission-Session-Id': token})

            if response.status_code == 200:
                break

            elif response.status_code == 409:
                token = response.headers['X-Transmission-Session-Id']

            else:
                raise RuntimeError('unexpected status code %s for torrent %s' % (response.status_code, url))

def list_torrents(torrent_urls):
    '''
    Print list of found torrent URLs.
    '''
    print('\n'.join(torrent_urls))


commands = dict(
    download=download_torrents,
    enqueue=add_torrents_to_transmission,
    list=list_torrents)

if __name__ == '__main__':
    config = load_config()
    
    parser = argparse.ArgumentParser(description='LostFilm torrents download tool')
    parser.add_argument('command')
    parser.add_argument('--email', '--username', help='LostFilm user email', type=str, default=config['username'])
    parser.add_argument('--password', help='LostFilm user password', type=str, default=config['password'])
    parser.add_argument('options', nargs=argparse.REMAINDER)

    args = parser.parse_args()

    command = commands[args.command]
    options = get_command_args(command, args.options)

    include = title_filter(*config['include'])
    exclude = title_filter(*config['exclude'])
    torrent_urls = get_torrent_urls(login(args.email, args.password), include, exclude)
    command(torrent_urls, **options)

lostfilm.yml

include:
    - (Grimm)
    - (Continuum)
    - (The Tomorrow People)
    - (Almost Human)
    - (Arrow)
    - (Castle)
    - (Elementary)
    - (Falling Skies)
    - (Marvel's Agents of S.H.I.E.L.D.)
    - (Warehouse 13)

exclude:
    - '[MP4]'
    - '[1080p]'

username: login
password: 12345678