sicksubs.py

#!/usr/bin/env python
import logging
import datetime
import glob
import db
import os
import sys
import time
import shlex
import sqlite3
import periscope
import subprocess

#***********************************<CONFIG>***********************************
# These are parameters that you may want to configure, although the defaults
# are quite sane

# location where sicksubs should store the queue file
DATABASE_FILE = u'~/.sicksubs/sicksubs.db'

# the language of the downloaded subs, can be nl or en
SUB_LANG = 'en'

# post-processing script(s), will be called with one argument on successful
# sub download, the name of the ep for which a sub was found
#
# multiple scripts can be separated by comma's. do _not_ use unicode strings
# since the shlex module does not support unicode prior to 2.7.3
POST_CALL = ''  # '/home/peter/test.sh,/home/peter/test2.sh'

# if you want the subtitle file to include the language in the file name, set
# this option to True. This would result in a subtitle of the form:
#
# White.Collar.S02E04.DVDRip.XviD-SAiNTS.nl.srt
APPEND_LANG = False

# database entries are cleaned automatically after this number of days of not
# finding anything
CLEAN_AFTER = 30
#***********************************</CONFIG>*********************************


quiet = False


def sickbeard_run(conn):
    """
    This function will be called when the script is executed by sickbeard. This
    will add a final_location to the correct item in the queue, to make sure
    the subtitle file can be moved there after downloading.
    """
    # It passes 6 parameters to these scripts:
    # 1 final full path to the episode file
    # 2 original name of the episode file
    # 3 show tvdb id
    # 4 season number
    # 5 episode number
    # 6 episode air date
    # example call:
    # ['/home/sickbeard/sicksubs/sicksubs.py',
    # u'/media/media/Series/Qi/Season 09/QI.S09E12.Illumination.avi',
    # u'/media/bin2/usenet_downloads/tv/QI.S09E12.HDTV.XviD-FTP/qi.s09e12.
    # hdtv.xvid-ftp.avi',
    # '72716', '9', '12', '2011-11-25']
    final_loc = os.path.expanduser(os.path.abspath(sys.argv[1]))

    ep = db.add_ep(conn, final_loc, datetime.date.today())
    return find_subs([ep])


def cron_run(conn):
    """
    This function will be called when the script is executed by cron. This will
    read the jobs and try to find sub downloads for each of them
    """
    # get all eps
    all_eps = db.get_all_eps(conn)
    return find_subs(all_eps)


def find_subs(eps):
    """
    This function searches for available subtitle files in each of the Ep instances
    in eps, using the periscope subtitle download library.
    """
    to_download = {}
    try:
        import xdg.BaseDirectory as bd
        cache_folder = os.path.join(bd.xdg_config_home, "periscope")
    except ImportError:
        cache_folder = os.path.join(os.path.expanduser("~"), ".config", "periscope")

    subdl = periscope.Periscope(cache_folder)
    for ep in eps:
        ep_name = os.path.splitext(os.path.expanduser(ep.final_loc))[0]
        logging.info('Searching subs for {0}'.format(ep_name))
        if not os.path.exists(ep.final_loc):
            db.remove_single(conn, ep)
            logging.info(u'Cleaned up db because {0} is no longer available on disk!'.format(ep_name))
            continue

        if glob.glob(ep_name + '*.srt'):
            # Maybe user downloaded sub for this ep manually?
            db.remove_single(conn, ep)
            logging.info(u'Cleaned up db because {0} already has subs!'.format(ep_name))
            continue

        if (datetime.date.today() - ep.added_on).days > CLEAN_AFTER:
            # Clean this ep as it's too old
            db.remove_single(conn, ep)
            logging.info(u'Cleaned up db because {0} is older than {1} days, and there are still no subs!'.format(
                ep_name, CLEAN_AFTER))
            continue

        subs = subdl.listSubtitles(ep.final_loc, [SUB_LANG])
        if subs:
            to_download[ep] = subs
        time.sleep(3)

    if not to_download:
        logging.info("No subs available for any of your eps yet!")
        return

    successful = []
    for d in to_download:
        if subdl.attemptDownloadSubtitle(to_download[d], [SUB_LANG]) is not None:
            successful.append(d)

    # append languages
    if APPEND_LANG:
        for s in successful:
            base = os.path.splitext(s.final_loc)[0]
            old_sub = base + '.srt'
            if os.path.isfile(old_sub):
                os.rename(old_sub, '.'.join((base, SUB_LANG, 'srt')))

    # remove successfully downloaded files from db
    db.remove_downloaded(conn, successful)
    # call post-processing for successfully downloaded files
    if POST_CALL:
        for d in successful:
            for script in POST_CALL.split(','):
                to_call = shlex.split(script)
                to_call.append(d.final_loc)
                subprocess.call(to_call)


if __name__ == '__main__':
    if '-q' in sys.argv:
        quiet = True
        sys.argv.remove('-q')

    if quiet:
        logging.basicConfig(level=logging.ERROR)
    else:
        logging.basicConfig(level=logging.INFO)

    conn = db.initialize(DATABASE_FILE)

    # two or seven arguments, two included for debugging purposes
    if len(sys.argv) in (2, 7):
        sickbeard_run(conn)
    else:
        cron_run(conn)