From b5683a3e774b460c756d82c73541442febd386b3 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Mon, 27 Oct 2025 19:54:09 +0300 Subject: [PATCH] Migrate inline cookie creation to standalone function --- app/library/Download.py | 20 +- app/library/DownloadQueue.py | 6 +- app/library/ItemDTO.py | 17 +- app/library/Presets.py | 33 +++- app/library/Utils.py | 25 +++ app/routes/api/notifications.py | 1 - app/tests/test_utils.py | 327 ++++++++++++++++++++++++++++++++ 7 files changed, 409 insertions(+), 20 deletions(-) diff --git a/app/library/Download.py b/app/library/Download.py index 22da3829..a69a1a97 100644 --- a/app/library/Download.py +++ b/app/library/Download.py @@ -17,7 +17,7 @@ from .config import Config from .Events import EventBus, Events from .ffprobe import ffprobe from .ItemDTO import ItemDTO -from .Utils import delete_dir, extract_info, extract_ytdlp_logs, load_cookies +from .Utils import create_cookies_file, delete_dir, extract_info, extract_ytdlp_logs from .ytdlp import YTDLP @@ -138,10 +138,11 @@ class Download: "status": data.get("status"), "filename": data.get("filename"), "info_dict": { - k: v for k, v in data.get("info_dict", {}).items() + k: v + for k, v in data.get("info_dict", {}).items() if k not in ["formats", "thumbnails", "description", "tags", "_format_sort_fields"] and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable - } + }, } self.logger.debug(f"PG Hook: {d_safe}") except Exception as e: @@ -162,10 +163,11 @@ class Download: "postprocessor": data.get("postprocessor"), "status": data.get("status"), "info_dict": { - k: v for k, v in data.get("info_dict", {}).items() + k: v + for k, v in data.get("info_dict", {}).items() if k not in ["formats", "thumbnails", "description", "tags", "_format_sort_fields"] and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable - } + }, } self.logger.debug(f"PP Hook: {d_safe}") except Exception as e: @@ -238,10 +240,7 @@ class Download: self.logger.debug( f"Creating cookie file for '{self.info.id}: {self.info.title}' - '{cookie_file}'." ) - cookie_file.write_text(self.info.cookies) - params["cookiefile"] = str(cookie_file.as_posix()) - - load_cookies(cookie_file) + params["cookiefile"] = str(create_cookies_file(self.info.cookies, cookie_file).as_posix()) except Exception as e: err_msg: str = f"Failed to create cookie file for '{self.info.id}: {self.info.title}'. '{e!s}'." self.logger.error(err_msg) @@ -536,8 +535,7 @@ class Download: # If still alive, use hard kill if self.proc.is_alive(): self.logger.warning( - f"Process PID={self.proc.pid} did not respond to terminate(), " - f"killing forcefully." + f"Process PID={self.proc.pid} did not respond to terminate(), killing forcefully." ) self.proc.kill() # Give it final moment diff --git a/app/library/DownloadQueue.py b/app/library/DownloadQueue.py index 274a247e..0cd7c1fb 100644 --- a/app/library/DownloadQueue.py +++ b/app/library/DownloadQueue.py @@ -29,10 +29,10 @@ from .Utils import ( archive_add, arg_converter, calc_download_path, + create_cookies_file, dt_delta, extract_info, extract_ytdlp_logs, - load_cookies, merge_dict, str_to_dt, ytdlp_reject, @@ -731,9 +731,7 @@ class DownloadQueue(metaclass=Singleton): if item.cookies: try: - cookie_file.write_text(item.cookies) - yt_conf["cookiefile"] = str(cookie_file.as_posix()) - load_cookies(cookie_file) + yt_conf["cookiefile"] = str(create_cookies_file(item.cookies, cookie_file).as_posix()) except Exception as e: msg = f"Failed to create cookie file for '{item.url}'. '{e!s}'." LOG.error(msg) diff --git a/app/library/ItemDTO.py b/app/library/ItemDTO.py index d638d45c..0ff99fa4 100644 --- a/app/library/ItemDTO.py +++ b/app/library/ItemDTO.py @@ -5,7 +5,7 @@ import uuid from dataclasses import dataclass, field from email.utils import formatdate from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from app.library.encoder import Encoder from app.library.Utils import ( @@ -19,6 +19,9 @@ from app.library.Utils import ( ) from app.library.YTDLPOpts import YTDLPOpts +if TYPE_CHECKING: + from app.library.Presets import Preset + LOG: logging.Logger = logging.getLogger("ItemDTO") @@ -207,6 +210,18 @@ class Item: return Item(**data) + def get_preset(self) -> "Preset": + """ + Get the preset for the item. + + Returns: + Preset: The preset for the item. If not found, None. + + """ + from .Presets import Presets + + return Presets.get_instance().get(self.preset if self.preset else self._default_preset()) + def get_archive_id(self) -> str | None: """ Get the archive ID for the download URL. diff --git a/app/library/Presets.py b/app/library/Presets.py index 22477f74..e53f5562 100644 --- a/app/library/Presets.py +++ b/app/library/Presets.py @@ -10,7 +10,7 @@ from aiohttp import web from .config import Config from .Events import EventBus, Events from .Singleton import Singleton -from .Utils import arg_converter, init_class +from .Utils import arg_converter, create_cookies_file, init_class LOG = logging.getLogger("presets") @@ -103,8 +103,11 @@ class Preset: default: bool = False """If True, the preset is a default preset.""" + _cookies_file: Path | None = field(init=False, default=None) + """The path to the cookies file.""" + def serialize(self) -> dict: - return self.__dict__ + return {k: v for k, v in self.__dict__.items() if not k.startswith("_")} def json(self) -> str: from .encoder import Encoder @@ -112,7 +115,31 @@ class Preset: return Encoder().encode(self.serialize()) def get(self, key: str, default: Any = None) -> Any: - return self.serialize().get(key, default) + return self.__dict__.get(key, default) + + def get_cookies_file(self, config: Config | None = None) -> Path | None: + """ + Get the path to the cookies file. + + Args: + config (Config|None): The config instance. + + Returns: + Path|None: The path to the cookies file. + + """ + if self._cookies_file: + return self._cookies_file + + if not self.cookies or not self.id: + return None + + if not config: + config = Config.get_instance() + + self._cookies_file = create_cookies_file(self.cookies, Path(config.config_path) / "cookies" / f"{self.id}.txt") + + return self._cookies_file class Presets(metaclass=Singleton): diff --git a/app/library/Utils.py b/app/library/Utils.py index d23f5169..a09095db 100644 --- a/app/library/Utils.py +++ b/app/library/Utils.py @@ -1634,3 +1634,28 @@ def get_channel_images(thumbnails: list[dict]) -> dict: artwork["poster"] = artwork["thumb"] # optional fallback return artwork + + +def create_cookies_file(cookies: str, file: Path | None = None) -> Path: + """ + Create a cookies file from a string of cookies. + + Args: + cookies (str): The cookie string. + file (Path|None): The path to the cookie file. If None, a temporary file is created. + + Returns: + Path: The path to the created cookie file. + + """ + if file is None: + from .config import Config + + file = Path(Config.get_instance().temp_path, f"c_{uuid.uuid4().hex}.txt") + + file.parent.mkdir(parents=True, exist_ok=True) + file.write_text(cookies) + + load_cookies(file) + + return file diff --git a/app/routes/api/notifications.py b/app/routes/api/notifications.py index b1604627..6174b62b 100644 --- a/app/routes/api/notifications.py +++ b/app/routes/api/notifications.py @@ -43,7 +43,6 @@ async def notification_add(request: Request, encoder: Encoder) -> Response: Args: request (Request): The request object. encoder (Encoder): The encoder instance. - notify (EventBus): The event bus instance. Returns: Response: The response object. diff --git a/app/tests/test_utils.py b/app/tests/test_utils.py index aeeda927..2897777a 100644 --- a/app/tests/test_utils.py +++ b/app/tests/test_utils.py @@ -1701,3 +1701,330 @@ class TestLoadModules: except Exception: # Module loading might fail in test environment assert True + + +class TestGetChannelImages: + """Test the get_channel_images function.""" + + def test_get_channel_images_poster_from_portrait(self): + """Test extracting poster image from portrait ratio thumbnail.""" + from app.library.Utils import get_channel_images + + thumbnails = [ + {"url": "http://example.com/poster.jpg", "width": 200, "height": 300, "id": "some_id"} + ] + + result = get_channel_images(thumbnails) + + assert "poster" in result + assert result["poster"] == "http://example.com/poster.jpg" + + def test_get_channel_images_thumb_from_square(self): + """Test extracting thumbnail from square ratio.""" + from app.library.Utils import get_channel_images + + thumbnails = [{"url": "http://example.com/thumb.jpg", "width": 300, "height": 300, "id": "id"}] + + result = get_channel_images(thumbnails) + + assert "thumb" in result + assert result["thumb"] == "http://example.com/thumb.jpg" + + def test_get_channel_images_banner_from_wide(self): + """Test extracting banner from very wide image.""" + from app.library.Utils import get_channel_images + + thumbnails = [{"url": "http://example.com/banner.jpg", "width": 1920, "height": 200, "id": "id"}] + + result = get_channel_images(thumbnails) + + assert "banner" in result + assert result["banner"] == "http://example.com/banner.jpg" + + def test_get_channel_images_icon_from_avatar(self): + """Test extracting icon from avatar uncropped.""" + from app.library.Utils import get_channel_images + + thumbnails = [{"url": "http://example.com/icon.jpg", "id": "avatar_uncropped"}] + + result = get_channel_images(thumbnails) + + assert "icon" in result + assert result["icon"] == "http://example.com/icon.jpg" + + def test_get_channel_images_landscape_from_banner_uncropped(self): + """Test extracting landscape from banner uncropped.""" + from app.library.Utils import get_channel_images + + thumbnails = [{"url": "http://example.com/landscape.jpg", "id": "banner_uncropped"}] + + result = get_channel_images(thumbnails) + + assert "landscape" in result + assert result["landscape"] == "http://example.com/landscape.jpg" + + def test_get_channel_images_empty_list(self): + """Test with empty thumbnail list.""" + from app.library.Utils import get_channel_images + + result = get_channel_images([]) + + assert result == {} + + def test_get_channel_images_fallbacks(self): + """Test fallback logic for missing images.""" + from app.library.Utils import get_channel_images + + # Create image with fanart but no banner + thumbnails = [{"url": "http://example.com/fanart.jpg", "width": 1920, "height": 200, "id": "id"}] + + result = get_channel_images(thumbnails) + + assert "fanart" in result + assert "banner" in result # Should fallback to fanart + assert result["banner"] == result["fanart"] + + def test_get_channel_images_no_url(self): + """Test handling of thumbnail without URL.""" + from app.library.Utils import get_channel_images + + thumbnails = [ + {"width": 300, "height": 300, "id": "id"}, # Missing URL + {"url": "http://example.com/thumb.jpg", "width": 300, "height": 300, "id": "id"}, + ] + + result = get_channel_images(thumbnails) + + assert len(result) > 0 + assert result.get("thumb") == "http://example.com/thumb.jpg" + + +class TestIsSafeKey: + """Test the _is_safe_key function (indirectly through merge_dict).""" + + def test_safe_key_normal_string(self): + """Test that normal strings are considered safe keys.""" + from app.library.Utils import merge_dict + + source = {"normal_key": "value"} + dest = {} + + result = merge_dict(source, dest) + + assert "normal_key" in result + assert result["normal_key"] == "value" + + def test_unsafe_key_dunder_attributes(self): + """Test that dunder attributes are blocked.""" + from app.library.Utils import merge_dict + + source = {"__class__": "should_not_merge", "normal_key": "value"} + dest = {} + + result = merge_dict(source, dest) + + assert "__class__" not in result + assert "normal_key" in result + + def test_unsafe_key_empty_string(self): + """Test that empty keys are blocked.""" + from app.library.Utils import merge_dict + + source = {"": "empty_key_value", "normal_key": "value"} + dest = {} + + result = merge_dict(source, dest) + + assert "" not in result + assert "normal_key" in result + + def test_unsafe_key_whitespace_only(self): + """Test that whitespace-only keys are blocked.""" + from app.library.Utils import merge_dict + + source = {" ": "whitespace_key", "normal_key": "value"} + dest = {} + + result = merge_dict(source, dest) + + assert " " not in result + assert "normal_key" in result + + def test_unsafe_key_non_string(self): + """Test that non-string keys are blocked.""" + from app.library.Utils import merge_dict + + source = {123: "numeric_key", "normal_key": "value"} + dest = {} + + result = merge_dict(source, dest) + + assert 123 not in result + assert "normal_key" in result + + +class TestDecryptDataCornerCases: + """Additional tests for decrypt_data edge cases.""" + + def test_decrypt_data_invalid_base64(self): + """Test decrypting invalid base64 data.""" + from app.library.Utils import decrypt_data + + result = decrypt_data("not_valid_base64!!!", b"k" * 32) + + assert result is None + + def test_decrypt_data_wrong_key(self): + """Test decrypting with wrong key returns None.""" + from app.library.Utils import decrypt_data, encrypt_data + + correct_key = b"a" * 32 + wrong_key = b"z" * 32 + + encrypted = encrypt_data("test data", correct_key) + result = decrypt_data(encrypted, wrong_key) + + # Should return None on decryption failure + assert result is None + + def test_decrypt_data_truncated(self): + """Test decrypting truncated encrypted data.""" + from app.library.Utils import decrypt_data + + result = decrypt_data("YQ==", b"k" * 32) + + assert result is None + + +class TestArgConverterAdvanced: + """Advanced tests for arg_converter function.""" + + def test_arg_converter_with_removed_options(self): + """Test arg_converter with removed options tracking.""" + try: + removed = [] + result = arg_converter("--quiet --skip-download", level=True, removed_options=removed) + + assert isinstance(result, dict) + except (ModuleNotFoundError, AttributeError, ImportError): + assert True + + def test_arg_converter_dumps_enabled(self): + """Test arg_converter with dumps flag enabled.""" + try: + result = arg_converter("--format best", dumps=True) + + # With dumps=True, result should still be a dict with JSON-serializable content + assert isinstance(result, (dict, list)) + except (ModuleNotFoundError, AttributeError, ImportError): + assert True + + +class TestCreateCookiesFile: + """Test the create_cookies_file function.""" + + def setup_method(self): + """Set up test environment.""" + self.temp_dir = tempfile.mkdtemp() + self.test_path = Path(self.temp_dir) + + def teardown_method(self): + """Clean up after tests.""" + import shutil + + shutil.rmtree(self.temp_dir, ignore_errors=True) + + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_with_path(self, mock_load_cookies): + """Test creating a cookies file with a specific path.""" + from app.library.Utils import create_cookies_file + + mock_load_cookies.return_value = (True, MagicMock()) + cookie_path = self.test_path / "cookies" / "test_cookies.txt" + + result = create_cookies_file("session=abc123", file=cookie_path) + + assert result == cookie_path + assert cookie_path.exists() + assert cookie_path.is_file() + assert cookie_path.read_text() == "session=abc123" + mock_load_cookies.assert_called_once_with(cookie_path) + + @patch("app.library.config.Config") + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_without_path(self, mock_load_cookies, mock_config): + """Test creating a cookies file without a specific path (auto temp file).""" + from app.library.Utils import create_cookies_file + + mock_config_inst = MagicMock() + mock_config_inst.temp_path = self.temp_dir + mock_config.get_instance.return_value = mock_config_inst + + mock_load_cookies.return_value = (True, MagicMock()) + + result = create_cookies_file("session=def456") + + assert result.exists() + assert result.is_file() + assert result.read_text() == "session=def456" + assert result.parent == Path(self.temp_dir) + mock_load_cookies.assert_called_once() + + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_invalid_cookies(self, mock_load_cookies): + """Test creating a cookies file with invalid cookies raises error.""" + from app.library.Utils import create_cookies_file + + mock_load_cookies.side_effect = ValueError("Invalid cookies") + cookie_path = self.test_path / "bad_cookies.txt" + + with pytest.raises(ValueError, match="Invalid cookies"): + create_cookies_file("invalid_data", file=cookie_path) + + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_creates_parent_directory(self, mock_load_cookies): + """Test that create_cookies_file creates parent directories as needed.""" + from app.library.Utils import create_cookies_file + + mock_load_cookies.return_value = (True, MagicMock()) + # Use a deeply nested path that doesn't exist yet + cookie_path = self.test_path / "a" / "b" / "c" / "cookies.txt" + + result = create_cookies_file("test_data", file=cookie_path) + + assert result == cookie_path + assert cookie_path.exists() + assert cookie_path.parent == Path(self.test_path / "a" / "b" / "c") + + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_with_special_characters(self, mock_load_cookies): + """Test create_cookies_file with special characters in cookie data.""" + from app.library.Utils import create_cookies_file + + mock_load_cookies.return_value = (True, MagicMock()) + cookie_data = "session=abc123; path=/; domain=.example.com; secure; httponly" + cookie_path = self.test_path / "special_cookies.txt" + + result = create_cookies_file(cookie_data, file=cookie_path) + + assert result == cookie_path + assert cookie_path.read_text() == cookie_data + + @patch("app.library.Utils.load_cookies") + def test_create_cookies_file_overwrites_existing(self, mock_load_cookies): + """Test that create_cookies_file overwrites existing files.""" + from app.library.Utils import create_cookies_file + + mock_load_cookies.return_value = (True, MagicMock()) + cookie_path = self.test_path / "cookies.txt" + + # Create initial file + cookie_path.parent.mkdir(parents=True, exist_ok=True) + cookie_path.write_text("old_data") + + # Overwrite with new data + result = create_cookies_file("new_data", file=cookie_path) + + assert result == cookie_path + assert cookie_path.read_text() == "new_data" +