From 2aee875642d5eadf10bc6ceeda19928b138daa20 Mon Sep 17 00:00:00 2001 From: Chocobo1 Date: Sat, 19 Apr 2025 07:11:50 +0800 Subject: [PATCH] Enforce SOCKS proxy setting in search engine plugins Previously it require each plugin to import helpers.py to setup SOCKS proxy. Now it is enforced by default for all plugins. Also added a function for plugins to ignore/restore the socket to default state. PR #22554. --- src/searchengine/nova3/helpers.py | 48 +++++++++++++++---------------- src/searchengine/nova3/nova2.py | 17 +++++++---- src/searchengine/nova3/nova2dl.py | 9 ++++-- 3 files changed, 41 insertions(+), 33 deletions(-) diff --git a/src/searchengine/nova3/helpers.py b/src/searchengine/nova3/helpers.py index 85eec1541..5e4457235 100644 --- a/src/searchengine/nova3/helpers.py +++ b/src/searchengine/nova3/helpers.py @@ -1,4 +1,4 @@ -#VERSION: 1.52 +# VERSION: 1.53 # Author: # Christophe DUMEZ (chris@qbittorrent.org) @@ -61,32 +61,32 @@ def _getBrowserUserAgent() -> str: _headers: dict[str, Any] = {'User-Agent': _getBrowserUserAgent()} +_original_socket = socket.socket -def injectSOCKSProxySocket() -> None: - socksURL = os.environ.get("qbt_socks_proxy") - if socksURL is not None: - parts = urllib.parse.urlsplit(socksURL) - resolveHostname = (parts.scheme == "socks4a") or (parts.scheme == "socks5h") - if (parts.scheme == "socks4") or (parts.scheme == "socks4a"): - socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS4, parts.hostname, parts.port, resolveHostname) - socket.socket = socks.socksocket # type: ignore[misc] - elif (parts.scheme == "socks5") or (parts.scheme == "socks5h"): - socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5, parts.hostname, parts.port, resolveHostname, parts.username, parts.password) - socket.socket = socks.socksocket # type: ignore[misc] +def enable_socks_proxy(enable: bool) -> None: + if enable: + socksURL = os.environ.get("qbt_socks_proxy") + if socksURL is not None: + parts = urllib.parse.urlsplit(socksURL) + resolveHostname = (parts.scheme == "socks4a") or (parts.scheme == "socks5h") + if (parts.scheme == "socks4") or (parts.scheme == "socks4a"): + socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS4, parts.hostname, parts.port, resolveHostname) + socket.socket = socks.socksocket # type: ignore[misc] + elif (parts.scheme == "socks5") or (parts.scheme == "socks5h"): + socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5, parts.hostname, parts.port, resolveHostname, parts.username, parts.password) + socket.socket = socks.socksocket # type: ignore[misc] + else: + # the following code provide backward compatibility for older qbt versions + # TODO: scheduled be removed with qbt >= 5.3 + legacySocksURL = os.environ.get("sock_proxy") + if legacySocksURL is not None: + legacySocksURL = f"socks5h://{legacySocksURL.strip()}" + parts = urllib.parse.urlsplit(legacySocksURL) + socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5, parts.hostname, parts.port, True, parts.username, parts.password) + socket.socket = socks.socksocket # type: ignore[misc] else: - # the following code provide backward compatibility for older qbt versions - # TODO: scheduled be removed with qbt >= 5.3 - legacySocksURL = os.environ.get("sock_proxy") - if legacySocksURL is not None: - legacySocksURL = f"socks5h://{legacySocksURL.strip()}" - parts = urllib.parse.urlsplit(legacySocksURL) - socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5, parts.hostname, parts.port, True, parts.username, parts.password) - socket.socket = socks.socksocket # type: ignore[misc] - - -# monkey patching, run it now -injectSOCKSProxySocket() + socket.socket = _original_socket # type: ignore[misc] # This is only provided for backward compatibility, new code should not use it diff --git a/src/searchengine/nova3/nova2.py b/src/searchengine/nova3/nova2.py index fef072113..6df365a07 100644 --- a/src/searchengine/nova3/nova2.py +++ b/src/searchengine/nova3/nova2.py @@ -1,4 +1,4 @@ -#VERSION: 1.48 +# VERSION: 1.49 # Author: # Fabien Devaux @@ -46,6 +46,16 @@ from multiprocessing import Pool, cpu_count from os import path from typing import Optional +# qbt tend to run this script in 'isolate mode' so append the current path manually +current_path = str(pathlib.Path(__file__).parent.resolve()) +if current_path not in sys.path: + sys.path.append(current_path) + +import helpers + +# enable SOCKS proxy for all plugins by default +helpers.enable_socks_proxy(True) + THREADED: bool = True try: MAX_THREADS: int = cpu_count() @@ -182,11 +192,6 @@ def run_search(search_params: tuple[type[Engine], str, Category]) -> bool: if __name__ == "__main__": def main() -> int: - # qbt tend to run this script in 'isolate mode' so append the current path manually - current_path = str(pathlib.Path(__file__).parent.resolve()) - if current_path not in sys.path: - sys.path.append(current_path) - # https://docs.python.org/3/library/sys.html#sys.exit class ExitCode(Enum): OK = 0 diff --git a/src/searchengine/nova3/nova2dl.py b/src/searchengine/nova3/nova2dl.py index c35646654..8a45b3724 100644 --- a/src/searchengine/nova3/nova2dl.py +++ b/src/searchengine/nova3/nova2dl.py @@ -1,4 +1,4 @@ -#VERSION: 1.24 +# VERSION: 1.25 # Author: # Christophe DUMEZ (chris@qbittorrent.org) @@ -38,7 +38,10 @@ current_path = str(pathlib.Path(__file__).parent.resolve()) if current_path not in sys.path: sys.path.append(current_path) -from helpers import download_file +import helpers + +# enable SOCKS proxy for all plugins by default +helpers.enable_socks_proxy(True) if __name__ == '__main__': if len(sys.argv) < 3: @@ -58,6 +61,6 @@ if __name__ == '__main__': if hasattr(engine, 'download_torrent'): engine.download_torrent(download_param) else: - print(download_file(download_param)) + print(helpers.download_file(download_param)) sys.exit(0)