[ie/youtube] Add a PO Token Provider Framework (yt-dlp#12840)

https://github.com/yt-dlp/yt-dlp/tree/master/yt_dlp/extractor/youtube/pot/README.md

Authored by: coletdjnz

Author: coletdjnz

README.md

@@ -1795,6 +1795,7 @@ The following extractors use this feature:
 * `player_client`: Clients to extract video data from. The currently available clients are `web`, `web_safari`, `web_embedded`, `web_music`, `web_creator`, `mweb`, `ios`, `android`, `android_vr`, `tv` and `tv_embedded`. By default, `tv,ios,web` is used, or `tv,web` is used when authenticating with cookies. The `web_music` client is added for `music.youtube.com` URLs when logged-in cookies are used. The `web_embedded` client is added for age-restricted videos but only works if the video is embeddable. The `tv_embedded` and `web_creator` clients are added for age-restricted videos if account age-verification is required. Some clients, such as `web` and `web_music`, require a `po_token` for their formats to be downloadable. Some clients, such as `web_creator`, will only work with authentication. Not all clients support authentication via cookies. You can use `default` for the default clients, or you can use `all` for all clients (not recommended). You can prefix a client with `-` to exclude it, e.g. `youtube:player_client=default,-ios`
 * `player_skip`: Skip some network requests that are generally needed for robust extraction. One or more of `configs` (skip client configs), `webpage` (skip initial webpage), `js` (skip js player), `initial_data` (skip initial data/next ep request). While these options can help reduce the number of requests needed or avoid some rate-limiting, they could cause issues such as missing formats or metadata.  See [#860](https://github.com/yt-dlp/yt-dlp/pull/860) and [#12826](https://github.com/yt-dlp/yt-dlp/issues/12826) for more details
 * `player_params`: YouTube player parameters to use for player requests. Will overwrite any default ones set by yt-dlp.
+* `player_js_variant`: The player javascript variant to use for signature and nsig deciphering. The known variants are: `main`, `tce`, `tv`, `tv_es6`, `phone`, `tablet`. Only `main` is recommended as a possible workaround; the others are for debugging purposes. The default is to use what is prescribed by the site, and can be selected with `actual`
 * `comment_sort`: `top` or `new` (default) - choose comment sorting mode (on YouTube's side)
 * `max_comments`: Limit the amount of comments to gather. Comma-separated list of integers representing `max-comments,max-parents,max-replies,max-replies-per-thread`. Default is `all,all,all,all`
     * E.g. `all,all,1000,10` will get a maximum of 1000 replies total, with up to 10 replies per thread. `1000,all,100` will get a maximum of 1000 comments, with a maximum of 100 replies total
@@ -1805,7 +1806,11 @@ The following extractors use this feature:
 * `data_sync_id`: Overrides the account Data Sync ID used in Innertube API requests. This may be needed if you are using an account with `youtube:player_skip=webpage,configs` or `youtubetab:skip=webpage`
 * `visitor_data`: Overrides the Visitor Data used in Innertube API requests. This should be used with `player_skip=webpage,configs` and without cookies. Note: this may have adverse effects if used improperly. If a session from a browser is wanted, you should pass cookies instead (which contain the Visitor ID)
 * `po_token`:  Proof of Origin (PO) Token(s) to use. Comma seperated list of PO Tokens in the format `CLIENT.CONTEXT+PO_TOKEN`, e.g. `youtube:po_token=web.gvs+XXX,web.player=XXX,web_safari.gvs+YYY`. Context can be either `gvs` (Google Video Server URLs) or `player` (Innertube player request)
-* `player_js_variant`: The player javascript variant to use for signature and nsig deciphering. The known variants are: `main`, `tce`, `tv`, `tv_es6`, `phone`, `tablet`. Only `main` is recommended as a possible workaround; the others are for debugging purposes. The default is to use what is prescribed by the site, and can be selected with `actual`
+* `pot_trace`: Enable debug logging for PO Token fetching. Either `true` or `false` (default)
+* `fetch_pot`: Policy to use for fetching a PO Token from providers. One of `always` (always try fetch a PO Token regardless if the client requires one for the given context), `never` (never fetch a PO Token), or `auto` (default; only fetch a PO Token if the client requires one for the given context)
+
+#### youtubepot-webpo
+* `bind_to_visitor_id`: Whether to use the Visitor ID instead of Visitor Data for caching WebPO tokens. Either `true` (default) or `false`
 
 #### youtubetab (YouTube playlists, channels, feeds, etc.)
 * `skip`: One or more of `webpage` (skip initial webpage download), `authcheck` (allow the download of playlists requiring authentication when no initial webpage is downloaded. This may cause unwanted behavior, see [#1122](https://github.com/yt-dlp/yt-dlp/pull/1122) for more details)

test/test_YoutubeDL.py

@@ -1435,6 +1435,27 @@ def test_load_plugins_compat(self):
         FakeYDL().close()
         assert all_plugins_loaded.value
 
+    def test_close_hooks(self):
+        # Should call all registered close hooks on close
+        close_hook_called = False
+        close_hook_two_called = False
+
+        def close_hook():
+            nonlocal close_hook_called
+            close_hook_called = True
+
+        def close_hook_two():
+            nonlocal close_hook_two_called
+            close_hook_two_called = True
+
+        ydl = FakeYDL()
+        ydl.add_close_hook(close_hook)
+        ydl.add_close_hook(close_hook_two)
+
+        ydl.close()
+        self.assertTrue(close_hook_called, 'Close hook was not called')
+        self.assertTrue(close_hook_two_called, 'Close hook two was not called')
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_networking_utils.py

@@ -20,15 +20,14 @@
     add_accept_encoding_header,
     get_redirect_method,
     make_socks_proxy_opts,
-    select_proxy,
     ssl_load_certs,
 )
 from yt_dlp.networking.exceptions import (
     HTTPError,
     IncompleteRead,
 )
 from yt_dlp.socks import ProxyType
-from yt_dlp.utils.networking import HTTPHeaderDict
+from yt_dlp.utils.networking import HTTPHeaderDict, select_proxy
 
 TEST_DIR = os.path.dirname(os.path.abspath(__file__))
 

test/test_pot/conftest.py

@@ -0,0 +1,71 @@
+import collections
+
+import pytest
+
+from yt_dlp import YoutubeDL
+from yt_dlp.cookies import YoutubeDLCookieJar
+from yt_dlp.extractor.common import InfoExtractor
+from yt_dlp.extractor.youtube.pot._provider import IEContentProviderLogger
+from yt_dlp.extractor.youtube.pot.provider import PoTokenRequest, PoTokenContext
+from yt_dlp.utils.networking import HTTPHeaderDict
+
+
+class MockLogger(IEContentProviderLogger):
+
+    log_level = IEContentProviderLogger.LogLevel.TRACE
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.messages = collections.defaultdict(list)
+
+    def trace(self, message: str):
+        self.messages['trace'].append(message)
+
+    def debug(self, message: str):
+        self.messages['debug'].append(message)
+
+    def info(self, message: str):
+        self.messages['info'].append(message)
+
+    def warning(self, message: str, *, once=False):
+        self.messages['warning'].append(message)
+
+    def error(self, message: str):
+        self.messages['error'].append(message)
+
+
+@pytest.fixture
+def ie() -> InfoExtractor:
+    ydl = YoutubeDL()
+    return ydl.get_info_extractor('Youtube')
+
+
+@pytest.fixture
+def logger() -> MockLogger:
+    return MockLogger()
+
+
+@pytest.fixture()
+def pot_request() -> PoTokenRequest:
+    return PoTokenRequest(
+        context=PoTokenContext.GVS,
+        innertube_context={'client': {'clientName': 'WEB'}},
+        innertube_host='youtube.com',
+        session_index=None,
+        player_url=None,
+        is_authenticated=False,
+        video_webpage=None,
+
+        visitor_data='example-visitor-data',
+        data_sync_id='example-data-sync-id',
+        video_id='example-video-id',
+
+        request_cookiejar=YoutubeDLCookieJar(),
+        request_proxy=None,
+        request_headers=HTTPHeaderDict(),
+        request_timeout=None,
+        request_source_address=None,
+        request_verify_tls=True,
+
+        bypass_cache=False,
+    )

test/test_pot/test_pot_builtin_memorycache.py

@@ -0,0 +1,117 @@
+import threading
+import time
+from collections import OrderedDict
+import pytest
+from yt_dlp.extractor.youtube.pot._provider import IEContentProvider, BuiltinIEContentProvider
+from yt_dlp.utils import bug_reports_message
+from yt_dlp.extractor.youtube.pot._builtin.memory_cache import MemoryLRUPCP, memorylru_preference, initialize_global_cache
+from yt_dlp.version import __version__
+from yt_dlp.extractor.youtube.pot._registry import _pot_cache_providers, _pot_memory_cache
+
+
+class TestMemoryLRUPCS:
+
+    def test_base_type(self):
+        assert issubclass(MemoryLRUPCP, IEContentProvider)
+        assert issubclass(MemoryLRUPCP, BuiltinIEContentProvider)
+
+    @pytest.fixture
+    def pcp(self, ie, logger) -> MemoryLRUPCP:
+        return MemoryLRUPCP(ie, logger, {}, initialize_cache=lambda max_size: (OrderedDict(), threading.Lock(), max_size))
+
+    def test_is_registered(self):
+        assert _pot_cache_providers.value.get('MemoryLRU') == MemoryLRUPCP
+
+    def test_initialization(self, pcp):
+        assert pcp.PROVIDER_NAME == 'memory'
+        assert pcp.PROVIDER_VERSION == __version__
+        assert pcp.BUG_REPORT_MESSAGE == bug_reports_message(before='')
+        assert pcp.is_available()
+
+    def test_store_and_get(self, pcp):
+        pcp.store('key1', 'value1', int(time.time()) + 60)
+        assert pcp.get('key1') == 'value1'
+        assert len(pcp.cache) == 1
+
+    def test_store_ignore_expired(self, pcp):
+        pcp.store('key1', 'value1', int(time.time()) - 1)
+        assert len(pcp.cache) == 0
+        assert pcp.get('key1') is None
+        assert len(pcp.cache) == 0
+
+    def test_store_override_existing_key(self, ie, logger):
+        MAX_SIZE = 2
+        pcp = MemoryLRUPCP(ie, logger, {}, initialize_cache=lambda max_size: (OrderedDict(), threading.Lock(), MAX_SIZE))
+        pcp.store('key1', 'value1', int(time.time()) + 60)
+        pcp.store('key2', 'value2', int(time.time()) + 60)
+        assert len(pcp.cache) == 2
+        pcp.store('key1', 'value2', int(time.time()) + 60)
+        # Ensure that the override key gets added to the end of the cache instead of in the same position
+        pcp.store('key3', 'value3', int(time.time()) + 60)
+        assert pcp.get('key1') == 'value2'
+
+    def test_store_ignore_expired_existing_key(self, pcp):
+        pcp.store('key1', 'value2', int(time.time()) + 60)
+        pcp.store('key1', 'value1', int(time.time()) - 1)
+        assert len(pcp.cache) == 1
+        assert pcp.get('key1') == 'value2'
+        assert len(pcp.cache) == 1
+
+    def test_get_key_expired(self, pcp):
+        pcp.store('key1', 'value1', int(time.time()) + 60)
+        assert pcp.get('key1') == 'value1'
+        assert len(pcp.cache) == 1
+        pcp.cache['key1'] = ('value1', int(time.time()) - 1)
+        assert pcp.get('key1') is None
+        assert len(pcp.cache) == 0
+
+    def test_lru_eviction(self, ie, logger):
+        MAX_SIZE = 2
+        provider = MemoryLRUPCP(ie, logger, {}, initialize_cache=lambda max_size: (OrderedDict(), threading.Lock(), MAX_SIZE))
+        provider.store('key1', 'value1', int(time.time()) + 5)
+        provider.store('key2', 'value2', int(time.time()) + 5)
+        assert len(provider.cache) == 2
+
+        assert provider.get('key1') == 'value1'
+
+        provider.store('key3', 'value3', int(time.time()) + 5)
+        assert len(provider.cache) == 2
+
+        assert provider.get('key2') is None
+
+        provider.store('key4', 'value4', int(time.time()) + 5)
+        assert len(provider.cache) == 2
+
+        assert provider.get('key1') is None
+        assert provider.get('key3') == 'value3'
+        assert provider.get('key4') == 'value4'
+
+    def test_delete(self, pcp):
+        pcp.store('key1', 'value1', int(time.time()) + 5)
+        assert len(pcp.cache) == 1
+        assert pcp.get('key1') == 'value1'
+        pcp.delete('key1')
+        assert len(pcp.cache) == 0
+        assert pcp.get('key1') is None
+
+    def test_use_global_cache_default(self, ie, logger):
+        pcp = MemoryLRUPCP(ie, logger, {})
+        assert pcp.max_size == _pot_memory_cache.value['max_size'] == 25
+        assert pcp.cache is _pot_memory_cache.value['cache']
+        assert pcp.lock is _pot_memory_cache.value['lock']
+
+        pcp2 = MemoryLRUPCP(ie, logger, {})
+        assert pcp.max_size == pcp2.max_size == _pot_memory_cache.value['max_size'] == 25
+        assert pcp.cache is pcp2.cache is _pot_memory_cache.value['cache']
+        assert pcp.lock is pcp2.lock is _pot_memory_cache.value['lock']
+
+    def test_fail_max_size_change_global(self, ie, logger):
+        pcp = MemoryLRUPCP(ie, logger, {})
+        assert pcp.max_size == _pot_memory_cache.value['max_size'] == 25
+        with pytest.raises(ValueError, match='Cannot change max_size of initialized global memory cache'):
+            initialize_global_cache(50)
+
+        assert pcp.max_size == _pot_memory_cache.value['max_size'] == 25
+
+    def test_memory_lru_preference(self, pcp, ie, pot_request):
+        assert memorylru_preference(pcp, pot_request) == 10000

test/test_pot/test_pot_builtin_utils.py

@@ -0,0 +1,46 @@
+import pytest
+from yt_dlp.extractor.youtube.pot.provider import (
+    PoTokenContext,
+
+)
+
+from yt_dlp.extractor.youtube.pot.utils import get_webpo_content_binding, ContentBindingType
+
+
+class TestGetWebPoContentBinding:
+
+    @pytest.mark.parametrize('client_name, context, is_authenticated, expected', [
+        *[(client, context, is_authenticated, expected) for client in [
+            'WEB', 'MWEB', 'TVHTML5', 'WEB_EMBEDDED_PLAYER', 'WEB_CREATOR', 'TVHTML5_SIMPLY_EMBEDDED_PLAYER']
+          for context, is_authenticated, expected in [
+            (PoTokenContext.GVS, False, ('example-visitor-data', ContentBindingType.VISITOR_DATA)),
+            (PoTokenContext.PLAYER, False, ('example-video-id', ContentBindingType.VIDEO_ID)),
+            (PoTokenContext.GVS, True, ('example-data-sync-id', ContentBindingType.DATASYNC_ID)),
+        ]],
+        ('WEB_REMIX', PoTokenContext.GVS, False, ('example-visitor-data', ContentBindingType.VISITOR_DATA)),
+        ('WEB_REMIX', PoTokenContext.PLAYER, False, ('example-visitor-data', ContentBindingType.VISITOR_DATA)),
+        ('ANDROID', PoTokenContext.GVS, False, (None, None)),
+        ('IOS', PoTokenContext.GVS, False, (None, None)),
+    ])
+    def test_get_webpo_content_binding(self, pot_request, client_name, context, is_authenticated, expected):
+        pot_request.innertube_context['client']['clientName'] = client_name
+        pot_request.context = context
+        pot_request.is_authenticated = is_authenticated
+        assert get_webpo_content_binding(pot_request) == expected
+
+    def test_extract_visitor_id(self, pot_request):
+        pot_request.visitor_data = 'CgsxMjNhYmNYWVpfLSiA4s%2DqBg%3D%3D'
+        assert get_webpo_content_binding(pot_request, bind_to_visitor_id=True) == ('123abcXYZ_-', ContentBindingType.VISITOR_ID)
+
+    def test_invalid_visitor_id(self, pot_request):
+        # visitor id not alphanumeric (i.e. protobuf extraction failed)
+        pot_request.visitor_data = 'CggxMjM0NTY3OCiA4s-qBg%3D%3D'
+        assert get_webpo_content_binding(pot_request, bind_to_visitor_id=True) == (pot_request.visitor_data, ContentBindingType.VISITOR_DATA)
+
+    def test_no_visitor_id(self, pot_request):
+        pot_request.visitor_data = 'KIDiz6oG'
+        assert get_webpo_content_binding(pot_request, bind_to_visitor_id=True) == (pot_request.visitor_data, ContentBindingType.VISITOR_DATA)
+
+    def test_invalid_base64(self, pot_request):
+        pot_request.visitor_data = 'invalid-base64'
+        assert get_webpo_content_binding(pot_request, bind_to_visitor_id=True) == (pot_request.visitor_data, ContentBindingType.VISITOR_DATA)