Migrate inline cookie creation to standalone function

This commit is contained in:
arabcoders 2025-10-27 19:54:09 +03:00
parent 24493885aa
commit b5683a3e77
7 changed files with 409 additions and 20 deletions

View file

@ -17,7 +17,7 @@ from .config import Config
from .Events import EventBus, Events from .Events import EventBus, Events
from .ffprobe import ffprobe from .ffprobe import ffprobe
from .ItemDTO import ItemDTO 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 from .ytdlp import YTDLP
@ -138,10 +138,11 @@ class Download:
"status": data.get("status"), "status": data.get("status"),
"filename": data.get("filename"), "filename": data.get("filename"),
"info_dict": { "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"] 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 and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable
} },
} }
self.logger.debug(f"PG Hook: {d_safe}") self.logger.debug(f"PG Hook: {d_safe}")
except Exception as e: except Exception as e:
@ -162,10 +163,11 @@ class Download:
"postprocessor": data.get("postprocessor"), "postprocessor": data.get("postprocessor"),
"status": data.get("status"), "status": data.get("status"),
"info_dict": { "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"] 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 and not isinstance(v, (type(None).__bases__[0], type(lambda: None))) # Skip non-serializable
} },
} }
self.logger.debug(f"PP Hook: {d_safe}") self.logger.debug(f"PP Hook: {d_safe}")
except Exception as e: except Exception as e:
@ -238,10 +240,7 @@ class Download:
self.logger.debug( self.logger.debug(
f"Creating cookie file for '{self.info.id}: {self.info.title}' - '{cookie_file}'." f"Creating cookie file for '{self.info.id}: {self.info.title}' - '{cookie_file}'."
) )
cookie_file.write_text(self.info.cookies) params["cookiefile"] = str(create_cookies_file(self.info.cookies, cookie_file).as_posix())
params["cookiefile"] = str(cookie_file.as_posix())
load_cookies(cookie_file)
except Exception as e: except Exception as e:
err_msg: str = f"Failed to create cookie file for '{self.info.id}: {self.info.title}'. '{e!s}'." err_msg: str = f"Failed to create cookie file for '{self.info.id}: {self.info.title}'. '{e!s}'."
self.logger.error(err_msg) self.logger.error(err_msg)
@ -536,8 +535,7 @@ class Download:
# If still alive, use hard kill # If still alive, use hard kill
if self.proc.is_alive(): if self.proc.is_alive():
self.logger.warning( self.logger.warning(
f"Process PID={self.proc.pid} did not respond to terminate(), " f"Process PID={self.proc.pid} did not respond to terminate(), killing forcefully."
f"killing forcefully."
) )
self.proc.kill() self.proc.kill()
# Give it final moment # Give it final moment

View file

@ -29,10 +29,10 @@ from .Utils import (
archive_add, archive_add,
arg_converter, arg_converter,
calc_download_path, calc_download_path,
create_cookies_file,
dt_delta, dt_delta,
extract_info, extract_info,
extract_ytdlp_logs, extract_ytdlp_logs,
load_cookies,
merge_dict, merge_dict,
str_to_dt, str_to_dt,
ytdlp_reject, ytdlp_reject,
@ -731,9 +731,7 @@ class DownloadQueue(metaclass=Singleton):
if item.cookies: if item.cookies:
try: try:
cookie_file.write_text(item.cookies) yt_conf["cookiefile"] = str(create_cookies_file(item.cookies, cookie_file).as_posix())
yt_conf["cookiefile"] = str(cookie_file.as_posix())
load_cookies(cookie_file)
except Exception as e: except Exception as e:
msg = f"Failed to create cookie file for '{item.url}'. '{e!s}'." msg = f"Failed to create cookie file for '{item.url}'. '{e!s}'."
LOG.error(msg) LOG.error(msg)

View file

@ -5,7 +5,7 @@ import uuid
from dataclasses import dataclass, field from dataclasses import dataclass, field
from email.utils import formatdate from email.utils import formatdate
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any
from app.library.encoder import Encoder from app.library.encoder import Encoder
from app.library.Utils import ( from app.library.Utils import (
@ -19,6 +19,9 @@ from app.library.Utils import (
) )
from app.library.YTDLPOpts import YTDLPOpts from app.library.YTDLPOpts import YTDLPOpts
if TYPE_CHECKING:
from app.library.Presets import Preset
LOG: logging.Logger = logging.getLogger("ItemDTO") LOG: logging.Logger = logging.getLogger("ItemDTO")
@ -207,6 +210,18 @@ class Item:
return Item(**data) 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: def get_archive_id(self) -> str | None:
""" """
Get the archive ID for the download URL. Get the archive ID for the download URL.

View file

@ -10,7 +10,7 @@ from aiohttp import web
from .config import Config from .config import Config
from .Events import EventBus, Events from .Events import EventBus, Events
from .Singleton import Singleton 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") LOG = logging.getLogger("presets")
@ -103,8 +103,11 @@ class Preset:
default: bool = False default: bool = False
"""If True, the preset is a default preset.""" """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: 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: def json(self) -> str:
from .encoder import Encoder from .encoder import Encoder
@ -112,7 +115,31 @@ class Preset:
return Encoder().encode(self.serialize()) return Encoder().encode(self.serialize())
def get(self, key: str, default: Any = None) -> Any: 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): class Presets(metaclass=Singleton):

View file

@ -1634,3 +1634,28 @@ def get_channel_images(thumbnails: list[dict]) -> dict:
artwork["poster"] = artwork["thumb"] # optional fallback artwork["poster"] = artwork["thumb"] # optional fallback
return artwork 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

View file

@ -43,7 +43,6 @@ async def notification_add(request: Request, encoder: Encoder) -> Response:
Args: Args:
request (Request): The request object. request (Request): The request object.
encoder (Encoder): The encoder instance. encoder (Encoder): The encoder instance.
notify (EventBus): The event bus instance.
Returns: Returns:
Response: The response object. Response: The response object.

View file

@ -1701,3 +1701,330 @@ class TestLoadModules:
except Exception: except Exception:
# Module loading might fail in test environment # Module loading might fail in test environment
assert True 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"