fix: restrict sites in config browser cookies settings (#333)

* use enums for supported sites and forum * fix: bugs from previous changes * fix: domain fix * refactor: remove aenum and use inherentance * refactor: remove import * refactor: dynamically decide which sites to allow * update poetry * ran ruff * refactor: remove sites that appear in supportforums from supported hosts * refactor: add aenum again use extend_enum to add members * fix: apply fix from ruff * ran ruff --------- Co-authored-by: datawhores <[email protected]>
jbsparrow · Nov 29, 2024 · dfed245 · dfed245
1 parent 4efd392
commit dfed245
Show file tree

Hide file tree

Showing 7 changed files with 131 additions and 111 deletions.
diff --git a/cyberdrop_dl/config_definitions/config_settings.py b/cyberdrop_dl/config_definitions/config_settings.py
@@ -3,11 +3,17 @@
 
 from pydantic import BaseModel, ByteSize, Field, NonNegativeInt, field_serializer
 
-from cyberdrop_dl.utils.constants import APP_STORAGE, BROWSERS, DOWNLOAD_STORAGE
+from cyberdrop_dl import __version__ as current_version
+from cyberdrop_dl.utils.constants import APP_STORAGE, BROWSERS, DOWNLOAD_STORAGE, PRERELEASE_TAGS
 from cyberdrop_dl.utils.data_enums_classes.hash import Hashing
 
 from .custom_types import AliasModel, HttpAppriseURLModel, NonEmptyStr
 
+if next((tag for tag in PRERELEASE_TAGS if tag in current_version), False):
+    from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedSitesDebug as SupportedSites
+else:
+    from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedSites
+
 
 class DownloadOptions(BaseModel):
     block_download_sub_folders: bool = False
@@ -94,7 +100,7 @@ class Sorting(BaseModel):
 class BrowserCookies(BaseModel):
     browsers: list[BROWSERS] = [BROWSERS.chrome]
     auto_import: bool = False
-    sites: list[NonEmptyStr] = []
+    sites: list[SupportedSites] = [domain.value for domain in SupportedSites]
 
 
 class DupeCleanupOptions(BaseModel):

diff --git a/cyberdrop_dl/managers/manager.py b/cyberdrop_dl/managers/manager.py
@@ -19,7 +19,7 @@
 from cyberdrop_dl.managers.progress_manager import ProgressManager
 from cyberdrop_dl.managers.realdebrid_manager import RealDebridManager
 from cyberdrop_dl.utils.args import ParsedArgs
-from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedDomains
+from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedForums
 from cyberdrop_dl.utils.logger import log
 from cyberdrop_dl.utils.transfer.first_time_setup import TransitionManager
 
@@ -172,10 +172,11 @@ def args_logging(self) -> None:
         auth_data_forums = self.config_manager.authentication_data.forums.model_dump()
         auth_data_others = self.config_manager.authentication_data.model_dump(exclude="forums")
 
-        for forum in SupportedDomains.supported_forums_map.values():
-            forum_xf_cookies_provided[forum] = bool(auth_data_forums[f"{forum}_xf_user_cookie"])
-            forum_credentials_provided[forum] = bool(
-                auth_data_forums[f"{forum}_username"] and auth_data_forums[f"{forum}_password"],
+        for forum in SupportedForums:
+            forum_name = forum.name
+            forum_xf_cookies_provided[forum_name] = bool(auth_data_forums[f"{forum_name}_xf_user_cookie"])
+            forum_credentials_provided[forum_name] = bool(
+                auth_data_forums[f"{forum_name}_username"] and auth_data_forums[f"{forum_name}_password"],
             )
 
         auth_provided = {

diff --git a/cyberdrop_dl/ui/prompts/user_prompts.py b/cyberdrop_dl/ui/prompts/user_prompts.py
@@ -11,7 +11,7 @@
 from cyberdrop_dl.ui.prompts.defaults import ALL_CHOICE, DONE_CHOICE, EXIT_CHOICE
 from cyberdrop_dl.utils.constants import BROWSERS, RESERVED_CONFIG_NAMES
 from cyberdrop_dl.utils.cookie_extraction import get_cookies_from_browsers
-from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedDomains
+from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedForums, SupportedHosts
 from cyberdrop_dl.utils.utilities import clear_term
 
 if TYPE_CHECKING:
@@ -132,8 +132,13 @@ def extract_cookies(manager: Manager, *, dry_run: bool = False) -> None:
     if domain_type == DONE_CHOICE.value:
         return
 
-    all_domains = SupportedDomains.supported_forums_map.keys() if domain_type == 1 else SupportedDomains.supported_hosts
+    all_domains = (
+        [domain.value for domain in SupportedForums]
+        if domain_type == 1
+        else [domain.value for domain in SupportedHosts]
+    )
     domain_choices = [Choice(site) for site in all_domains] + [ALL_CHOICE]
+
     domains = basic_prompts.ask_checkbox(domain_choices, message="Select site(s) to import cookies from:")
     browsers = browser_prompt()
 

diff --git a/cyberdrop_dl/utils/cookie_extraction.py b/cyberdrop_dl/utils/cookie_extraction.py
@@ -9,7 +9,7 @@
 from rich.console import Console
 
 from cyberdrop_dl.dependencies import browser_cookie3
-from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedDomains
+from cyberdrop_dl.utils.data_enums_classes.supported_domains import SupportedForums
 
 if TYPE_CHECKING:
     from http.cookiejar import CookieJar
@@ -63,8 +63,8 @@ def get_cookies_from_browsers(
 
     browsers = browsers or manager.config_manager.settings_data.browser_cookies.browsers
     domains = domains or manager.config_manager.settings_data.browser_cookies.sites
-    browsers = list(map(str.lower, browsers))
-    domains = list(map(str.lower, domains))
+    browsers = [str.lower(x) for x in browsers]
+    domains = [x.value for x in domains]
 
     extractors = [getattr(browser_cookie3, b) for b in browsers if hasattr(browser_cookie3, b)]
 
@@ -86,9 +86,11 @@ def get_cookies_from_browsers(
 
 def update_forum_config_cookies(manager: Manager, forum: str, cookie: CookieJar) -> None:
     auth_args: dict = manager.config_manager.authentication_data
-    if forum not in SupportedDomains.supported_forums_map:
+    forum_dict = {member.name.lower(): member.value for member in SupportedForums}
+
+    if forum not in forum_dict.keys():
         return
-    forum = f"{SupportedDomains.supported_forums_map[forum]}"
+    forum = f"{forum_dict[forum]}"
     with contextlib.suppress(KeyError):
         auth_args["Forums"][f"{forum}_xf_user_cookie"] = cookie._cookies[forum]["/"]["xf_user"].value
         auth_args["Forums"][f"{forum}_xf_user_cookie"] = cookie._cookies["www." + forum]["/"]["xf_user"].value
diff --git a/cyberdrop_dl/utils/data_enums_classes/supported_domains.py b/cyberdrop_dl/utils/data_enums_classes/supported_domains.py
@@ -1,86 +1,94 @@
-from __future__ import annotations
+from aenum import StrEnum, extend_enum
 
-from dataclasses import dataclass
-from typing import ClassVar
+FORUMS = [
+    ("celebforum", "celebforum.to"),
+    ("f95zone", "f95zone.to"),
+    ("leakedmodels", "leakedmodels.com"),
+    ("nudostar", "nudostar.com"),
+    ("xbunker", "xbunker.nu"),
+]
 
+WEBSITES = [
+    ("bunkr", "bunkr"),
+    ("bunkrr", "bunkrr"),
+    ("coomer", "coomer"),
+    ("cyberdrop", "cyberdrop"),
+    ("cyberfile", "cyberfile"),
+    ("e-hentai", "e-hentai"),
+    ("erome", "erome"),
+    ("fapello", "fapello"),
+    ("gofile", "gofile"),
+    ("host.church", "host.church"),
+    ("hotpic", "hotpic"),
+    ("ibb.co", "ibb.co"),
+    ("imageban", "imageban"),
+    ("imagepond.net", "imagepond.net"),
+    ("img.kiwi", "img.kiwi"),
+    ("imgbox", "imgbox"),
+    ("imgur", "imgur"),
+    ("jpeg.pet", "jpeg.pet"),
+    ("jpg.church", "jpg.church"),
+    ("jpg.fish", "jpg.fish"),
+    ("jpg.fishing", "jpg.fishing"),
+    ("jpg.homes", "jpg.homes"),
+    ("jpg.pet", "jpg.pet"),
+    ("jpg1.su", "jpg1.su"),
+    ("jpg2.su", "jpg2.su"),
+    ("jpg3.su", "jpg3.su"),
+    ("jpg4.su", "jpg4.su"),
+    ("jpg5.su", "jpg5.su"),
+    ("kemono", "kemono"),
+    ("mediafire", "mediafire"),
+    ("nudostar.tv", "nudostar.tv"),
+    ("omegascans", "omegascans"),
+    ("pimpandhost", "pimpandhost"),
+    ("pixeldrain", "pixeldrain"),
+    ("postimg", "postimg"),
+    ("realbooru", "realbooru"),
+    ("real-debrid", "real-debrid"),
+    ("redd.it", "redd.it"),
+    ("reddit", "reddit"),
+    ("redgifs", "redgifs"),
+    ("rule34.xxx", "rule34.xxx"),
+    ("rule34.xyz", "rule34.xyz"),
+    ("rule34vault", "rule34vault"),
+    ("saint", "saint"),
+    ("scrolller", "scrolller"),
+    ("socialmediagirls", "socialmediagirls"),
+    ("toonily", "toonily"),
+    ("tokyomotion.net", "tokyomotion.net"),
+    ("xbunkr", "xbunkr"),
+    ("xxxbunker", "xxxbunker"),
+]
 
-@dataclass
-class SupportedDomains:
-    """The allows options for domains to skip when scraping and mappings."""
 
-    supported_hosts: ClassVar[tuple[str, ...]] = (
-        "bunkr",
-        "bunkrr",
-        "celebforum",
-        "coomer",
-        "cyberdrop",
-        "cyberfile",
-        "e-hentai",
-        "erome",
-        "f95zone",
-        "fapello",
-        "gofile",
-        "host.church",
-        "hotpic",
-        "ibb.co",
-        "imageban",
-        "imagepond.net",
-        "img.kiwi",
-        "imgbox",
-        "imgur",
-        "jpeg.pet",
-        "jpg.church",
-        "jpg.fish",
-        "jpg.fishing",
-        "jpg.homes",
-        "jpg.pet",
-        "jpg1.su",
-        "jpg2.su",
-        "jpg3.su",
-        "jpg4.su",
-        "jpg5.su",
-        "kemono",
-        "leakedmodels",
-        "mediafire",
-        "nudostar.com",
-        "nudostar.tv",
-        "omegascans",
-        "pimpandhost",
-        "pixeldrain",
-        "postimg",
-        "realbooru",
-        "real-debrid",
-        "redd.it",
-        "reddit",
-        "redgifs",
-        "rule34.xxx",
-        "rule34.xyz",
-        "rule34vault",
-        "saint",
-        "scrolller",
-        "socialmediagirls",
-        "toonily",
-        "tokyomotion.net",
-        "xbunker",
-        "xbunkr",
-        "xxxbunker",
-    )
+class SupportedForums(StrEnum):
+    pass
 
-    supported_forums: ClassVar[tuple[str, ...]] = (
-        "celebforum.to",
-        "f95zone.to",
-        "forums.socialmediagirls.com",
-        "leakedmodels.com",
-        "nudostar.com",
-        "xbunker.nu",
-    )
-    supported_forums_map: ClassVar[dict[str, str]] = {
-        "celebforum.to": "celebforum",
-        "f95zone.to": "f95zone",
-        "forums.socialmediagirls.com": "socialmediagirls",
-        "leakedmodels.com": "leakedmodels",
-        "nudostar.com": "nudostar",
-        "xbunker.nu": "xbunker",
-    }
 
-    sites: list[str]
+class SupportedHosts(StrEnum):
+    pass
+
+
+class SupportedHostsDebug(StrEnum):
+    pass
+
+
+class SupportedSites(StrEnum):
+    pass
+
+
+class SupportedSitesDebug(StrEnum):
+    pass
+
+
+for site in FORUMS:
+    extend_enum(SupportedForums, site[0], site[1])
+for site in WEBSITES:
+    extend_enum(SupportedHosts, site[0], site[1])
+for site in [*WEBSITES, ("simpcity", "simpcity")]:
+    extend_enum(SupportedHostsDebug, site[0], site[1])
+for site in [*WEBSITES, *FORUMS]:
+    extend_enum(SupportedSites, site[0], site[1])
+for site in [*WEBSITES, *FORUMS, ("simpcity", "simpcity")]:
+    extend_enum(SupportedSitesDebug, site[0], site[1])
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cyberdrop-dl-patched"
-version = "5.7.2"
+version = "5.7.3.dev1"
 description = "Bulk downloader for multiple file hosts"
 authors = ["Jacob B <[email protected]>"]
 readme = "README.md"
@@ -35,9 +35,9 @@ lz4 = "^4.3.3"
 pycryptodomex = "^3.21.0"
 jeepney =   [{platform = "linux", version = "^0.8.0"},{platform = "bsd", version = "^0.8.0"}]
 xxhash = "^3.5.0"
-humanfriendly = "^10.0"
 pyreadline3 = "^3.5.4"
 pydantic = "^2.10.2"
+aenum = "^3.1.15"
 
 [tool.poetry.scripts]
 cyberdrop-dl = "cyberdrop_dl.main:main"