plugin.video.torrenter/resources/btclient/opensubtitle.py

264 lines
9.4 KiB
Python

'''
Created on Apr 2, 2015
@author: ivan
'''
import xmlrpclib
import urllib2
import os.path
import gzip
from StringIO import StringIO
import logging
import subprocess
import struct
import time
from _version import __version__
logger = logging.getLogger('opensubtitles')
class Urllib2Transport(xmlrpclib.Transport):
def __init__(self, opener=None, https=False, use_datetime=0):
xmlrpclib.Transport.__init__(self, use_datetime)
self.opener = opener or urllib2.build_opener()
self.https = https
def request(self, host, handler, request_body, verbose=0):
proto = ('http', 'https')[bool(self.https)]
req = urllib2.Request('%s://%s%s' % (proto, host, handler), request_body)
req.add_header('User-agent', self.user_agent)
self.verbose = verbose
return self.parse_response(self.opener.open(req))
class OpenSubtitles(object):
USER_AGENT = 'BTClient v%s' % __version__
def __init__(self, lang, user='', pwd=''):
self._lang = lang
self._proxy = xmlrpclib.ServerProxy('http://api.opensubtitles.org/xml-rpc',
Urllib2Transport(use_datetime=True),
allow_none=True, use_datetime=True)
self._token = None
self._user = user
self._pwd = pwd
def login(self):
res = self._proxy.LogIn(self._user, self._pwd, 'en', self.USER_AGENT)
self._parse_status(res)
token = res.get('token')
if token:
self._token = token
else:
raise xmlrpclib.Fault('NO_TOKEN', 'No token!')
def _parse_status(self, res):
if res.has_key('status'):
code = res['status'].split()[0]
if code != '200':
raise xmlrpclib.Fault('ERROR_CODE_RETURENED', 'Returned error status: %s (%s)' % (code, res))
return True
else:
raise xmlrpclib.Fault('NO_STATUS', 'No status!')
def search(self, filename, filesize=None, filehash=None, limit=20):
filename = os.path.split(filename)[1]
name = os.path.splitext(filename)[0]
query = []
if filehash and filesize:
query.append({'sublanguageid': self._lang, 'moviehash': filehash, 'moviebytesize': str(filesize)})
query.append({'sublanguageid': self._lang, 'tag': filename})
query.append({'sublanguageid': self._lang, 'query': name})
res = self._proxy.SearchSubtitles(self._token, query, {'limit': limit})
self._parse_status(res)
data = res.get('data')
return data if data else []
@staticmethod
def _sub_file(filename, lang, ext):
lang = lang.lower()
path, fname = os.path.split(filename)
fname = os.path.splitext(fname)[0]
return os.path.join(path, fname + '.' + lang + '.' + ext)
@staticmethod
def _base_name(filename):
fname = os.path.split(filename)[1]
return os.path.splitext(fname)[0]
@staticmethod
def hash_file(f, filesize):
longlongformat = '<q' # little-endian long long
bytesize = struct.calcsize(longlongformat)
hash = filesize # @ReservedAssignment
if filesize < 65536 * 2:
raise ValueError("SizeError")
for _x in range(65536 / bytesize):
buffer = f.read(bytesize) # @ReservedAssignment
(l_value,) = struct.unpack(longlongformat, buffer)
hash += l_value
hash = hash & 0xFFFFFFFFFFFFFFFF # to remain as 64bit number @ReservedAssignment
f.seek(max(0, filesize - 65536))
for _x in range(65536 / bytesize):
buffer = f.read(bytesize) # @ReservedAssignment
(l_value,) = struct.unpack(longlongformat, buffer)
hash += l_value
hash = hash & 0xFFFFFFFFFFFFFFFF # @ReservedAssignment
returnedhash = "%016x" % hash
return returnedhash
def choose(self, data):
null = f = open(os.devnull, "w")
items = []
for l in data:
items.append(l['SubDownloadLink'])
items.append(l['SubFileName'])
items.append(l['SubDownloadsCnt'])
p = subprocess.Popen(
'zenity --list --title "Select subtitles" --text "Select best matching subtitles" --width 1024 --height 600 --column Link --column Name --column Downloads --hide-column=1',
stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=null, shell=True)
res, _ = p.communicate(u'\n'.join(items).encode('utf-8'))
null.close()
return res if res.startswith('http') else None
@staticmethod
def download_if_not_exists(filename, lang, filesize=None, filehash=None, sub_ext='srt',
can_choose=True, overwrite=False, retries=3):
sfile = OpenSubtitles._sub_file(filename, lang, sub_ext)
if os.path.exists(sfile) and os.stat(sfile).st_size > 0 and not overwrite:
logger.debug('subs %s are already downloaded', sfile)
return sfile
else:
while True:
try:
with OpenSubtitles(lang) as opensub:
res = opensub.download(filename, filesize, filehash, can_choose)
if res:
logger.debug('Subtitles %s downloaded', res)
return res
else:
logger.debug('No subtitles found for file %s in language %s', filename, lang)
return
except urllib2.HTTPError, e:
retries -= 1
if retries <= 0:
raise e
logger.debug('Retrying to load subtitles due to HTTP error %d, remains %d attempts', e.code,
retries)
time.sleep(1)
def download(self, filename, filesize=None, filehash=None, can_choose=True):
data = self.search(filename, filesize, filehash)
if not data:
return None
media_file = OpenSubtitles._base_name(filename).lower()
def same_name(b):
return media_file == OpenSubtitles._base_name(b['SubFileName']).lower()
if filehash and filesize:
match = filter(lambda x: x.get('QueryNumber', 0) == 0, data)
logger.debug('Got results by filehash')
else:
match = filter(same_name, data)
if match and can_choose != 'always':
sub = match[0]
link = sub['SubDownloadLink']
ext = sub['SubFormat']
logger.debug('Find exact match for media file, subtitle is %s', sub['SubFileName'])
elif can_choose:
link = self.choose(data)
ext = 'srt'
else:
sub = data[0]
link = sub['SubDownloadLink']
ext = sub['SubFormat']
if link:
return self.download_link(filename, link, ext)
def download_link(self, filename, link, ext):
out_file = OpenSubtitles._sub_file(filename, self._lang, ext)
res = urllib2.urlopen(link, timeout=10)
data = StringIO(res.read())
data.seek(0)
res.close()
z = gzip.GzipFile(fileobj=data)
with open(out_file, 'wb') as f:
while True:
d = z.read(1024)
if not d:
break
f.write(d)
z.close()
return out_file
def logout(self):
try:
res = self._proxy.LogOut(self._token)
self._parse_status(res)
except urllib2.HTTPError:
logger.warn('Failed to logout')
def __enter__(self):
self.login()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.logout()
def down(f, lang, overwrite=False):
filesize, filehash = calc_hash(f)
OpenSubtitles.download_if_not_exists(f, lang, filesize=filesize,
filehash=filehash, can_choose=True, overwrite=overwrite)
def calc_hash(f):
if not os.access(f, os.R_OK):
raise ValueError('Cannot read from file %s' % f)
filesize = os.stat(f).st_size
with open(f, 'rb') as fs:
filehash = OpenSubtitles.hash_file(fs, filesize)
return filesize, filehash
def list_subs(f, lang):
import pprint
filesize, filehash = calc_hash(f)
with OpenSubtitles(lang) as opensub:
res = opensub.search(f, filesize, filehash)
res = map(lambda x: {'SubFileName': x['SubFileName'],
'SubDownloadsCnt': x['SubDownloadsCnt'],
'QueryNumber': x.get('QueryNumber', 0),
},
res)
pprint.pprint(res)
if __name__ == '__main__':
from argparse import ArgumentParser
p = ArgumentParser()
p.add_argument("video_file", help="video file")
p.add_argument("-d", "--download", action="store_true", help="Download subtitles for video files")
p.add_argument("-l", "--list", action="store_true", help="List available subtitles")
p.add_argument("--lang", default='eng', help="Language")
p.add_argument("--debug", action="store_true", help="Print debug messages")
p.add_argument("--overwrite", action="store_true", help="Overwrite existing subtitles ")
args = p.parse_args()
if args.debug:
logging.basicConfig(level=logging.DEBUG)
if args.download:
down(args.video_file, args.lang, args.overwrite)
else:
list_subs(args.video_file, args.lang)