mirror of
https://github.com/morpheus65535/bazarr
synced 2024-12-23 00:03:33 +00:00
320 lines
10 KiB
Python
320 lines
10 KiB
Python
# coding=utf-8
|
|
|
|
from __future__ import absolute_import
|
|
import logging
|
|
import types
|
|
import os
|
|
import datetime
|
|
|
|
from guessit import guessit
|
|
from requests.compat import urljoin, quote
|
|
from subliminal import Episode, Movie, region
|
|
from subliminal_patch.core import remove_crap_from_fn
|
|
from subliminal_patch.http import CertifiSession
|
|
import six
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DroneAPIClient(object):
|
|
api_url = None
|
|
_fill_attrs = None
|
|
|
|
def __init__(self, version=1, session=None, headers=None, timeout=10, base_url=None, api_key=None,
|
|
ssl_no_verify=False):
|
|
headers = dict(headers or {}, **{"X-Api-Key": api_key})
|
|
|
|
#: Session for the requests
|
|
self.session = session or CertifiSession()
|
|
if ssl_no_verify:
|
|
self.session.verify = False
|
|
|
|
self.session.timeout = timeout
|
|
self.session.headers.update(headers or {})
|
|
|
|
if not base_url.endswith("/"):
|
|
base_url += "/"
|
|
|
|
if not base_url.startswith("http"):
|
|
base_url = "http://%s" % base_url
|
|
|
|
if not base_url.endswith("api/"):
|
|
self.api_url = urljoin(base_url, "api/")
|
|
|
|
def get_guess(self, video, scene_name):
|
|
raise NotImplemented
|
|
|
|
def get_additional_data(self, video):
|
|
raise NotImplemented
|
|
|
|
def build_params(self, params):
|
|
"""
|
|
quotes values and converts keys of params to camelCase from underscore
|
|
:param params: dict
|
|
:return:
|
|
"""
|
|
out = {}
|
|
for key, value in six.iteritems(params):
|
|
if not isinstance(value, (str,)):
|
|
value = str(value)
|
|
|
|
elif isinstance(value, six.text_type):
|
|
value = value.encode("utf-8")
|
|
|
|
key = key.split('_')[0] + ''.join(x.capitalize() for x in key.split('_')[1:])
|
|
out[key] = quote(value)
|
|
return out
|
|
|
|
def get(self, endpoint, requests_kwargs=None, **params):
|
|
url = urljoin(self.api_url, endpoint)
|
|
params = self.build_params(params)
|
|
|
|
# perform the request
|
|
r = self.session.get(url, params=params, **(requests_kwargs or {}))
|
|
r.raise_for_status()
|
|
|
|
# get the response as json
|
|
j = r.json()
|
|
|
|
# check response status
|
|
if j:
|
|
return j
|
|
return []
|
|
|
|
def status(self, **kwargs):
|
|
return self.get("system/status", requests_kwargs=kwargs)
|
|
|
|
def update_video(self, video, scene_name):
|
|
"""
|
|
update video attributes based on scene_name
|
|
:param video:
|
|
:param scene_name:
|
|
:return:
|
|
"""
|
|
scene_fn, guess = self.get_guess(video, scene_name)
|
|
video_fn = os.path.basename(video.name)
|
|
for attr in self._fill_attrs:
|
|
if attr in guess:
|
|
value = guess.get(attr)
|
|
logger.debug(u"%s: Filling attribute %s: %s", video_fn, attr, value)
|
|
setattr(video, attr, value)
|
|
|
|
video.original_name = scene_fn
|
|
|
|
|
|
def sonarr_series_cache_key(namespace, fn, **kw):
|
|
def generate_key(*arg):
|
|
return "sonarr_series"
|
|
return generate_key
|
|
|
|
|
|
class SonarrClient(DroneAPIClient):
|
|
needs_attrs_to_work = ("series", "season", "episode",)
|
|
_fill_attrs = ("release_group", "format",)
|
|
cfg_name = "sonarr"
|
|
|
|
def __init__(self, base_url="http://127.0.0.1:8989/", **kwargs):
|
|
super(SonarrClient, self).__init__(base_url=base_url, **kwargs)
|
|
|
|
@region.cache_on_arguments(should_cache_fn=lambda x: bool(x),
|
|
function_key_generator=sonarr_series_cache_key)
|
|
def get_all_series(self):
|
|
return self.get("series")
|
|
|
|
def get_show_id(self, video):
|
|
def is_correct_show(s):
|
|
return s["title"] == video.series or (video.series_tvdb_id and "tvdbId" in s and
|
|
s["tvdbId"] == video.series_tvdb_id)
|
|
|
|
for show in self.get_all_series():
|
|
if is_correct_show(show):
|
|
return show["id"]
|
|
|
|
logger.debug(u"%s: Show not found, refreshing cache: %s", video.name, video.series)
|
|
for show in self.get_all_series.refresh(self):
|
|
if is_correct_show(show):
|
|
return show["id"]
|
|
|
|
def get_additional_data(self, video):
|
|
for attr in self.needs_attrs_to_work:
|
|
if getattr(video, attr, None) is None:
|
|
logger.debug(u"%s: Not enough data available for Sonarr", video.name)
|
|
return
|
|
|
|
found_show_id = self.get_show_id(video)
|
|
|
|
if not found_show_id:
|
|
logger.debug(u"%s: Show not found in Sonarr: %s", video.name, video.series)
|
|
return
|
|
|
|
episode_fn = os.path.basename(video.name)
|
|
|
|
for episode in self.get("episode", series_id=found_show_id):
|
|
episode_file = episode.get("episodeFile", {})
|
|
if os.path.basename(episode_file.get("relativePath", "")) == episode_fn:
|
|
scene_name = episode_file.get("sceneName")
|
|
original_filepath = episode_file.get("originalFilePath")
|
|
|
|
data = {}
|
|
if scene_name:
|
|
logger.debug(u"%s: Got scene filename from Sonarr: %s", episode_fn, scene_name)
|
|
data["scene_name"] = scene_name
|
|
|
|
if original_filepath:
|
|
logger.debug(u"%s: Got original file path from Sonarr: %s", episode_fn, original_filepath)
|
|
data["original_filepath"] = original_filepath
|
|
|
|
if data:
|
|
return data
|
|
|
|
logger.debug(u"%s: Can't get original filename, sceneName-attribute not set", episode_fn)
|
|
return
|
|
|
|
logger.debug(u"%s: Episode not found in Sonarr: S%02dE%02d", episode_fn, video.season, video.episode)
|
|
|
|
def get_guess(self, video, scene_name):
|
|
"""
|
|
run guessit on scene_name
|
|
:param video:
|
|
:param scene_name:
|
|
:return:
|
|
"""
|
|
ext = os.path.splitext(video.name)[1]
|
|
guess_from = remove_crap_from_fn(scene_name + ext)
|
|
|
|
# guess
|
|
hints = {
|
|
"single_value": True,
|
|
"type": "episode",
|
|
}
|
|
|
|
return guess_from, guessit(guess_from, options=hints)
|
|
|
|
|
|
def radarr_movies_cache_key(namespace, fn, **kw):
|
|
def generate_key(*arg):
|
|
return "radarr_movies"
|
|
return generate_key
|
|
|
|
|
|
class RadarrClient(DroneAPIClient):
|
|
needs_attrs_to_work = ("title",)
|
|
_fill_attrs = ("release_group", "format",)
|
|
cfg_name = "radarr"
|
|
|
|
def __init__(self, base_url="http://127.0.0.1:7878/", **kwargs):
|
|
super(RadarrClient, self).__init__(base_url=base_url, **kwargs)
|
|
|
|
@region.cache_on_arguments(should_cache_fn=lambda x: bool(x["data"]), function_key_generator=radarr_movies_cache_key)
|
|
def get_all_movies(self):
|
|
return {"d": datetime.datetime.now(), "data": self.get("movie")}
|
|
|
|
def get_movie(self, movie_fn, movie_path):
|
|
def is_correct_movie(m):
|
|
movie_file = movie.get("movieFile", {})
|
|
if os.path.basename(movie_file.get("relativePath", "")) == movie_fn:
|
|
return m
|
|
|
|
res = self.get_all_movies()
|
|
try:
|
|
# get creation date of movie_path to see whether our cache is still valid
|
|
ctime = os.path.getctime(movie_path)
|
|
created = datetime.datetime.fromtimestamp(ctime)
|
|
if created < res["d"]:
|
|
for movie in res["data"]:
|
|
if is_correct_movie(movie):
|
|
return movie
|
|
except TypeError:
|
|
# legacy cache data
|
|
pass
|
|
|
|
logger.debug(u"%s: Movie not found, refreshing cache", movie_fn)
|
|
res = self.get_all_movies.refresh(self)
|
|
for movie in res["data"]:
|
|
if is_correct_movie(movie):
|
|
return movie
|
|
|
|
def get_additional_data(self, video):
|
|
for attr in self.needs_attrs_to_work:
|
|
if getattr(video, attr, None) is None:
|
|
logger.debug(u"%s: Not enough data available for Radarr")
|
|
return
|
|
movie_fn = os.path.basename(video.name)
|
|
|
|
movie = self.get_movie(movie_fn, video.name)
|
|
if not movie:
|
|
logger.debug(u"%s: Movie not found", movie_fn)
|
|
|
|
else:
|
|
movie_file = movie.get("movieFile", {})
|
|
scene_name = movie_file.get("sceneName")
|
|
release_group = movie_file.get("releaseGroup")
|
|
|
|
additional_data = {}
|
|
if scene_name:
|
|
logger.debug(u"%s: Got scene filename from Radarr: %s", movie_fn, scene_name)
|
|
additional_data["scene_name"] = scene_name
|
|
|
|
if release_group:
|
|
logger.debug(u"%s: Got release group from Radarr: %s", movie_fn, release_group)
|
|
additional_data["release_group"] = release_group
|
|
|
|
return additional_data
|
|
|
|
def get_guess(self, video, scene_name):
|
|
"""
|
|
run guessit on scene_name
|
|
:param video:
|
|
:param scene_name:
|
|
:return:
|
|
"""
|
|
ext = os.path.splitext(video.name)[1]
|
|
guess_from = remove_crap_from_fn(scene_name + ext)
|
|
|
|
# guess
|
|
hints = {
|
|
"single_value": True,
|
|
"type": "movie",
|
|
}
|
|
|
|
return guess_from, guessit(guess_from, options=hints)
|
|
|
|
|
|
class DroneManager(object):
|
|
registry = {
|
|
Episode: SonarrClient,
|
|
Movie: RadarrClient,
|
|
}
|
|
|
|
@classmethod
|
|
def get_client(cls, video, cfg_kwa):
|
|
media_type = type(video)
|
|
client_cls = cls.registry.get(media_type)
|
|
if not client_cls:
|
|
raise NotImplementedError("Media type not supported: %s", media_type)
|
|
|
|
return client_cls(**cfg_kwa[client_cls.cfg_name])
|
|
|
|
|
|
def refine(video, **kwargs):
|
|
"""
|
|
|
|
:param video:
|
|
:param embedded_subtitles:
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
|
|
client = DroneManager.get_client(video, kwargs)
|
|
|
|
additional_data = client.get_additional_data(video)
|
|
|
|
if additional_data:
|
|
if "scene_name" in additional_data:
|
|
client.update_video(video, additional_data["scene_name"])
|
|
|
|
elif "original_filepath" in additional_data:
|
|
client.update_video(video, os.path.splitext(additional_data["original_filepath"])[0])
|
|
|
|
if "release_group" in additional_data and not video.release_group:
|
|
video.release_group = remove_crap_from_fn(additional_data["release_group"])
|