Added season pack download feature to titlovi

5 files changed, 283 insertions, 7 deletions

diff --git a/custom_libs/subliminal_patch/providers/titlovi.py b/custom_libs/subliminal_patch/providers/titlovi.py
index 88782522c..c7682ec9b 100644
--- a/custom_libs/subliminal_patch/providers/titlovi.py
+++ b/custom_libs/subliminal_patch/providers/titlovi.py
@@ -56,7 +56,7 @@ class TitloviSubtitle(Subtitle):
     provider_name = 'titlovi'
 
     def __init__(self, language, download_link, sid, releases, title, alt_title=None, season=None,
-                 episode=None, year=None, rating=None, download_count=None, asked_for_release_group=None, asked_for_episode=None):
+                 episode=None, year=None, rating=None, download_count=None, asked_for_release_group=None, asked_for_episode=None, is_pack=False):
         super(TitloviSubtitle, self).__init__(language)
         self.sid = sid
         self.releases = self.release_info = releases
@@ -71,6 +71,7 @@ class TitloviSubtitle(Subtitle):
         self.matches = None
         self.asked_for_release_group = asked_for_release_group
         self.asked_for_episode = asked_for_episode
+        self.is_pack = is_pack
 
     def __repr__(self):
         if self.season and self.episode:
@@ -216,7 +217,7 @@ class TitloviProvider(Provider, ProviderSubtitleArchiveMixin):
         is_episode = False
         if season and episode:
             is_episode = True
-            #search_params['season'] = season
+            search_params['season'] = season
             #search_params['episode'] = episode
         #if year:
         #    search_params['year'] = year
@@ -238,6 +239,18 @@ class TitloviProvider(Provider, ProviderSubtitleArchiveMixin):
             resp_json = response.json()
             if resp_json['SubtitleResults']:
                 query_results.extend(resp_json['SubtitleResults'])
+            
+            # if there are more pages, loop through them. If there is more than 3 pages, stop at 3
+            if resp_json['PagesAvailable'] > 1:
+                for page in range(2, min(4, resp_json['PagesAvailable'] + 1)):
+                    search_params['pg'] = page
+                    response = self.get_result(self.api_search_url, search_params)
+                    resp_json = response.json()
+                    if resp_json['SubtitleResults']:
+                        query_results.extend(resp_json['SubtitleResults'])
+                    else:
+                        break
+
         except TooManyRequests:
             raise
         except Exception as e:
@@ -258,15 +271,19 @@ class TitloviProvider(Provider, ProviderSubtitleArchiveMixin):
                 # skip if season and episode number does not match
                 if season and season != sub.get('Season'):
                     continue
-                elif episode and episode != sub.get('Episode'):
+                elif episode and episode != sub.get('Episode') and sub.get('Episode') != 0:
                     continue
 
+                is_pack = False
+                if sub.get('Episode') == 0:
+                    is_pack = True
+                    
                 subtitle = self.subtitle_class(Language.fromtitlovi(sub.get('Lang')), sub.get('Link'), sub.get('Id'), sub.get('Release'), _title,
-                                               alt_title=alt_title, season=sub.get('Season'), episode=sub.get('Episode'),
+                                               alt_title=alt_title, season=sub.get('Season'), episode=episode,
                                                year=sub.get('Year'), rating=sub.get('Rating'),
                                                download_count=sub.get('DownloadCount'),
                                                asked_for_release_group=video.release_group,
-                                               asked_for_episode=episode)
+                                               asked_for_episode=episode, is_pack=is_pack)
             else:
                 subtitle = self.subtitle_class(Language.fromtitlovi(sub.get('Lang')), sub.get('Link'), sub.get('Id'), sub.get('Release'), _title,
                                                alt_title=alt_title, year=sub.get('Year'), rating=sub.get('Rating'),
@@ -321,13 +338,25 @@ class TitloviProvider(Provider, ProviderSubtitleArchiveMixin):
 
         subs_in_archive = archive.namelist()
 
-        # if Serbian lat and cyr versions are packed together, try to find right version
-        if len(subs_in_archive) > 1 and (subtitle.language == 'sr' or subtitle.language == 'sr-Cyrl'):
+        if len(subs_in_archive) > 1 and subtitle.is_pack:
+            # if subtitle is a pack, try to find the right subtitle by format SSxEE or SxxEyy
+            self.get_subtitle_from_pack(subtitle, subs_in_archive, archive)
+        elif len(subs_in_archive) > 1 and (subtitle.language == 'sr' or subtitle.language == 'sr-Cyrl'):
+            # if Serbian lat and cyr versions are packed together, try to find right version
             self.get_subtitle_from_bundled_archive(subtitle, subs_in_archive, archive)
         else:
             # use default method for everything else
             subtitle.content = self.get_subtitle_from_archive(subtitle, archive)
 
+    def get_subtitle_from_pack(self, subtitle, subs_in_archive, archive):
+        # try to find the right subtitle, it should contain season and episode number in format SSxEE or SxxEyy
+        format1 = '%.2dx%.2d' % (subtitle.season, subtitle.episode)
+        format2 = 's%.2de%.2d' % (subtitle.season, subtitle.episode)
+        for sub_name in subs_in_archive:
+            if format1 in sub_name.lower() or format2 in sub_name.lower():
+                subtitle.content = fix_line_ending(archive.read(sub_name))
+                return
+    
     def get_subtitle_from_bundled_archive(self, subtitle, subs_in_archive, archive):
         sr_lat_subs = []
         sr_cyr_subs = []
diff --git a/tests/subliminal_patch/data/titlovi_gettoken_response.json b/tests/subliminal_patch/data/titlovi_gettoken_response.json
new file mode 100644
index 000000000..ab1627590
--- /dev/null
+++ b/tests/subliminal_patch/data/titlovi_gettoken_response.json
@@ -0,0 +1,6 @@
+{
+  "ExpirationDate": "2024-10-06T19:05:13.5",
+  "Token": "asdf1234",
+  "UserId": 111,
+  "UserName": "user1"
+}
\ No newline at end of file
diff --git a/tests/subliminal_patch/data/titlovi_search_response.json b/tests/subliminal_patch/data/titlovi_search_response.json
new file mode 100644
index 000000000..2bbe054ac
--- /dev/null
+++ b/tests/subliminal_patch/data/titlovi_search_response.json
@@ -0,0 +1,172 @@
+{
+  "ResultsFound": 11,
+  "PagesAvailable": 1,
+  "CurrentPage": 1,
+  "SubtitleResults": [
+      {
+          "Id": 346305,
+          "Title": "Nikita aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=346305",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2022-04-11T14:03:30.59",
+          "DownloadCount": 415,
+          "Rating": 0.0,
+          "Release": "PROPER.FRENCH.1080p.BluRay.x264.TrueHD.5.1-FGT"
+      },
+      {
+          "Id": 323824,
+          "Title": "Nikita Aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=323824",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2021-02-21T23:53:51.257",
+          "DownloadCount": 397,
+          "Rating": 0.0,
+          "Release": "720p BluRay x264 DTS-PRoDJi"
+      },
+      {
+          "Id": 120571,
+          "Title": "Nikita Aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=120571",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2011-02-28T22:54:45.7",
+          "DownloadCount": 3543,
+          "Rating": 0.0,
+          "Release": "720p.BD rip"
+      },
+      {
+          "Id": 91576,
+          "Title": "La Femme Nikita",
+          "Year": 1997,
+          "Type": 2,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=91576",
+          "Season": 5,
+          "Episode": 0,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2009-12-21T23:13:20.407",
+          "DownloadCount": 3227,
+          "Rating": 0.0,
+          "Release": ""
+      },
+      {
+          "Id": 81025,
+          "Title": "La Femme Nikita",
+          "Year": 1997,
+          "Type": 2,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=81025",
+          "Season": 4,
+          "Episode": 0,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2009-06-05T03:09:19.77",
+          "DownloadCount": 3799,
+          "Rating": 0.0,
+          "Release": ""
+      },
+      {
+          "Id": 81024,
+          "Title": "La Femme Nikita",
+          "Year": 1997,
+          "Type": 2,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=81024",
+          "Season": 3,
+          "Episode": 0,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2009-06-05T03:07:39.683",
+          "DownloadCount": 3842,
+          "Rating": 0.0,
+          "Release": ""
+      },
+      {
+          "Id": 81023,
+          "Title": "La Femme Nikita",
+          "Year": 1997,
+          "Type": 2,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=81023",
+          "Season": 2,
+          "Episode": 0,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2009-06-05T03:06:06.21",
+          "DownloadCount": 4310,
+          "Rating": 0.0,
+          "Release": ""
+      },
+      {
+          "Id": 81022,
+          "Title": "La Femme Nikita",
+          "Year": 1997,
+          "Type": 2,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=81022",
+          "Season": 1,
+          "Episode": 0,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2009-06-05T03:04:40.14",
+          "DownloadCount": 3924,
+          "Rating": 0.0,
+          "Release": ""
+      },
+      {
+          "Id": 69118,
+          "Title": "Nikita Aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=69118",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2008-12-07T18:48:22.087",
+          "DownloadCount": 4950,
+          "Rating": 5.0,
+          "Release": "720p.BluRay.x264-SiNNERS"
+      },
+      {
+          "Id": 14697,
+          "Title": "Nikita Aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=14697",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2006-03-14T11:29:44.45",
+          "DownloadCount": 2188,
+          "Rating": 5.0,
+          "Release": ""
+      },
+      {
+          "Id": 10582,
+          "Title": "Nikita Aka La Femme Nikita",
+          "Year": 1990,
+          "Type": 1,
+          "Link": "https://titlovi.com/download/?type=1&mediaid=10582",
+          "Season": -1,
+          "Episode": -1,
+          "Special": -1,
+          "Lang": "Srpski",
+          "Date": "2005-09-24T19:40:34.233",
+          "DownloadCount": 1214,
+          "Rating": 0.0,
+          "Release": ""
+      }
+  ]
+}
\ No newline at end of file
diff --git a/tests/subliminal_patch/data/titlovi_some_subtitle_pack.zip b/tests/subliminal_patch/data/titlovi_some_subtitle_pack.zip
new file mode 100644
index 000000000..9c66352b1
--- /dev/null
+++ b/tests/subliminal_patch/data/titlovi_some_subtitle_pack.zip
diff --git a/tests/subliminal_patch/test_titlovi.py b/tests/subliminal_patch/test_titlovi.py
new file mode 100644
index 000000000..886c43be0
--- /dev/null
+++ b/tests/subliminal_patch/test_titlovi.py
@@ -0,0 +1,69 @@
+import pytest
+import subliminal
+import datetime
+import tempfile
+import os
+
+from subliminal_patch.providers.titlovi import TitloviProvider 
+from subliminal_patch.providers.titlovi import TitloviSubtitle 
+from dogpile.cache.region import register_backend as register_cache_backend
+from subliminal_patch.core import Episode
+from subzero.language import Language
+from subliminal.subtitle import fix_line_ending
+
+from zipfile import ZipFile
+
[email protected](scope="session")
+def titlovi_episodes():
+    return {
+        "la_femme_nikita_s01e13": Episode(
+            "La Femme Nikita (1997) - S01E13 - Recruit [HDTV-720p][Opus 2.0][x265].mkv",
+            "La Femme Nikita",
+            1,
+            13,
+            series_imdb_id="tt21209876",
+            video_codec="x265",
+        ),
+    }
+
[email protected](scope="session")
+def region():
+    register_cache_backend("subzero.cache.file", "subzero.cache_backends.file", "SZFileBackend")
+    subliminal.region.configure(
+        "subzero.cache.file",
+        expiration_time=datetime.timedelta(days=30),
+        arguments={"appname": "sz_cache", "app_cache_dir": tempfile.gettempdir()},
+        replace_existing_backend=True,
+    )
+    subliminal.region.backend.sync()
+
+def test_list_subtitles_and_download_from_pack(region, titlovi_episodes, requests_mock, data):
+    language = Language.fromietf('sr-Latn')
+    item = titlovi_episodes["la_femme_nikita_s01e13"]
+
+    with open(os.path.join(data, 'titlovi_gettoken_response.json'), "rb") as f:
+        response = f.read()
+        requests_mock.post('https://kodi.titlovi.com/api/subtitles/gettoken?username=user1&password=pass1&json=True', content=response)
+
+    with open(os.path.join(data, 'titlovi_search_response.json'), "rb") as f:
+        response = f.read()
+        requests_mock.get('https://kodi.titlovi.com/api/subtitles/search?token=asdf1234&userid=111&&query=la femme nikita&lang=Srpski&json=True', content=response)
+        
+    with open(os.path.join(data, 'titlovi_some_subtitle_pack.zip'), "rb") as f:
+        response = f.read()
+        requests_mock.get('https://titlovi.com/download/?type=1&mediaid=81022', content=response)
+        
+    with TitloviProvider('user1','pass1') as provider:
+        subtitles = provider.list_subtitles(item, languages={language})
+
+        assert len(subtitles) == 1
+
+        subtitle = subtitles[0]
+        provider.download_subtitle(subtitle)
+        with open(os.path.join(data, 'titlovi_some_subtitle_pack.zip'), "rb") as f:
+          archive = ZipFile(f)
+          # subs_in_archive = archive.namelist()
+          subtitle_content = fix_line_ending(archive.read('La Femme Nikita - 01x13 - Recruit.srt'))
+          assert(subtitle.content == subtitle_content)
+
+