diff --git a/test/test_pot/test_pot_director.py b/test/test_pot/test_pot_director.py index 53cc1fd5bc..a9ddf9ca38 100644 --- a/test/test_pot/test_pot_director.py +++ b/test/test_pot/test_pot_director.py @@ -66,7 +66,6 @@ class ExamplePTP(BaseMockPoTokenProvider): _SUPPORTED_CLIENTS = ('WEB',) _SUPPORTED_CONTEXTS = (PoTokenContext.GVS, ) - _SUPPORTED_PROXY_SCHEMES = ('socks5', 'http') def _real_request_pot(self, request: PoTokenRequest) -> PoTokenResponse: if request.data_sync_id == 'example': @@ -82,7 +81,6 @@ def success_ptp(response: PoTokenResponse | None = None, key: str | None = None) _SUPPORTED_CLIENTS = ('WEB',) _SUPPORTED_CONTEXTS = (PoTokenContext.GVS,) - _SUPPORTED_PROXY_SCHEMES = ('socks5', 'http') def _real_request_pot(self, request: PoTokenRequest) -> PoTokenResponse: return response or PoTokenResponse(EXAMPLE_PO_TOKEN) diff --git a/test/test_pot/test_pot_framework.py b/test/test_pot/test_pot_framework.py index fabfc92c09..e1ea1ee318 100644 --- a/test/test_pot/test_pot_framework.py +++ b/test/test_pot/test_pot_framework.py @@ -6,6 +6,7 @@ from yt_dlp.utils.networking import HTTPHeaderDict from yt_dlp.extractor.youtube.pot.provider import ( PoTokenRequest, PoTokenContext, + ExternalRequestFeature, ) @@ -38,7 +39,11 @@ class ExamplePTP(PoTokenProvider): _SUPPORTED_CLIENTS = ('WEB',) _SUPPORTED_CONTEXTS = (PoTokenContext.GVS, ) - _SUPPORTED_PROXY_SCHEMES = ('socks5', 'http') + + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = ( + ExternalRequestFeature.PROXY_SCHEME_HTTP, + ExternalRequestFeature.PROXY_SCHEME_SOCKS5H, + ) def is_available(self) -> bool: return True @@ -146,9 +151,88 @@ class TestPoTokenProvider: provider = ExamplePTP(ie=ie, logger=logger, settings={}) pot_request.request_proxy = 'socks4://example.com' - with pytest.raises(PoTokenProviderRejectedRequest): + with pytest.raises( + PoTokenProviderRejectedRequest, + match='External requests by "example" provider do not support proxy scheme "socks4". Supported proxy ' + 'schemes: http, socks5h', + ): provider.request_pot(pot_request) + pot_request.request_proxy = 'http://example.com' + + assert provider.request_pot(pot_request) + + def test_provider_ignore_external_request_features(self, ie, logger, pot_request): + class InternalPTP(ExamplePTP): + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = None + + provider = InternalPTP(ie=ie, logger=logger, settings={}) + + pot_request.request_proxy = 'socks5://example.com' + assert provider.request_pot(pot_request) + pot_request.request_source_address = '0.0.0.0' + assert provider.request_pot(pot_request) + + def test_provider_unsupported_external_request_source_address(self, ie, logger, pot_request): + class InternalPTP(ExamplePTP): + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = tuple() + + provider = InternalPTP(ie=ie, logger=logger, settings={}) + + pot_request.request_source_address = None + assert provider.request_pot(pot_request) + + pot_request.request_source_address = '0.0.0.0' + with pytest.raises( + PoTokenProviderRejectedRequest, + match='External requests by "example" provider do not support setting source address', + ): + provider.request_pot(pot_request) + + def test_provider_supported_external_request_source_address(self, ie, logger, pot_request): + class InternalPTP(ExamplePTP): + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = ( + ExternalRequestFeature.SOURCE_ADDRESS, + ) + + provider = InternalPTP(ie=ie, logger=logger, settings={}) + + pot_request.request_source_address = None + assert provider.request_pot(pot_request) + + pot_request.request_source_address = '0.0.0.0' + assert provider.request_pot(pot_request) + + def test_provider_unsupported_external_request_tls_verification(self, ie, logger, pot_request): + class InternalPTP(ExamplePTP): + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = tuple() + + provider = InternalPTP(ie=ie, logger=logger, settings={}) + + pot_request.request_verify_tls = True + assert provider.request_pot(pot_request) + + pot_request.request_verify_tls = False + with pytest.raises( + PoTokenProviderRejectedRequest, + match='External requests by "example" provider do not support ignoring TLS certificate failures', + ): + provider.request_pot(pot_request) + + def test_provider_supported_external_request_tls_verification(self, ie, logger, pot_request): + class InternalPTP(ExamplePTP): + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = ( + ExternalRequestFeature.DISABLE_TLS_VERIFICATION, + ) + + provider = InternalPTP(ie=ie, logger=logger, settings={}) + + pot_request.request_verify_tls = True + assert provider.request_pot(pot_request) + + pot_request.request_verify_tls = False + assert provider.request_pot(pot_request) + def test_provider_request_webpage(self, ie, logger, pot_request): provider = ExamplePTP(ie=ie, logger=logger, settings={}) diff --git a/yt_dlp/extractor/youtube/pot/README.md b/yt_dlp/extractor/youtube/pot/README.md index 19de1eebcb..51096d5e0f 100644 --- a/yt_dlp/extractor/youtube/pot/README.md +++ b/yt_dlp/extractor/youtube/pot/README.md @@ -31,6 +31,7 @@ from yt_dlp.extractor.youtube.pot.provider import ( PoTokenProviderRejectedRequest, register_provider, register_preference, + ExternalRequestFeature, ) from yt_dlp.networking.common import Request from yt_dlp.extractor.youtube.pot.utils import get_webpo_content_binding @@ -58,10 +59,14 @@ class MyPoTokenProviderPTP(PoTokenProvider): # Provider name must end with "PTP PoTokenContext.GVS, ) - # Possible values: http, https, socks4, socks4a, socks5, socks5h - # If your provider makes requests outside PoTokenProvider._urlopen, you should set this to any proxy schemes supported. - # If you use PoTokenProvider._urlopen to make requests, set to None. - _SUPPORTED_PROXY_SCHEMES = ('http',) + # If your provider makes external requests to websites (i.e. to youtube.com) using another library or service (i.e., not _request_webpage), + # set the request features that are supported here. + # If only using _request_webpage to make external requests, set this to None. + _SUPPORTED_EXTERNAL_REQEUST_FEATURES = ( + ExternalRequestFeature.PROXY_SCHEME_HTTP, + ExternalRequestFeature.SOURCE_ADDRESS, + ExternalRequestFeature.DISABLE_TLS_VERIFICATION + ) def is_available(self) -> bool: """ diff --git a/yt_dlp/extractor/youtube/pot/provider.py b/yt_dlp/extractor/youtube/pot/provider.py index 6cddb46a5d..1a96e478a8 100644 --- a/yt_dlp/extractor/youtube/pot/provider.py +++ b/yt_dlp/extractor/youtube/pot/provider.py @@ -6,6 +6,7 @@ import abc import copy import dataclasses import enum +import functools import typing import urllib.parse @@ -22,6 +23,7 @@ from yt_dlp.utils import traverse_obj from yt_dlp.utils.networking import HTTPHeaderDict __all__ = [ + 'ExternalRequestFeature', 'PoTokenContext', 'PoTokenProvider', 'PoTokenProviderError', @@ -90,6 +92,17 @@ class PoTokenProviderError(IEContentProviderError): """An error occurred while fetching a PO Token""" +class ExternalRequestFeature(enum.Enum): + PROXY_SCHEME_HTTP = enum.auto() + PROXY_SCHEME_HTTPS = enum.auto() + PROXY_SCHEME_SOCKS4 = enum.auto() + PROXY_SCHEME_SOCKS4A = enum.auto() + PROXY_SCHEME_SOCKS5 = enum.auto() + PROXY_SCHEME_SOCKS5H = enum.auto() + SOURCE_ADDRESS = enum.auto() + DISABLE_TLS_VERIFICATION = enum.auto() + + class PoTokenProvider(IEContentProvider, abc.ABC, suffix='PTP'): # Set to None to disable the check @@ -101,8 +114,10 @@ class PoTokenProvider(IEContentProvider, abc.ABC, suffix='PTP'): # Also see yt_dlp.extractor.youtube._base.INNERTUBE_CLIENTS for a list of client names currently supported by the YouTube extractor. _SUPPORTED_CLIENTS: tuple[str] | None = () - # Possible values: http, https, socks4, socks4a, socks5, socks5h - _SUPPORTED_PROXY_SCHEMES: tuple[str] | None = () + # If making external requests to websites (i.e. to youtube.com) using another library or service (i.e., not _request_webpage), + # add the request features that are supported. + # If only using _request_webpage to make external requests, set this to None. + _SUPPORTED_EXTERNAL_REQEUST_FEATURES: tuple[ExternalRequestFeature] | None = () def __validate_request(self, request: PoTokenRequest): if not self.is_available(): @@ -118,11 +133,40 @@ class PoTokenProvider(IEContentProvider, abc.ABC, suffix='PTP'): raise PoTokenProviderRejectedRequest( f'Client "{client_name}" is not supported by {self.PROVIDER_NAME}. Supported clients: {", ".join(self._SUPPORTED_CLIENTS) or "none"}') - if self._SUPPORTED_PROXY_SCHEMES is not None and request.request_proxy: + self.__validate_external_request_features(request) + + @functools.cached_property + def _supported_proxy_schemes(self): + return { + scheme: feature + for scheme, feature in { + 'http': ExternalRequestFeature.PROXY_SCHEME_HTTP, + 'https': ExternalRequestFeature.PROXY_SCHEME_HTTPS, + 'socks4': ExternalRequestFeature.PROXY_SCHEME_SOCKS4, + 'socks4a': ExternalRequestFeature.PROXY_SCHEME_SOCKS4A, + 'socks5': ExternalRequestFeature.PROXY_SCHEME_SOCKS5, + 'socks5h': ExternalRequestFeature.PROXY_SCHEME_SOCKS5H, + }.items() + if feature in (self._SUPPORTED_EXTERNAL_REQEUST_FEATURES or []) + } + + def __validate_external_request_features(self, request: PoTokenRequest): + if self._SUPPORTED_EXTERNAL_REQEUST_FEATURES is None: + return + + if request.request_proxy: scheme = urllib.parse.urlparse(request.request_proxy).scheme - if scheme.lower() not in self._SUPPORTED_PROXY_SCHEMES: + if scheme.lower() not in self._supported_proxy_schemes: raise PoTokenProviderRejectedRequest( - f'Proxy scheme "{scheme}" is not supported by {self.PROVIDER_NAME}. Supported proxy schemes: {", ".join(self._SUPPORTED_PROXY_SCHEMES) or "none"}') + f'External requests by "{self.PROVIDER_NAME}" provider do not support proxy scheme "{scheme}". Supported proxy schemes: {", ".join(self._supported_proxy_schemes) or "none"}') + + if request.request_source_address and ExternalRequestFeature.SOURCE_ADDRESS not in self._SUPPORTED_EXTERNAL_REQEUST_FEATURES: + raise PoTokenProviderRejectedRequest( + f'External requests by "{self.PROVIDER_NAME}" provider do not support setting source address') + + if not request.request_verify_tls and ExternalRequestFeature.DISABLE_TLS_VERIFICATION not in self._SUPPORTED_EXTERNAL_REQEUST_FEATURES: + raise PoTokenProviderRejectedRequest( + f'External requests by "{self.PROVIDER_NAME}" provider do not support ignoring TLS certificate failures') def request_pot(self, request: PoTokenRequest) -> PoTokenResponse: self.__validate_request(request)