diff --git a/docs/guide/cli-reference.md b/docs/guide/cli-reference.md index 4e0ec417..90e4b2bd 100644 --- a/docs/guide/cli-reference.md +++ b/docs/guide/cli-reference.md @@ -146,6 +146,43 @@ Keep only certain track types, or skip certain track types. Attachments are alwa | `--tag` | Group tag override. | | `--repack` | Add a `REPACK` tag to the filename. | +### Cross-service muxing + +Take some track types from a different service and mux them with the rest. Each `--cross-*` flag +takes a service tag and a URL. Video and chapters replace what the service you download provides; +audio and subtitles are merged in, so the best-per-language selection runs across both services' +tracks together. See [Cross-service muxing](cross-service-muxing.md) for the full guide. + +| Flag | Description | +|---|---| +| `--cross-video` | Take video tracks from `SERVICE URL` (replaces the download service's video). | +| `--cross-audio` | Merge in audio tracks from `SERVICE URL`. | +| `--cross-subtitles` | Merge in subtitle tracks from `SERVICE URL`. | +| `--cross-chapters` | Take chapters from `SERVICE URL` (replaces the download service's chapters). | +| `--cross-audio-offset` | Shift the cross-sourced audio in the mux, e.g. `10s`, `500ms`, `-5.5s`. A bare number is milliseconds. | +| `--cross-subtitle-offset` | Shift the cross-sourced subtitles in the mux. Same format as above. | +| `--cross-profile` | Profile used for the cross-service credentials. Defaults to `--profile`. | +| `--cross-proxy` | Proxy for every cross-service. Overrides each cross-service's own proxy configuration. | +| `--cross-wanted` | Episode to take from the cross-services, e.g. `S01E02`. Fallback for the per-type flags below. Defaults to the same season and episode as the title you download. | +| `--cross-video-wanted` | Episode to take video from, e.g. `S01E02`. Overrides `--cross-wanted` for video. | +| `--cross-audio-wanted` | Episode to take audio from. Overrides `--cross-wanted` for audio. | +| `--cross-subtitles-wanted` | Episode to take subtitles from. Overrides `--cross-wanted` for subtitles. | +| `--cross-chapters-wanted` | Episode to take chapters from. Overrides `--cross-wanted` for chapters. | + +```bash +unshackle dl --cross-audio EXAMPLE2 "https://example2.com/show/1" EXAMPLE "https://example.com/show/1" +``` + +Each cross-service authenticates, licenses and decrypts its own tracks with its own CDM, key +vaults and (unless `--cross-proxy` is set) its own proxy, so you need credentials for it as well. +Offsets are applied by the muxer, which does not re-time the media itself. + +!!! note "Titles must match" + For a series, unshackle looks for the same season and episode number on the cross-service. + Use `--cross-wanted` (or a per-type flag such as `--cross-audio-wanted`) when the two services + number the episodes differently. If no title matches, unshackle logs a warning and keeps the + tracks of the service you download. + ### Metadata & tagging | Flag | Description | diff --git a/docs/guide/cross-service-muxing.md b/docs/guide/cross-service-muxing.md new file mode 100644 index 00000000..4602f551 --- /dev/null +++ b/docs/guide/cross-service-muxing.md @@ -0,0 +1,101 @@ +# Cross-service muxing + +Some titles are on more than one service, and each service is strong at something +different: one carries the 4K HDR video, another the better dub, a third the fuller +subtitles. Cross-service muxing lets you download the video from one service and pull +the audio, subtitles, or chapters from others, then mux the result into a single file. + +!!! warning "This is a power-user feature" + unshackle does not check that the services carry the same cut of the title. You are + responsible for pairing services that share the same edit. A different cut, a + different frame rate, or an extra recap will not line up, and a single offset cannot + fix it. See [Limitations](#limitations). + +## How it works + +You download from a **primary** service as normal, the positional `SERVICE URL` at the +end of the command. Each `--cross-*` flag names **another** service tag and URL to take a +track type from: + +- **Video** and **chapters** *replace* what the primary provides. There is only ever one + video source and one set of chapters. +- **Audio** and **subtitles** *merge* with the primary's own tracks. Your normal track + selection (`-l`, quality, codec, channels) then runs across the combined set, so the + best version of each language wins whichever service it came from. + +```bash +# Download EXAMPLE's video, but take the audio from EXAMPLE2 +unshackle dl --cross-audio EXAMPLE2 "https://example2.com/show/1" EXAMPLE "https://example.com/show/1" +``` + +You can combine several sources in one run: + +```bash +# Video from the primary, a dub from one service, subtitles from another +unshackle dl \ + --cross-audio EXAMPLE2 "https://example2.com/show/1" \ + --cross-subtitles EXAMPLE3 "https://example3.com/show/1" \ + EXAMPLE "https://example.com/show/1" +``` + +The full flag list is in the [CLI reference](cli-reference.md#cross-service-muxing). + +## Authentication, DRM, and proxies + +Each cross-service behaves like a download of its own. It authenticates with **its own** +credentials, licenses and decrypts with **its own** CDM, and stores keys in **its own** +key vault. A PlayReady-only cross-service works even when the primary uses Widevine, +because it loads the CDM its own configuration asks for. + +!!! note "You need credentials for every service" + A cross-service authenticates the same way it would on a normal download, so its + profile, cookies, or credentials must be set up. Use `--cross-profile` to pick a + profile for the cross-services (it defaults to `--profile`). + +By default each cross-service resolves **its own** proxy from its configuration, so a +geo-restricted source still reaches the right region. Set `--cross-proxy` to force one +proxy for every cross-service instead. + +## Choosing the episode + +For a TV title, each cross-service is matched to the **same season and episode** as the +title you download. When a service numbers its episodes differently, override the match: + +- `--cross-wanted S01E03` sets a shared override for every cross-service. +- `--cross-audio-wanted`, `--cross-video-wanted`, `--cross-subtitles-wanted`, and + `--cross-chapters-wanted` are per-type overrides that win over `--cross-wanted`. + +If a cross-service has no title matching the request, unshackle warns and keeps the +primary's own tracks for that type. + +## Offsets + +If a cross-service begins with a short buffer the primary does not have, such as an advert +or a splash screen, its tracks start a few seconds late. Nudge them back into line: + +```bash +# EXAMPLE2's audio starts 3 seconds late; shift it earlier +unshackle dl --cross-audio EXAMPLE2 "https://example2.com/show/1" \ + --cross-audio-offset -3s EXAMPLE "https://example.com/show/1" +``` + +`--cross-audio-offset` and `--cross-subtitle-offset` accept `10s`, `500ms`, or `-5.5s`; a +bare number is milliseconds. The offset is passed to the muxer, which shifts the track's +start. It does not re-time or resample the media, so a single offset only corrects a +constant lead or lag, not drift that grows through the title. + +## Limitations + +Cross-service muxing assumes the sources are the same cut of the title. The current +release does not: + +- **Sync mismatched cuts.** Different edits, ad breaks, or frame rates are out of scope; + the offset is a manual nudge for a constant lead, not automatic synchronisation. +- **Pin a language to a service.** Merged audio and subtitles are chosen by the normal + best-per-language selection. You cannot yet force the English from one service when + another service's English scores higher. +- **Reuse a cross-service across a season.** Each title in a season download loads and + authenticates its cross-services again, so a login-rate-limited service is contacted + once per episode. +- **Run a cross-service's own post-download steps.** Only the primary service's + post-download hooks fire. diff --git a/mkdocs.yml b/mkdocs.yml index 8fa8d712..b5ac292b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -113,6 +113,7 @@ nav: - User Guide: - Downloading: guide/downloading.md - CLI Reference: guide/cli-reference.md + - Cross-service Muxing: guide/cross-service-muxing.md - DRM & CDM Setup: guide/drm-and-cdm.md - Proxies & VPN: guide/proxies-and-vpn.md - Key Vaults: guide/vaults.md diff --git a/tests/commands/__init__.py b/tests/commands/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/commands/test_cross_mux.py b/tests/commands/test_cross_mux.py new file mode 100644 index 00000000..ed4db286 --- /dev/null +++ b/tests/commands/test_cross_mux.py @@ -0,0 +1,433 @@ +"""Unit and offline end-to-end tests for cross-service muxing. + +The ``dl`` cross-mux layer takes some track types from a second service and muxes them with the +first. These tests cover the behaviour that is easy to get wrong: + +- ``drm_source`` routes each track to the service, title, CDM and vaults that licence it, so a + cross-sourced track licenses through its own service and a primary track through the primary's. +- Keys from a cross-service cache under that service's own vault, never the primary's. +- Audio and subtitles merge into the primary's own tracks (video and chapters replace), and an + explicit id shared by both services is namespaced so the two never collide. +- A ``sync_offset_ms`` on a track becomes ``--sync 0:`` in the built mkvmerge command, and is + absent when no offset is set. +- ``match_cross_title`` honours the per-type ``--cross-*-wanted`` override, falling back to the + shared ``--cross-wanted`` and then to the matching season/episode. +- The new CLI options parse and bind. + +The offline end-to-end test stands up a synthetic cross-service (its titles and tracks held in +memory, so no network or manifest parsing is needed) and drives ``fetch_cross_tracks`` -> +``apply_cross_tracks`` -> ``Tracks.mux``, asserting the cross track is merged, muxed with its +offset, and routed through the correct service for DRM. +""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace +from typing import Any, Optional + +import click +import pytest + +from unshackle.commands.dl import dl +from unshackle.core.titles.episode import Episode +from unshackle.core.tracks import Audio, Subtitle, Tracks, Video +from unshackle.core.tracks import tracks as tracks_mod + +pytestmark = pytest.mark.unit + + +class Svc: + """Stand-in for a service class, only used to satisfy the Title ``service`` type check.""" + + +class CrossSvc: + """A minimal cross-service: its class name drives namespacing and the ``cross_source`` tag.""" + + def __init__(self, titles: Optional[list] = None, tracks: Optional[Tracks] = None) -> None: + self.session = object() + self._titles = titles or [] + self._tracks = tracks or Tracks() + + def get_titles(self) -> list: + return self._titles + + def get_tracks(self, title: Any) -> Tracks: + return self._tracks + + def get_chapters(self, title: Any) -> list: + return [] + + +class FakeVaults: + """Records added keys so we can prove which vault a key landed in.""" + + def __init__(self, name: str) -> None: + self.name = name + self.keys: dict[str, str] = {} + + def add_key(self, kid: str, key: str, excluding: Any = None) -> int: + self.keys[kid] = key + return 1 + + def get_key(self, kid: str) -> tuple[Optional[str], None]: + return self.keys.get(kid), None + + +def make_audio(track_id: str, *, language: str = "en", offset: Optional[int] = None, path: Any = None) -> Audio: + track = Audio( + id_=track_id, + url=f"https://example.test/{track_id}.m4a", + language=language, + codec=Audio.Codec.AAC, + bitrate=128000, + ) + if offset is not None: + track.data["sync_offset_ms"] = offset + if path is not None: + path.touch() # mkvmerge command building requires the track file to exist + track.path = path + return track + + +def make_subtitle(track_id: str, *, language: str = "en", offset: Optional[int] = None, path: Any = None) -> Subtitle: + track = Subtitle( + id_=track_id, + url=f"https://example.test/{track_id}.vtt", + language=language, + codec=Subtitle.Codec.WebVTT, + ) + if offset is not None: + track.data["sync_offset_ms"] = offset + if path is not None: + path.touch() # mkvmerge command building requires the track file to exist + track.path = path + return track + + +def make_video(track_id: str, *, language: str = "en") -> Video: + return Video( + id_=track_id, + url=f"https://example.test/{track_id}.m3u8", + language=language, + codec=Video.Codec.HEVC, + range_=Video.Range.SDR, + width=1920, + height=1080, + bitrate=5000000, + ) + + +def make_dl(**overrides: Any) -> dl: + """Build a bare ``dl`` with only the attributes the cross-mux methods read (no heavy __init__).""" + inst = dl.__new__(dl) + inst.log = logging.getLogger("test.cross") + inst.cross_video = None + inst.cross_audio = None + inst.cross_subtitles = None + inst.cross_chapters = None + inst.cross_audio_offset = None + inst.cross_subtitle_offset = None + inst.cross_wanted = None + inst.cross_wanted_by_type = {"video": None, "audio": None, "subtitles": None, "chapters": None} + inst.cross_track_sources = {} + inst.cdm = object() + inst.vaults = object() + for key, value in overrides.items(): + setattr(inst, key, value) + return inst + + +def capture_mux_command(monkeypatch: pytest.MonkeyPatch) -> dict: + """Patch mkvmerge discovery and Popen so ``Tracks.mux`` builds a command without running it.""" + captured: dict = {} + + class FakeStdout: + def readline(self) -> str: + return "" + + class FakePopen: + def __init__(self, cmd: list, **_: Any) -> None: + captured["cmd"] = cmd + self.stdout = FakeStdout() + + def wait(self) -> int: + return 0 + + monkeypatch.setattr(tracks_mod.binaries, "MKVToolNix", "/usr/bin/mkvmerge") + monkeypatch.setattr(tracks_mod.subprocess, "Popen", FakePopen) + return captured + + +def test_drm_source_routes_cross_and_primary() -> None: + d = make_dl() + primary_service = SimpleNamespace(session=object()) + primary_title = object() + cross_service = SimpleNamespace(session=object()) + cross_title = object() + cross_cdm = object() + cross_vaults = object() + + cross_track = make_audio("cross-1", language="ja") + d.cross_track_sources[cross_track.id] = (cross_service, cross_title, cross_cdm, cross_vaults) + + assert d.drm_source(cross_track, primary_service, primary_title) == ( + cross_service, + cross_title, + cross_cdm, + cross_vaults, + ) + + primary_track = make_audio("prim-1") + assert d.drm_source(primary_track, primary_service, primary_title) == ( + primary_service, + primary_title, + d.cdm, + d.vaults, + ) + + +def test_cross_key_caches_under_owning_vault() -> None: + primary_vaults = FakeVaults("primary") + cross_vaults = FakeVaults("cross") + d = make_dl(vaults=primary_vaults) + + cross_service = CrossSvc() + cross_track = make_audio("cross-1", language="ja") + d.cross_track_sources[cross_track.id] = (cross_service, object(), object(), cross_vaults) + + # download_track resolves the owning vault through drm_source and threads it into prepare_drm, + # which caches licensed keys to exactly that vault. + _, _, _, vaults = d.drm_source(cross_track, SimpleNamespace(session=object()), object()) + assert vaults is cross_vaults + + vaults.add_key("kid-1", "deadbeef") + assert cross_vaults.keys == {"kid-1": "deadbeef"} + assert primary_vaults.keys == {} + + +def test_apply_cross_tracks_merges_audio_and_selection_sees_it() -> None: + d = make_dl() + title = Episode(id_="primary-ep-merge", service=Svc, title="Show", season=1, number=1, language="en") + title.tracks.audio.append(make_audio("prim-en", language="en")) + + cross_service = CrossSvc() + cross_tracks = Tracks() + cross_tracks.audio.append(make_audio("cross-ja", language="ja")) + sources = {"audio": (cross_service, object(), cross_tracks, object(), object())} + + d.apply_cross_tracks(title, sources) + + assert sorted(str(a.language) for a in title.tracks.audio) == ["en", "ja"] + assert any(a.id == "prim-en" for a in title.tracks.audio) # primary retained, not replaced + picked = Tracks.by_language(title.tracks.audio, ["ja"]) + assert picked and picked[0].id == "cross-ja" + + +def test_apply_cross_tracks_namespaces_colliding_ids() -> None: + d = make_dl() + title = Episode(id_="primary-ep-clash", service=Svc, title="Show", season=1, number=1, language="en") + title.tracks.audio.append(make_audio("dup", language="en")) + + cross_service = CrossSvc() + cross_tracks = Tracks() + cross_tracks.audio.append(make_audio("dup", language="ja")) # explicit id shared with the primary's + sources = {"audio": (cross_service, object(), cross_tracks, object(), object())} + + d.apply_cross_tracks(title, sources) + + assert sorted(a.id for a in title.tracks.audio) == ["dup", "dup_crosssvc"] + # The cross track is routed by its namespaced id; the primary track is not recorded as cross. + assert set(d.cross_track_sources) == {"dup_crosssvc"} + + +def test_apply_cross_tracks_namespaces_video_against_primary_ids() -> None: + # A cross video sharing an id with a surviving primary track must not shadow that track's + # drm_source entry, or the primary track would license through the cross service. + d = make_dl() + title = Episode(id_="primary-ep-vclash", service=Svc, title="Show", season=1, number=1, language="en") + title.tracks.audio.append(make_audio("dup", language="en")) + + cross_tracks = Tracks() + cross_tracks.videos.append(make_video("dup")) + sources = {"video": (CrossSvc(), object(), cross_tracks, object(), object())} + + d.apply_cross_tracks(title, sources) + + assert [v.id for v in title.tracks.videos] == ["dup_crosssvc"] + assert set(d.cross_track_sources) == {"dup_crosssvc"} + prim = title.tracks.audio[0] + assert d.drm_source(prim, "PRIMARY", title) == ("PRIMARY", title, d.cdm, d.vaults) + + +def test_fetch_cross_tracks_clears_stale_sources(monkeypatch: pytest.MonkeyPatch) -> None: + # Titles are processed one after another; an id recorded for the previous episode must not + # linger and shadow one of the next episode's own tracks. + import unshackle.commands.dl as dl_mod + + monkeypatch.setattr(dl_mod.Services, "get_tag", lambda name: name) + + d = make_dl(cross_audio=("CROSS", "https://cross/1")) + d.cross_track_sources["stale"] = (object(), object(), object(), object()) + monkeypatch.setattr(d, "load_cross_service", lambda tag, url: (CrossSvc(), object(), object())) + + title = Episode(id_="primary-ep-stale", service=Svc, title="Show", season=1, number=2, language="en") + assert d.fetch_cross_tracks(title) == {} # cross catalogue has no S01E02 + assert d.cross_track_sources == {} + + +def test_apply_cross_tracks_replaces_video() -> None: + d = make_dl() + title = Episode(id_="primary-ep-video", service=Svc, title="Show", season=1, number=1, language="en") + title.tracks.videos.append(make_video("prim-v")) + + cross_tracks = Tracks() + cross_tracks.videos.append(make_video("cross-v")) + sources = {"video": (CrossSvc(), object(), cross_tracks, object(), object())} + + d.apply_cross_tracks(title, sources) + + assert [v.id for v in title.tracks.videos] == ["cross-v"] + + +def test_apply_cross_tracks_records_offset() -> None: + d = make_dl(cross_audio_offset=7500) + title = Episode(id_="primary-ep-offset", service=Svc, title="Show", season=1, number=1, language="en") + + cross_service = CrossSvc() + cross_tracks = Tracks() + cross_tracks.audio.append(make_audio("cross-ja", language="ja")) + sources = {"audio": (cross_service, object(), cross_tracks, object(), object())} + + d.apply_cross_tracks(title, sources) + + cross = next(a for a in title.tracks.audio if a.id == "cross-ja") + assert cross.data["sync_offset_ms"] == 7500 + + +def test_audio_offset_emits_sync(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + captured = capture_mux_command(monkeypatch) + track = make_audio("a-off", offset=10000, path=tmp_path / "a.m4a") + Tracks(track).mux("Title", delete=False) + + cmd = captured["cmd"] + assert "--sync" in cmd + assert cmd[cmd.index("--sync") + 1] == "0:10000" + + +def test_subtitle_offset_emits_sync(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + captured = capture_mux_command(monkeypatch) + track = make_subtitle("s-off", offset=-2500, path=tmp_path / "s.vtt") + Tracks(track).mux("Title", delete=False) + + cmd = captured["cmd"] + assert "--sync" in cmd + assert cmd[cmd.index("--sync") + 1] == "0:-2500" + + +def test_no_offset_no_sync(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + captured = capture_mux_command(monkeypatch) + track = make_audio("a-plain", path=tmp_path / "a.m4a") + Tracks(track).mux("Title", delete=False) + + assert "--sync" not in captured["cmd"] + + +def test_match_cross_title_per_type_override_and_fallback() -> None: + primary = Episode(id_="primary-ep-match", service=Svc, title="Show", season=1, number=1, language="en") + catalogue = [ + Episode(id_="cross-ep-0101", service=Svc, title="Show", season=1, number=1, language="en"), + Episode(id_="cross-ep-0205", service=Svc, title="Show", season=2, number=5, language="en"), + ] + + # Default: the same season and episode as the title being downloaded. + assert make_dl().match_cross_title(primary, catalogue, "audio").id == "cross-ep-0101" + + # Per-type override wins over the default. + per_type = {"video": None, "audio": "S02E05", "subtitles": None, "chapters": None} + assert make_dl(cross_wanted_by_type=per_type).match_cross_title(primary, catalogue, "audio").id == "cross-ep-0205" + + # A type with no per-type override falls back to the shared --cross-wanted. + shared = make_dl(cross_wanted="S02E05") + assert shared.match_cross_title(primary, catalogue, "video").id == "cross-ep-0205" + + # Per-type override beats the shared fallback for its own type. + mixed = make_dl( + cross_wanted="S02E05", + cross_wanted_by_type={"video": None, "audio": "S01E01", "subtitles": None, "chapters": None}, + ) + assert mixed.match_cross_title(primary, catalogue, "audio").id == "cross-ep-0101" + + +def test_match_cross_title_rejects_bad_wanted() -> None: + primary = Episode(id_="primary-ep-bad", service=Svc, title="Show", season=1, number=1, language="en") + with pytest.raises(click.UsageError): + make_dl(cross_wanted="episode two").match_cross_title(primary, [], "audio") + + +def test_cross_cli_options_parse() -> None: + command = click.Command("dl", params=list(dl.cli.params), callback=lambda **_: None) + ctx = command.make_context( + "dl", + [ + "--cross-proxy", + "http://p:8080", + "--cross-audio-wanted", + "S02E05", + "--cross-audio-offset", + "10s", + "--cross-video", + "SVC", + "https://example.test/1", + ], + ) + assert ctx.params["cross_proxy"] == "http://p:8080" + assert ctx.params["cross_audio_wanted"] == "S02E05" + assert ctx.params["cross_audio_offset"] == 10000 + assert ctx.params["cross_video"] == ("SVC", "https://example.test/1") + + +def test_cross_mux_end_to_end(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + import unshackle.commands.dl as dl_mod + + monkeypatch.setattr(dl_mod.Services, "get_tag", lambda name: name) + + primary = Episode(id_="primary-ep-e2e", service=Svc, title="Show", season=1, number=1, language="en") + primary.tracks.audio.append(make_audio("prim-en", language="en", path=tmp_path / "prim_en.m4a")) + + cross_episode = Episode(id_="cross-ep-e2e", service=Svc, title="Show", season=1, number=1, language="ja") + cross_tracks = Tracks() + cross_tracks.audio.append(make_audio("cross-ja", language="ja", path=tmp_path / "cross_ja.m4a")) + cross_service = CrossSvc(titles=[cross_episode], tracks=cross_tracks) + cross_cdm = object() + cross_vaults = object() + + d = make_dl(cross_audio=("CROSS", "https://cross/1"), cross_audio_offset=5000) + monkeypatch.setattr(d, "load_cross_service", lambda tag, url: (cross_service, cross_cdm, cross_vaults)) + + sources = d.fetch_cross_tracks(primary) + assert set(sources) == {"audio"} + + d.apply_cross_tracks(primary, sources) + + # Merged, not replaced: both languages present. + assert {str(a.language) for a in primary.tracks.audio} == {"en", "ja"} + + # DRM routes the cross track through the cross service's own service/cdm/vaults. + cross_track = next(a for a in primary.tracks.audio if str(a.language) == "ja") + svc, ttl, cdm, vaults = d.drm_source(cross_track, SimpleNamespace(session=object()), primary) + assert svc is cross_service + assert ttl is cross_episode + assert cdm is cross_cdm + assert vaults is cross_vaults + + # The primary track still routes to the primary's own service/cdm/vaults. + prim_track = next(a for a in primary.tracks.audio if str(a.language) == "en") + assert d.drm_source(prim_track, "PRIMARY", primary) == ("PRIMARY", primary, d.cdm, d.vaults) + + # The built mux command carries the cross audio's offset. + captured = capture_mux_command(monkeypatch) + primary.tracks.mux("Show S01E01", delete=False) + cmd = captured["cmd"] + assert "--sync" in cmd + assert cmd[cmd.index("--sync") + 1] == "0:5000" diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index c26dd713..c07bb060 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -92,6 +92,7 @@ from unshackle.core.utils.bitrate import apply_real_bitrates from unshackle.core.utils.click_types import ( AUDIO_CODEC_LIST, LANGUAGE_RANGE, + OFFSET, QUALITY_LIST, SEASON_RANGE, SLOW_DELAY_RANGE, @@ -757,6 +758,93 @@ class dl: default=False, help="Continue with best available quality if requested resolutions are not available.", ) + @click.option( + "--cross-video", + nargs=2, + type=(str, str), + default=None, + metavar="SERVICE URL", + help="Take video tracks from another service instead of this one.", + ) + @click.option( + "--cross-audio", + nargs=2, + type=(str, str), + default=None, + metavar="SERVICE URL", + help="Take audio tracks from another service instead of this one.", + ) + @click.option( + "--cross-subtitles", + nargs=2, + type=(str, str), + default=None, + metavar="SERVICE URL", + help="Take subtitle tracks from another service instead of this one.", + ) + @click.option( + "--cross-chapters", + nargs=2, + type=(str, str), + default=None, + metavar="SERVICE URL", + help="Take chapters from another service instead of this one.", + ) + @click.option( + "--cross-audio-offset", + type=OFFSET, + default=None, + help="Shift cross-sourced audio in the mux, e.g. 10s, 500ms or -5.5s.", + ) + @click.option( + "--cross-subtitle-offset", + type=OFFSET, + default=None, + help="Shift cross-sourced subtitles in the mux, e.g. 10s, 500ms or -5.5s.", + ) + @click.option( + "--cross-profile", + type=str, + default=None, + help="Profile used for cross-service credentials. Defaults to --profile.", + ) + @click.option( + "--cross-proxy", + type=str, + default=None, + help="Proxy for every cross-service. Overrides each cross-service's own proxy configuration.", + ) + @click.option( + "--cross-video-wanted", + type=str, + default=None, + help="Episode to take video from the cross-service, e.g. S01E02. Overrides --cross-wanted for video.", + ) + @click.option( + "--cross-audio-wanted", + type=str, + default=None, + help="Episode to take audio from the cross-service, e.g. S01E02. Overrides --cross-wanted for audio.", + ) + @click.option( + "--cross-subtitles-wanted", + type=str, + default=None, + help="Episode to take subtitles from the cross-service, e.g. S01E02. Overrides --cross-wanted for subtitles.", + ) + @click.option( + "--cross-chapters-wanted", + type=str, + default=None, + help="Episode to take chapters from the cross-service, e.g. S01E02. Overrides --cross-wanted for chapters.", + ) + @click.option( + "--cross-wanted", + type=str, + default=None, + help="Episode to take from the cross-services, e.g. S01E02. Default fallback when no per-type override " + "is given; otherwise the matching season and episode.", + ) @click.option( "--remote", is_flag=True, @@ -795,6 +883,19 @@ class dl: animeapi_id: Optional[str] = None, enrich: bool = False, output_dir: Optional[Path] = None, + cross_video: Optional[tuple[str, str]] = None, + cross_audio: Optional[tuple[str, str]] = None, + cross_subtitles: Optional[tuple[str, str]] = None, + cross_chapters: Optional[tuple[str, str]] = None, + cross_audio_offset: Optional[int] = None, + cross_subtitle_offset: Optional[int] = None, + cross_profile: Optional[str] = None, + cross_proxy: Optional[str] = None, + cross_video_wanted: Optional[str] = None, + cross_audio_wanted: Optional[str] = None, + cross_subtitles_wanted: Optional[str] = None, + cross_chapters_wanted: Optional[str] = None, + cross_wanted: Optional[str] = None, *_: Any, **__: Any, ): @@ -848,6 +949,24 @@ class dl: self.animeapi_title: Optional[str] = None self.output_dir = output_dir + self.cross_video = ctx.params.get("cross_video", cross_video) + self.cross_audio = ctx.params.get("cross_audio", cross_audio) + self.cross_subtitles = ctx.params.get("cross_subtitles", cross_subtitles) + self.cross_chapters = ctx.params.get("cross_chapters", cross_chapters) + self.cross_audio_offset = ctx.params.get("cross_audio_offset", cross_audio_offset) + self.cross_subtitle_offset = ctx.params.get("cross_subtitle_offset", cross_subtitle_offset) + self.cross_profile = ctx.params.get("cross_profile", cross_profile) or profile + self.cross_proxy = ctx.params.get("cross_proxy", cross_proxy) + self.cross_wanted = ctx.params.get("cross_wanted", cross_wanted) + self.cross_wanted_by_type: dict[str, Optional[str]] = { + "video": ctx.params.get("cross_video_wanted", cross_video_wanted), + "audio": ctx.params.get("cross_audio_wanted", cross_audio_wanted), + "subtitles": ctx.params.get("cross_subtitles_wanted", cross_subtitles_wanted), + "chapters": ctx.params.get("cross_chapters_wanted", cross_chapters_wanted), + } + # Held outside track.data because the values are live objects, not exportable metadata. + self.cross_track_sources: dict[str, tuple[Service, Title_T, Any, Vaults]] = {} + if animeapi_id: from unshackle.core.utils.animeapi import resolve_animeapi @@ -1250,6 +1369,10 @@ class dl: config=self.service_config, cdm=self.cdm, proxy_providers=self.proxy_providers, profile=self.profile ) + # A service reads dl-level options off ctx.parent.params, so a cross-service is handed a + # copy of ours rather than a hand-written set of defaults that drifts as options are added. + self.dl_params: dict[str, Any] = dict(ctx.params) + if repack: config.repack = True @@ -1261,6 +1384,195 @@ class dl: # able to keep `self` as the first positional self.cli._result_callback = self.result + @property + def has_cross_mux(self) -> bool: + return any((self.cross_video, self.cross_audio, self.cross_subtitles, self.cross_chapters)) + + def build_vaults(self, vault_tag: str) -> Vaults: + """Build and load a Vaults keyed by a service tag, so its keys cache under that owner.""" + vaults = Vaults(vault_tag) + for vault in config.key_vaults: + vault_type = vault["type"] + vault_name = vault.get("name", vault_type) + vault_copy = vault.copy() + del vault_copy["type"] + + if vault_type.lower() == "api" and "decrypt_labs" in vault_name.lower(): + if not vault_copy.get("token") and config.decrypt_labs_api_key: + vault_copy["token"] = config.decrypt_labs_api_key + + if vault_type.lower() == "sqlite": + vaults.load_critical(vault_type, **vault_copy) + else: + vaults.load(vault_type, **vault_copy) + return vaults + + def load_cross_service(self, tag: str, url: str) -> tuple[Service, Any, Vaults]: + """Load and authenticate another service plus its own CDM and vaults, for muxing its tracks.""" + service_cls: Any = Services.load(tag) + + service_config_path = Services.get_path(tag) / config.filenames.config + service_config = ( + yaml.safe_load(service_config_path.read_text(encoding="utf8")) if service_config_path.exists() else {} + ) + + # Stand in for the dl group Click would have invoked the service under. + parent_ctx = click.Context(self.cli, info_name="dl") + parent_params = {**self.dl_params, "profile": self.cross_profile} + if self.cross_proxy: + parent_params["proxy"] = self.cross_proxy + parent_params["no_proxy"] = False + parent_ctx.params = parent_params + + # Keys licensed for this service cache under its own vault, not the primary's; its CDM binds + # to those same vaults so lookups and writes stay with the owning service. + cross_vaults = self.build_vaults(Services.get_vault_tag(tag)) + cross_cdm: Any = self.get_cdm(tag, self.cross_profile, vaults=cross_vaults) + + ctx = click.Context(service_cls.cli, parent=parent_ctx, info_name=tag) + ctx.obj = ContextData( + config=service_config, + cdm=cross_cdm, + proxy_providers=self.proxy_providers, + profile=self.cross_profile, + ) + + kwargs = {param.name: param.default for param in service_cls.cli.params if param.name} + kwargs["title"] = url + + service: Service = service_cls(ctx, **kwargs) + service.authenticate( + self.get_cookie_jar(tag, self.cross_profile), self.get_credentials(tag, self.cross_profile) + ) + return service, cross_cdm, cross_vaults + + def match_cross_title(self, title: Title_T, cross_titles: Iterable[Title_T], track_type: str) -> Optional[Title_T]: + """Find the counterpart of a title in another service's catalogue for a given track type.""" + if isinstance(title, Movie): + return next((x for x in cross_titles if isinstance(x, Movie)), None) + + if not isinstance(title, Episode): + return None + + season, number = title.season, title.number + # Per-type override wins; --cross-wanted is the shared fallback; else the matching S/E. + wanted = self.cross_wanted_by_type.get(track_type) or self.cross_wanted + if wanted: + match = re.match(r"S(\d+)E(\d+)$", wanted, re.IGNORECASE) + if not match: + raise click.UsageError(f"cross-wanted must look like S01E02, not {wanted!r}") + season, number = int(match.group(1)), int(match.group(2)) + + return next( + (x for x in cross_titles if isinstance(x, Episode) and x.season == season and x.number == number), None + ) + + def fetch_cross_tracks(self, title: Title_T) -> dict[str, tuple[Service, Title_T, Tracks, Any, Vaults]]: + """Resolve every --cross-* option to the service, title, tracks, CDM and vaults it names.""" + specs = ( + ("video", self.cross_video), + ("audio", self.cross_audio), + ("subtitles", self.cross_subtitles), + ("chapters", self.cross_chapters), + ) + + # Sources are per-title; a stale id recorded for a previous episode must not shadow one of + # this episode's own tracks in drm_source. + self.cross_track_sources.clear() + + # One service may supply several track types; load and authenticate it only once, and reuse + # a matched title's tracks (a per-type --cross-*-wanted may point types at different titles). + loaded: dict[tuple[str, str], tuple[Service, Any, Vaults, list[Title_T]]] = {} + track_pools: dict[int, Tracks] = {} + sources: dict[str, tuple[Service, Title_T, Tracks, Any, Vaults]] = {} + + for track_type, spec in specs: + if not spec: + continue + + key = (Services.get_tag(spec[0]), spec[1]) + if key not in loaded: + self.log.info(f"Cross-mux: loading {key[0]}") + service, cdm, vaults = self.load_cross_service(*key) + loaded[key] = (service, cdm, vaults, list(service.get_titles())) + service, cdm, vaults, titles = loaded[key] + + cross_title = self.match_cross_title(title, titles, track_type) + if not cross_title: + self.log.warning(f"Cross-mux: {key[0]} has no title matching {title}, keeping this service's own.") + continue + + if id(cross_title) not in track_pools: + self.log.info(f"Cross-mux: getting {track_type} from {key[0]}") + tracks = service.get_tracks(cross_title) + tracks.chapters = service.get_chapters(cross_title) + track_pools[id(cross_title)] = tracks + + sources[track_type] = (service, cross_title, track_pools[id(cross_title)], cdm, vaults) + + return sources + + def apply_cross_tracks( + self, title: Title_T, sources: dict[str, tuple[Service, Title_T, Tracks, Any, Vaults]] + ) -> None: + """Merge cross-sourced audio/subtitles into this title (video/chapters replace) and record sources.""" + + def adopt( + tracks: Iterable[AnyTrack], + service: Service, + cross_title: Title_T, + cdm: Any, + vaults: Vaults, + offset: Optional[int], + ) -> None: + for track in tracks: + # Track.download() already prefers track.session over the primary service's. + track.session = service.session + self.cross_track_sources[track.id] = (service, cross_title, cdm, vaults) + track.data["cross_source"] = service.__class__.__name__ + if offset: + track.data["sync_offset_ms"] = offset + + # Ids must stay unique across every track kept for this title, or a cross track's entry in + # cross_track_sources would shadow a primary track's and route its DRM to the wrong service. + seen_ids = {t.id for t in (*title.tracks.audio, *title.tracks.subtitles)} + + def namespace(tracks: Iterable[AnyTrack], service: Service) -> None: + for track in tracks: + if track.id in seen_ids: + track.id = f"{track.id}_{service.__class__.__name__.lower()}" + seen_ids.add(track.id) + + # Exactly one video source and chapters are singular, so those replace outright. + if "video" in sources: + service, cross_title, cross_tracks, cdm, vaults = sources["video"] + namespace(cross_tracks.videos, service) + title.tracks.videos = cross_tracks.videos + adopt(title.tracks.videos, service, cross_title, cdm, vaults, None) + + # Audio and subtitles merge so best-per-language selection runs across the primary's own + # tracks and the cross-sourced ones together. + seen_ids.update(t.id for t in title.tracks.videos) + + if "audio" in sources: + service, cross_title, cross_tracks, cdm, vaults = sources["audio"] + namespace(cross_tracks.audio, service) + adopt(cross_tracks.audio, service, cross_title, cdm, vaults, self.cross_audio_offset) + title.tracks.audio.extend(cross_tracks.audio) + + if "subtitles" in sources: + service, cross_title, cross_tracks, cdm, vaults = sources["subtitles"] + namespace(cross_tracks.subtitles, service) + adopt(cross_tracks.subtitles, service, cross_title, cdm, vaults, self.cross_subtitle_offset) + title.tracks.subtitles.extend(cross_tracks.subtitles) + + if "chapters" in sources: + title.tracks.chapters = sources["chapters"][2].chapters + + def drm_source(self, track: AnyTrack, service: Service, title: Title_T) -> tuple[Service, Title_T, Any, Vaults]: + """The service, title, CDM and vaults that licence a track; the primary's unless cross-sourced.""" + return self.cross_track_sources.get(track.id, (service, title, self.cdm, self.vaults)) + @with_task_temp def result( self, @@ -2360,6 +2672,22 @@ class dl: level="INFO", operation="get_tracks", service=self.service, context=tracks_info ) + if self.has_cross_mux: + with console.status("Getting Cross-Service Tracks...", spinner="dots"): + try: + cross_sources = self.fetch_cross_tracks(title) + except Exception as e: + if self.debug_logger: + self.debug_logger.log_error( + "cross_mux", e, service=self.service, context={"title": str(title)} + ) + raise + + if cross_sources: + self.apply_cross_tracks(title, cross_sources) + applied = ", ".join(f"{k} from {v[0].__class__.__name__}" for k, v in cross_sources.items()) + self.log.info(f"Cross-mux: using {applied}") + # strip SDH subs to non-SDH if no equivalent same-lang non-SDH is available # uses a loose check, e.g, wont strip en-US SDH sub if a non-SDH en-GB is available # Check if automatic SDH stripping is enabled in config @@ -3040,36 +3368,39 @@ class dl: with SyncLive(Padding(download_table, (1, 5)), console=console, refresh_per_second=20): def download_track(track: AnyTrack, i: int) -> None: + drm_service, drm_title, drm_cdm, drm_vaults = self.drm_source(track, service, title) track.download( session=track.session or service.session, no_proxy_download=no_proxy_download, prepare_drm=partial( partial(self.prepare_drm, table=download_table), track=track, - title=title, + title=drm_title, certificate=partial( - service.get_widevine_service_certificate, - title=title, + drm_service.get_widevine_service_certificate, + title=drm_title, track=track, ), licence=partial( - service.get_playready_license - if is_playready_cdm(self.cdm) - else service.get_widevine_license, - title=title, + drm_service.get_playready_license + if is_playready_cdm(drm_cdm) + else drm_service.get_widevine_license, + title=drm_title, track=track, ), clearkey_licence=partial( - service.get_clearkey_license, - title=title, + drm_service.get_clearkey_license, + title=drm_title, track=track, ), cdm_only=cdm_only, vaults_only=vaults_only, export=export_path, - service_session=service.session, + service_session=drm_service.session, + cdm=drm_cdm, + vaults=drm_vaults, ), - cdm=self.cdm, + cdm=drm_cdm, max_workers=workers, progress=tracks_progress_callables[i], ) @@ -3909,6 +4240,8 @@ class dl: vaults_only: bool = False, export: Optional[Path] = None, service_session: Optional[Any] = None, + cdm: Optional[object] = None, + vaults: Optional[Vaults] = None, ) -> None: """ Prepare the DRM by getting decryption data like KIDs, Keys, and such. @@ -3917,6 +4250,9 @@ class dl: if not drm: return + cdm = cdm if cdm is not None else self.cdm + vaults = vaults if vaults is not None else self.vaults + server_cdm = getattr(self, "server_cdm", False) if server_cdm: @@ -3977,26 +4313,38 @@ class dl: if not server_cdm: if isinstance(drm, Widevine): - if not is_widevine_cdm(self.cdm): - widevine_cdm = self.get_cdm(self.service, self.profile, drm="widevine", quality=track_quality) + if not is_widevine_cdm(cdm): + widevine_cdm = self.get_cdm( + self.service, self.profile, drm="widevine", quality=track_quality, vaults=vaults + ) if widevine_cdm and is_widevine_cdm(widevine_cdm): if track_quality: self.log.info(f"Switching to Widevine CDM for Widevine {track_quality}p content") else: self.log.info("Switching to Widevine CDM for Widevine content") - self.cdm = widevine_cdm + if cdm is self.cdm: + # Persist the switch only for the primary CDM so its later tracks skip the + # re-load; a cross-sourced CDM stays local and must not clobber the primary's. + self.cdm = widevine_cdm + cdm = widevine_cdm else: raise ValueError(f"Title needs a Widevine CDM but {self.service} is configured with PlayReady.") elif isinstance(drm, PlayReady): - if not is_playready_cdm(self.cdm): - playready_cdm = self.get_cdm(self.service, self.profile, drm="playready", quality=track_quality) + if not is_playready_cdm(cdm): + playready_cdm = self.get_cdm( + self.service, self.profile, drm="playready", quality=track_quality, vaults=vaults + ) if playready_cdm and is_playready_cdm(playready_cdm): if track_quality: self.log.info(f"Switching to PlayReady CDM for PlayReady {track_quality}p content") else: self.log.info("Switching to PlayReady CDM for PlayReady content") - self.cdm = playready_cdm + if cdm is self.cdm: + # Persist the switch only for the primary CDM so its later tracks skip the + # re-load; a cross-sourced CDM stays local and must not clobber the primary's. + self.cdm = playready_cdm + cdm = playready_cdm else: raise ValueError(f"Title needs a PlayReady CDM but {self.service} is configured with Widevine.") @@ -4061,15 +4409,13 @@ class dl: continue if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) + content_key, vault_used = vaults.get_key(kid) if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): cek_tree.add(label) - pending_vault_writes.append( - partial(self.vaults.add_key, kid, content_key, excluding=vault_used) - ) + pending_vault_writes.append(partial(vaults.add_key, kid, content_key, excluding=vault_used)) self.LICENSE_KEY_CACHE[kid] = content_key if self.debug_logger: @@ -4123,9 +4469,9 @@ class dl: try: if self.service == "NF": - drm.get_NF_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_NF_content_keys(cdm=cdm, licence=licence, certificate=certificate) else: - drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_content_keys(cdm=cdm, licence=licence, certificate=certificate) except Exception as e: if drm.content_keys: self.log.debug(f"License call failed but keys already in content_keys: {e}") @@ -4262,15 +4608,13 @@ class dl: continue if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) + content_key, vault_used = vaults.get_key(kid) if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): cek_tree.add(label) - pending_vault_writes.append( - partial(self.vaults.add_key, kid, content_key, excluding=vault_used) - ) + pending_vault_writes.append(partial(vaults.add_key, kid, content_key, excluding=vault_used)) self.LICENSE_KEY_CACHE[kid] = content_key if self.debug_logger: @@ -4313,7 +4657,7 @@ class dl: from_vaults = drm.content_keys.copy() try: - drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_content_keys(cdm=cdm, licence=licence, certificate=certificate) except Exception as e: if drm.content_keys: self.log.debug(f"License call failed but keys already in content_keys: {e}") @@ -4414,15 +4758,13 @@ class dl: continue if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) + content_key, vault_used = vaults.get_key(kid) if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): cek_tree.add(label) - pending_vault_writes.append( - partial(self.vaults.add_key, kid, content_key, excluding=vault_used) - ) + pending_vault_writes.append(partial(vaults.add_key, kid, content_key, excluding=vault_used)) self.LICENSE_KEY_CACHE[kid] = content_key elif vaults_only: msg = f"No Vault has a Key for {kid.hex} and --vaults-only was used" @@ -4612,6 +4954,7 @@ class dl: profile: Optional[str] = None, drm: Optional[str] = None, quality: Optional[int] = None, + vaults: Optional[Vaults] = None, ) -> Optional[object]: """ Get CDM for a specified service (either Local or Remote CDM). @@ -4704,4 +5047,4 @@ class dl: from unshackle.core.cdm import load_cdm - return load_cdm(cdm_name, service_name=service, vaults=self.vaults) + return load_cdm(cdm_name, service_name=service, vaults=vaults if vaults is not None else self.vaults) diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index faa5ce26..c3304f42 100644 --- a/unshackle/core/tracks/tracks.py +++ b/unshackle/core/tracks/tracks.py @@ -631,25 +631,25 @@ class Tracks: audio_default = i == preferred_audio_idx else: audio_default = at.is_original_lang - cl.extend( - [ - "--track-name", - f"0:{at.get_track_name() or ''}", - "--language", - f"0:{at.language}", - "--default-track", - f"0:{audio_default}", - "--visual-impaired-flag", - f"0:{at.descriptive}", - "--original-flag", - f"0:{at.is_original_lang}", - "--compression", - "0:none", # disable extra compression - "(", - str(at.path), - ")", - ] - ) + audio_args = [ + "--track-name", + f"0:{at.get_track_name() or ''}", + "--language", + f"0:{at.language}", + "--default-track", + f"0:{audio_default}", + "--visual-impaired-flag", + f"0:{at.descriptive}", + "--original-flag", + f"0:{at.is_original_lang}", + "--compression", + "0:none", # disable extra compression + ] + + if at.data.get("sync_offset_ms"): + audio_args.extend(["--sync", f"0:{at.data['sync_offset_ms']}"]) + + cl.extend(audio_args + ["(", str(at.path), ")"]) if not skip_subtitles: for i, st in enumerate(self.subtitles): @@ -660,29 +660,29 @@ class Tracks: default = i == preferred_subtitle_idx else: default = bool(self.audio and is_close_match(st.language, [self.audio[0].language]) and st.forced) - cl.extend( - [ - "--track-name", - f"0:{st.get_track_name() or ''}", - "--language", - f"0:{st.language}", - "--sub-charset", - "0:UTF-8", - "--forced-track", - f"0:{st.forced}", - "--default-track", - f"0:{default}", - "--hearing-impaired-flag", - f"0:{st.sdh}", - "--original-flag", - f"0:{st.is_original_lang}", - "--compression", - "0:none", # disable extra compression (probably zlib) - "(", - str(st.path), - ")", - ] - ) + subtitle_args = [ + "--track-name", + f"0:{st.get_track_name() or ''}", + "--language", + f"0:{st.language}", + "--sub-charset", + "0:UTF-8", + "--forced-track", + f"0:{st.forced}", + "--default-track", + f"0:{default}", + "--hearing-impaired-flag", + f"0:{st.sdh}", + "--original-flag", + f"0:{st.is_original_lang}", + "--compression", + "0:none", # disable extra compression (probably zlib) + ] + + if st.data.get("sync_offset_ms"): + subtitle_args.extend(["--sync", f"0:{st.data['sync_offset_ms']}"]) + + cl.extend(subtitle_args + ["(", str(st.path), ")"]) if self.chapters: chapters_path = config.directories.temp / config.filenames.chapters.format( diff --git a/unshackle/core/utils/click_types.py b/unshackle/core/utils/click_types.py index 9f0f4334..bbb09ac1 100644 --- a/unshackle/core/utils/click_types.py +++ b/unshackle/core/utils/click_types.py @@ -425,10 +425,30 @@ class SlowDelayRange(click.ParamType): return (low, high) +class OffsetType(click.ParamType): + """Parses a time offset like '10s', '500ms' or '-5.5s' into milliseconds. A bare number is milliseconds.""" + + name = "offset" + + _PATTERN = re.compile(r"^(-?\d+(?:\.\d+)?)\s*(s|ms)?$") + + def convert(self, value: Any, param: Optional[click.Parameter], ctx: Optional[click.Context]) -> int: + if isinstance(value, int): + return value + + match = self._PATTERN.match(str(value).strip()) + if not match: + self.fail(f"'{value}' is not a valid offset. Use format: 10s, 500ms or -5.5s", param, ctx) + + number = float(match.group(1)) + return int(number * 1000) if match.group(2) == "s" else int(number) + + SEASON_RANGE = SeasonRange() LANGUAGE_RANGE = LanguageRange() QUALITY_LIST = QualityList() AUDIO_CODEC_LIST = AudioCodecList(Audio.Codec) SLOW_DELAY_RANGE = SlowDelayRange() +OFFSET = OffsetType() # VIDEO_CODEC_CHOICE will be created dynamically when imported