subdl

#!/usr/bin/python

# subdl - command-line tool to download subtitles from opensubtitles.org.
#
# Uses code from subdownloader (a GUI app).

__doc__ = '''\
Syntax: subdl [options] moviefile.avi

Subdl is a command-line tool for downloading subtitles from opensubtitles.org.

By default, it will search for English subtitles, display the results,
download the highest-rated result in the requested language and save it to the
appropriate filename.

Options:
  --help               This text
  --version            Print version and exit
  --lang=LANGUAGES     Comma-separated list of languages in 3-letter code, e.g.
                       'eng,spa,fre', or 'all' for all.  Default is 'eng'.
  --list-languages     List available languages and exit.
  --download=ID        Download a particular subtitle by numeric ID.
  --download=first     Download the first search result [default].
  --download=all       Download all search results.
  --download=query     Query which search result to download.
  --download=none, -n  Display search results and exit.
  --output=OUTPUT      Output to specified output filename.  Can include the
                       following format specifiers:
                         %I subtitle id
                         %m movie file base      %M movie file extension
                         %s subtitle file base   %S subtitle file extension
                         %l language (English)   %L language (2-letter ISO639)
                       Default is "%m.%S"; if multiple languages are searched,
                       then the default is "%m.%L.%S"; if --download=all, then
                       the default is "%m.%L.%I.%S".
  --existing=abort     Abort if output filename already exists [default].
  --existing=bypass    Exit gracefully if output filename already exists.
  --existing=overwrite Overwrite if output filename already exists.
  --existing=query     Query whether to overwrite.
  --imdb-id=id         Query by IMDB id. Hash is tried first unless --force-imdb
                         is used. IMDB URLs are also accepted.
  --force-imdb         Force IMDB search. --imdb-id must be specified.
  --interactive, -i    Equivalent to --download=query --existing=query.
  '''

NAME = 'subdl'
VERSION = '1.0.3'

VERSION_INFO = '''\

This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

http://code.google.com/p/subdl/'''

import os, sys
import struct
import xmlrpclib
import StringIO, gzip, base64
import getopt
import re

osdb_server = "http://api.opensubtitles.org/xml-rpc"
xmlrpc_server = xmlrpclib.Server(osdb_server)
login = xmlrpc_server.LogIn("", "", "en", NAME+" "+VERSION)
osdb_token = login["token"]

class Options: pass
options = Options()
options.lang = 'eng'
options.download = 'first'
options.output = None
options.existing = 'abort'
options.imdb_id = None
options.force_imdb = False;

class SubtitleSearchResult:
    def __init__(self, dict):
        self.__dict__ = dict

def file_ext(filename):
    return filename[filename.rfind('.')+1:]

def file_base(filename):
    return filename[:filename.rfind('.')]

def gunzipstr(zs):
    return gzip.GzipFile(fileobj=StringIO.StringIO(zs)).read()

def writefile(filename, str):
    try:
        open(filename,'wb').write(str)
    except Exception, e:
        raise SystemExit ("Error writing to %s: %s" %(filename, e))

def query_num(s, min, max):
    while True:
        print s,
        try:
            n = raw_input()
        except KeyboardInterrupt:
            raise SystemExit("Aborted")
        try:
            n = int(n)
            if n >= min and n <= max:
                return n
        except:
            pass

def query_yn(s):
    while True:
        print s,
        try:
            s = raw_input().lower()
        except KeyboardInterrupt:
            raise SystemExit("Aborted")
        if s.startswith('y'):
            return True
        elif s.startswith('n'):
            return False

def movie_hash(name):
    longlongformat = '<Q'
    bytesize = struct.calcsize(longlongformat)
    assert bytesize == 8
    f = open(name, "rb")
    filesize = os.path.getsize(name)
    hash = filesize
    if filesize < 65536 * 2:
        raise Exception("Error hashing %s: file too small" %(name))
    for x in range(65536/bytesize):
        hash += struct.unpack(longlongformat, f.read(bytesize))[0]
        hash &= 0xFFFFFFFFFFFFFFFF
    f.seek(filesize-65536,0)
    for x in range(65536/bytesize):
        hash += struct.unpack(longlongformat, f.read(bytesize))[0]
        hash &= 0xFFFFFFFFFFFFFFFF
    f.close()
    return "%016x" % hash

def SearchSubtitlesByHash(filename, langs_search):
    moviehash = movie_hash(filename)
    moviebytesize = os.path.getsize(filename)
    searchlist = [({'sublanguageid':langs_search,
                    'moviehash':moviehash,
                    'moviebytesize':str(moviebytesize)})]
    print >>sys.stderr, "Searching for subtitles for moviehash=%s..." %(moviehash)
    try:
        results = xmlrpc_server.SearchSubtitles(osdb_token,searchlist)
    except Exception, e:
        raise SystemExit("Error in XMLRPC SearchSubtitles call: %s"%e)
    data = results['data']
    return data and [SubtitleSearchResult(d) for d in data]

def SearchSubtitlesByIMDBId(filename, langs_search, imdb_id):
    result = re.search("\d+", imdb_id);
    imdb_id = result.group(0);
    searchlist = [({'sublanguageid':langs_search,
                    'imdbid':imdb_id})]
    print >>sys.stderr, "Searching for subtitles for IMDB id=%s..." %(imdb_id)
    try:
        results = xmlrpc_server.SearchSubtitles(osdb_token,searchlist)
    except Exception, e:
        raise SystemExit("Error in XMLRPC SearchSubtitles call: %s"%e)
    data = results['data']
    return data and [SubtitleSearchResult(d) for d in data]

def format_movie_name(s):
    if s.startswith('"') and s.endswith('"'):
        s = s[1:-1]
    s = s.replace('"', "'")
    return '"%s"' %s

def DisplaySubtitleSearchResults(search_results):
    print "Found %d results:" %(len(search_results))
    idsubtitle_maxlen = 0
    moviename_maxlen = 0
    downloads_maxlen = 0
    for subtitle in search_results:
        idsubtitle = subtitle.IDSubtitleFile
        idsubtitle_maxlen = max(idsubtitle_maxlen, len(idsubtitle))
        moviename = format_movie_name(subtitle.MovieName)
        moviename_maxlen = max(moviename_maxlen, len(moviename))
        downloads = subtitle.SubDownloadsCnt
        downloads_maxlen = max(downloads_maxlen, len(downloads))

    n = 0
    count_maxlen = len(`len(search_results)`)
    for subtitle in search_results:
        n += 1
        idsubtitle = subtitle.IDSubtitleFile
        lang = subtitle.ISO639
        # langn = subtitle.LanguageName
        # str_uploader = subtitle.UserNickName or "Anonymous"
        moviename = format_movie_name(subtitle.MovieName)
        filename = subtitle.SubFileName
        rating = subtitle.SubRating
        downloads = subtitle.SubDownloadsCnt
        # idmovie = subtitle.IDMovie
        # idmovieimdb = subtitle.IDMovieImdb
        print '',
        if options.download == 'query':
            print "%s."%`n`.rjust(count_maxlen),
        print "#%s [%s] [Rat:%s DL:%s] %s %s "%(idsubtitle.rjust(idsubtitle_maxlen),
                                                lang,
                                                rating.rjust(4),
                                                downloads.rjust(downloads_maxlen),
                                                moviename.ljust(moviename_maxlen),
                                                filename.encode('utf8'))

def DownloadSubtitle(sub_id):
    '''Download subtitle #sub_id and return subtitle text as string.'''
    try:
        answer = xmlrpc_server.DownloadSubtitles(osdb_token,[sub_id])
        subtitle_compressed = answer['data'][0]['data']
    except Exception, e:
        raise SystemExit("Error in XMLRPC DownloadSubtitles call: %s"%e)
    return gunzipstr(base64.decodestring(subtitle_compressed))

def DownloadAndSaveSubtitle(sub_id,destfilename):
    if os.path.exists(destfilename):
        if options.existing == 'abort':
            print "Subtitle %s already exists; aborting (try --interactive)."%destfilename
            raise SystemExit(3)
        elif options.existing == 'bypass':
            print "Subtitle %s already exists; bypassing."%destfilename
            return
        elif options.existing == 'overwrite':
            print "Subtitle %s already exists; overwriting."%destfilename
        elif options.existing == 'query':
            if query_yn("Subtitle %s already exists.  Overwrite [y/n]?"%destfilename):
                pass
            else:
                raise SystemExit("File not overwritten.")
        else:
            raise Exception("internal error: bad option.existing=%s"%options.existing)
    print >>sys.stderr, "Downloading #%s to %s..." %(sub_id, destfilename),
    s = DownloadSubtitle(sub_id)
    writefile(destfilename, s)
    print >>sys.stderr, "done, wrote %d bytes."%(len(s))

def replace_fmt(input, replacements):
    output = ''
    i = 0
    while i < len(input):
        c = input[i]
        if c == '%':
            i += 1
            c = input[i]
            try:
                output += replacements[c]
            except:
                raise SystemExit("Bad '%%s' in format specifier" %c)
        else:
            output += c
        i += 1
    return output

def format_subtitle_output_filename(videoname, search_result):
    subname = search_result.SubFileName
    repl = {
        '%': '%',
        'I': search_result.IDSubtitleFile,
        'm': file_base(videoname), 'M': file_ext(videoname),
        's': file_base(subname),   'S': file_ext(subname),
        'l': search_result.LanguageName,
        'L': search_result.ISO639
        }
    output_filename = replace_fmt(options.output, repl)
    assert output_filename != videoname
    return output_filename

def AutoDownloadAndSave(videoname, search_result, downloaded = None):
    output_filename = format_subtitle_output_filename(videoname, search_result)
    if downloaded is not None:
        if output_filename in downloaded:
            raise SystemExit("Already wrote to %s!  Uniquify output filename format."%output_filename)
        downloaded[output_filename] = 1
    DownloadAndSaveSubtitle(search_result.IDSubtitleFile, output_filename)

def select_search_result_by_id(id, search_results):
    for search_result in search_results:
        if search_result.IDSubtitleFile == id:
            return search_result
    raise SystemExit("Search results did not contain subtitle with id %s"%id)

def help():
    print __doc__
    raise SystemExit

def isnumber(value):
    try:
        return int(value) > 0
    except:
        return False

def ListLanguages():
    languages = xmlrpc_server.GetSubLanguages('')['data']
    for language in languages:
        print language['SubLanguageID'], language['ISO639'], language['LanguageName']
    raise SystemExit

def default_output_fmt():
    if options.download == 'all':
        return "%m.%L.%I.%S"
    elif options.lang == 'all' or ',' in options.lang:
        return "%m.%L.%S"
    else:
        return "%m.%S"


def parseargs(args):
    try:
        opts, arguments = getopt.getopt(args, 'h?in', [
                'existing=', 'lang=', 'search-only=',
                'download=', 'output=', 'interactive',
                'list-languages', 'imdb-id=', 'force-imdb',
                'help', 'version', 'versionx'])
    except getopt.GetoptError, e:
        raise SystemExit("%s: %s (see --help)"%(sys.argv[0],e))
    for option, value in opts:
        if option == '--help' or option == '-h' or option == '-?':
            help()
        elif option == '--versionx':
            print VERSION
            raise SystemExit
        elif option == '--version':
            print "%s %s" %(NAME, VERSION)
            raise SystemExit
        elif option == '--existing':
            if value in ['abort', 'overwrite', 'bypass', 'query']:
                pass
            else:
                raise SystemExit("Argument to --existing must be one of: abort, overwrite, bypass, query")
            options.existing = value
        elif option == '--lang':
            options.lang = value
        elif option == '--download':
            if value in ['all', 'first', 'query', 'none'] or isnumber(value):
                pass
            else:
                raise SystemExit("Argument to --download must be numeric subtitle id or one: all, first, query, none")
            options.download = value
        elif option == '-n':
            options.download = 'none'
        elif option == '--output':
            options.output = value
        elif option == '--imdb-id':
            options.imdb_id = value
        elif option == '--force-imdb':
            options.force_imdb = True
        elif option == '--interactive' or option == '-i':
            options.download = 'query'
            options.existing = 'query'
        elif option == '--list-languages':
            ListLanguages()
        else:
            raise SystemExit("internal error: bad option '%s'"%option)
    if not options.output:
        options.output = default_output_fmt()
    if len(arguments) != 1:
        raise SystemExit("syntax: %s [options] filename.avi  (see --help)" %(sys.argv[0]))
    return arguments[0]

def main(args):
    file = parseargs(args)
    if not os.path.exists(file):
        raise SystemExit("can't find file '%s'"%file)
    if options.force_imdb:
        if options.imdb_id == None:
            raise SystemExit("With --force-imdb a --imdb-id must be provided.")
        search_results = SearchSubtitlesByIMDBId(file, options.lang, options.imdb_id)
    else:
        search_results = SearchSubtitlesByHash(file, options.lang)
        if not search_results and options.imdb_id != None:
            print "No results found by hash, trying IMDB id";
            search_results = SearchSubtitlesByIMDBId(file, options.lang, options.imdb_id)
    if not search_results:
        raise SystemExit("No results found.")
    DisplaySubtitleSearchResults(search_results)
    if options.download == 'none':
        raise SystemExit
    elif options.download == 'first':
        if len(search_results) > 1:
            print "Defaulting to first result (try --interactive)."
        AutoDownloadAndSave(file, search_results[0])
    elif options.download == 'all':
        downloaded = {}
        for search_result in search_results:
            AutoDownloadAndSave(file, search_result, downloaded)
    elif options.download == 'query':
        n = query_num("Enter result to download [1..%d]:"%(len(search_results)),
                      1, len(search_results))
        AutoDownloadAndSave(file, search_results[n-1])
    elif isnumber(options.download):
        search_result = select_search_result_by_id(options.download, search_results)
        AutoDownloadAndSave(file, search_result)
    else:
        raise Exception("internal error: bad option.download=%s"%options.download)

main(sys.argv[1:])