mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
imSp4rky
bd64f3a8c4
Catch SystemExit at service import, in the heartbeat task, the 5xx middleware, search, samples and chapters, and answer it as SERVICE_ERROR.
229 lines
8.6 KiB
Python
229 lines
8.6 KiB
Python
"""List and search requests on a --remote-only server run on a per-request cache directory.
|
|
|
|
The directory holds only what the client sent, goes back to the client when it logged in
|
|
itself, and is always removed when the request ends.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import types
|
|
import zlib
|
|
|
|
import click
|
|
import pytest
|
|
|
|
from unshackle.core.api import handlers
|
|
from unshackle.core.api.stats import stats
|
|
from unshackle.core.cacher import Cacher
|
|
from unshackle.core.config import config
|
|
|
|
|
|
class _CacheService:
|
|
fail_titles = False
|
|
|
|
def __init__(self, ctx, title):
|
|
self.seen = _CacheService.seen
|
|
self.seen["parent_params"] = ctx.parent.params
|
|
self.cache = Cacher("EXAMPLE")
|
|
|
|
def authenticate(self, cookies, credential):
|
|
self.seen["seeded"] = self.cache.get("tokens").data
|
|
self.seen["cache_dir"] = self.cache.get("tokens").path.parent
|
|
self.cache.get("tokens").set("new")
|
|
self.cache.get("titles_abc").set("titles")
|
|
|
|
def get_titles(self):
|
|
if self.fail_titles:
|
|
raise RuntimeError("boom")
|
|
return []
|
|
|
|
def search(self):
|
|
return iter(())
|
|
|
|
|
|
@pytest.fixture
|
|
def seen(tmp_path, monkeypatch):
|
|
import types
|
|
|
|
_CacheService.seen = {}
|
|
_CacheService.fail_titles = False
|
|
|
|
@click.command()
|
|
@click.argument("title")
|
|
@click.pass_context
|
|
def cli(ctx, **kwargs):
|
|
return _CacheService(ctx, **kwargs)
|
|
|
|
module = types.SimpleNamespace(cli=cli)
|
|
monkeypatch.setattr(config.directories, "cache", tmp_path)
|
|
monkeypatch.setattr(stats, "mode", "remote_only")
|
|
monkeypatch.setattr(handlers, "validate_service", lambda tag, request=None: tag)
|
|
monkeypatch.setattr(handlers, "load_service_yaml", lambda service: {})
|
|
monkeypatch.setattr(handlers, "resolve_handler_proxy", lambda *args: (None, []))
|
|
monkeypatch.setattr(handlers, "resolve_server_account", lambda *args: None)
|
|
monkeypatch.setattr(handlers, "load_full_cdm", lambda *args: None)
|
|
monkeypatch.setattr(handlers.Services, "load", staticmethod(lambda service: module))
|
|
monkeypatch.setattr("unshackle.commands.dl.dl.get_cookie_jar", staticmethod(lambda service, profile: None))
|
|
monkeypatch.setattr("unshackle.commands.dl.dl.get_credentials", staticmethod(lambda service, profile: None))
|
|
return _CacheService.seen
|
|
|
|
|
|
def client_cache(tmp_path, value):
|
|
Cacher("SRC").get("tokens").set(value)
|
|
text = (tmp_path / "SRC" / "tokens.json").read_bytes()
|
|
return {"tokens": base64.b64encode(zlib.compress(text)).decode("ascii")}
|
|
|
|
|
|
def request_dirs(tmp_path):
|
|
return list((tmp_path / "_requests").rglob("*")) if (tmp_path / "_requests").exists() else []
|
|
|
|
|
|
CREDENTIALS = {"credentials": {"username": "a@example.com", "password": "pw"}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_seeded_cache_is_read_and_returned_then_removed(seen, tmp_path):
|
|
data = {"service": "EXAMPLE", "title_id": "t", "cache": client_cache(tmp_path, "old")}
|
|
response = await handlers.list_titles_handler(data)
|
|
body = json.loads(response.body)
|
|
|
|
assert seen["seeded"] == "old"
|
|
assert seen["cache_dir"].is_relative_to(tmp_path / "_requests")
|
|
assert set(body["cache"]) == {"tokens"}
|
|
returned = json.loads(zlib.decompress(base64.b64decode(body["cache"]["tokens"])))
|
|
assert returned["data"] == "new"
|
|
assert not seen["cache_dir"].exists()
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_returns_cache_to_a_client_with_credentials(seen, tmp_path):
|
|
response = await handlers.search_handler({"service": "EXAMPLE", "query": "q", **CREDENTIALS})
|
|
body = json.loads(response.body)
|
|
assert body["count"] == 0
|
|
assert set(body["cache"]) == {"tokens"}
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_cache_returned_without_client_login(seen, tmp_path):
|
|
response = await handlers.search_handler({"service": "EXAMPLE", "query": "q"})
|
|
assert "cache" not in json.loads(response.body)
|
|
assert seen["cache_dir"].is_relative_to(tmp_path / "_requests")
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_dir_removed_on_error(seen, tmp_path):
|
|
_CacheService.fail_titles = True
|
|
data = {"service": "EXAMPLE", "title_id": "t", **CREDENTIALS}
|
|
response = await handlers.list_tracks_handler(data)
|
|
assert response.status >= 400
|
|
assert "cache" not in json.loads(response.body)
|
|
assert seen["cache_dir"].is_relative_to(tmp_path / "_requests")
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_answers_a_service_sys_exit_with_an_error(seen, tmp_path, monkeypatch):
|
|
def search(self):
|
|
raise SystemExit(1)
|
|
|
|
monkeypatch.setattr(_CacheService, "search", search)
|
|
response = await handlers.search_handler({"service": "EXAMPLE", "query": "q"})
|
|
assert json.loads(response.body)["error_code"] == "SERVICE_ERROR"
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_client_cache_still_removes_the_dir(seen, tmp_path):
|
|
data = {"service": "EXAMPLE", "query": "q", "cache": {"tokens": 1}}
|
|
with pytest.raises(handlers.APIError):
|
|
await handlers.search_handler(data)
|
|
assert not request_dirs(tmp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_server_account_uses_its_own_cache_and_returns_none(seen, tmp_path, monkeypatch):
|
|
monkeypatch.setattr(handlers, "resolve_server_account", lambda *args: "acct")
|
|
monkeypatch.setattr(handlers, "server_account_cookies", lambda *args: None)
|
|
data = {"service": "EXAMPLE", "query": "q", "cache": client_cache(tmp_path, "old"), **CREDENTIALS}
|
|
response = await handlers.search_handler(data)
|
|
|
|
assert "cache" not in json.loads(response.body)
|
|
assert not seen["seeded"]
|
|
assert seen["cache_dir"] == tmp_path / "_accounts" / "EXAMPLE" / "acct"
|
|
assert not (tmp_path / "_requests").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_anonymous_search_skips_the_login_and_keeps_the_account_cache(seen, tmp_path, monkeypatch):
|
|
def search(self):
|
|
seen["search_cache_dir"] = self.cache.get("tokens").path.parent
|
|
return iter(())
|
|
|
|
def no_login_material(*args):
|
|
raise AssertionError("an anonymous search must not read login material")
|
|
|
|
handlers.Services.load("EXAMPLE").ANONYMOUS_SEARCH = True
|
|
monkeypatch.setattr(_CacheService, "search", search)
|
|
monkeypatch.setattr(handlers, "resolve_server_account", lambda *args: "acct")
|
|
monkeypatch.setattr(handlers, "server_login_material", no_login_material)
|
|
response = await handlers.search_handler({"service": "EXAMPLE", "query": "q", **CREDENTIALS})
|
|
|
|
assert json.loads(response.body)["results"] == []
|
|
assert "seeded" not in seen
|
|
assert seen["parent_params"]["cookies_supplied"] is False
|
|
assert seen["search_cache_dir"] == tmp_path / "_accounts" / "EXAMPLE" / "acct"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_mode_keeps_the_server_cache(seen, tmp_path, monkeypatch):
|
|
monkeypatch.setattr(stats, "mode", "full")
|
|
data = {"service": "EXAMPLE", "query": "q", "cache": client_cache(tmp_path, "old"), **CREDENTIALS}
|
|
response = await handlers.search_handler(data)
|
|
|
|
assert "cache" not in json.loads(response.body)
|
|
assert not seen["seeded"]
|
|
assert seen["cache_dir"] == tmp_path / "EXAMPLE"
|
|
assert not (tmp_path / "_requests").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tracks_service_sees_dl_params(seen):
|
|
from unshackle.core.tracks import Audio
|
|
|
|
data = {"service": "EXAMPLE", "title_id": "t", "dl_params": {"a_lang": ["es-419"], "acodec": ["ddp"]}}
|
|
await handlers.list_tracks_handler(data)
|
|
assert seen["parent_params"]["a_lang"] == ["es-419"]
|
|
assert seen["parent_params"]["acodec"] == [Audio.Codec.EC3]
|
|
assert "dl_params" in handlers.LIST_HANDLER_TRANSPORT_KEYS
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_service_gets_no_dl_params(seen):
|
|
await handlers.search_handler({"service": "EXAMPLE", "query": "q", "dl_params": {"a_lang": ["es-419"]}})
|
|
assert "a_lang" not in seen["parent_params"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_samples_answers_a_service_exit_in_decrypts_samples(monkeypatch):
|
|
class Service:
|
|
def decrypts_samples(self, track):
|
|
raise SystemExit(1)
|
|
|
|
session = types.SimpleNamespace(
|
|
tracks={"t": object()}, service_instance=Service(), service_tag="EXAMPLE", lock=asyncio.Lock()
|
|
)
|
|
|
|
async def get_session(session_id, request=None):
|
|
return session
|
|
|
|
monkeypatch.setattr(handlers, "get_validated_session", get_session)
|
|
monkeypatch.setattr(handlers, "require_authenticated", lambda s: None)
|
|
monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
|
|
with pytest.raises(handlers.APIError) as e:
|
|
await handlers.session_samples_handler("sid", "t", "1", b"")
|
|
assert e.value.error_code.value == "SERVICE_ERROR"
|