bazarr/libs/subliminal_patch/refiners/drone.py

# 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"])