From 570520af290331a43f06ffa3f3e19aa58092cb58 Mon Sep 17 00:00:00 2001 From: ZhenShuo2021 <98386542+ZhenShuo2021@users.noreply.github.com> Date: Tue, 12 Nov 2024 02:49:45 +0800 Subject: [PATCH] refactor: move functions --- v2dl/__init__.py | 1 + v2dl/config.py | 86 ------------------------------------ v2dl/option.py | 77 ++++++++++++++++++++++++++++++++ v2dl/utils/security_utils.py | 1 + v2dl/utils/utils.py | 9 ++++ v2dl/v2dl.py | 5 ++- 6 files changed, 91 insertions(+), 88 deletions(-) create mode 100644 v2dl/option.py diff --git a/v2dl/__init__.py b/v2dl/__init__.py index 442c298..8d76880 100644 --- a/v2dl/__init__.py +++ b/v2dl/__init__.py @@ -20,4 +20,5 @@ "FileProcessingError", "DownloadError", "get_bot", + "__version__", ] diff --git a/v2dl/config.py b/v2dl/config.py index 2814c83..0372a70 100644 --- a/v2dl/config.py +++ b/v2dl/config.py @@ -1,8 +1,5 @@ -import argparse -import logging import os import platform -import sys from dataclasses import dataclass from pathlib import Path from typing import Any @@ -151,86 +148,3 @@ def _merge_config(base: dict[str, Any], custom: dict[str, Any]) -> dict: else: base[key] = value return base - - -def parse_arguments(): - formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=36) - parser = argparse.ArgumentParser(description="V2PH scraper.", formatter_class=formatter) - - input_group = parser.add_mutually_exclusive_group(required=True) - input_group.add_argument("url", nargs="?", help="URL to scrape") - input_group.add_argument( - "-i", - "--input-file", - metavar="PATH", - help="Path to txt file containing URL list to be downloaded", - ) - input_group.add_argument("-a", "--account", action="store_true", help="Manage account") - input_group.add_argument("--version", action="store_true", help="Show package version") - - parser.add_argument( - "--bot", - dest="bot_type", - default="drission", - type=str, - choices=["selenium", "drission"], - required=False, - help="Type of bot to use (default: drission)", - ) - - parser.add_argument( - "--chrome-args", - type=str, - help="Override Chrome arguments (example: --chrome-args='--arg1//--arg2//--arg3')", - ) - parser.add_argument( - "--user-agent", - type=str, - help="Override user-agent (example: --user-agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64)...')", - ) - - parser.add_argument("--dry-run", action="store_true", help="Dry run without downloading") - parser.add_argument("--no-skip", action="store_true", help="Do not skip downloaded files") - parser.add_argument("--terminate", action="store_true", help="Terminate chrome after scraping") - parser.add_argument( - "--use-default-chrome-profile", - action="store_true", - help="Use default chrome profile. Using default profile with an operating chrome is not valid", - ) - - log_group = parser.add_mutually_exclusive_group() - log_group.add_argument("-q", "--quiet", action="store_true", help="Quiet mode") - log_group.add_argument("-v", "--verbose", action="store_true", help="Verbose mode") - log_group.add_argument( - "--log-level", default=None, type=int, choices=range(1, 6), help="Set log level (1~5)" - ) - - args = parser.parse_args() - - if args.quiet: - log_level = logging.ERROR - elif args.verbose: - log_level = logging.DEBUG - elif args.log_level is not None: - log_level_mapping = { - 1: logging.DEBUG, - 2: logging.INFO, - 3: logging.WARNING, - 4: logging.WARNING, - 5: logging.CRITICAL, - } - log_level = log_level_mapping.get(args.log_level, logging.INFO) - else: - log_level = logging.INFO - - args.chrome_args = args.chrome_args.split("//") if args.chrome_args else None - - return args, log_level - - -def check_input_file(input_path: str): - if input_path and not os.path.isfile(input_path): - logging.error("Input file %s does not exist.", input_path) - sys.exit(1) - else: - logging.info("Input file %s exists and is accessible.", input_path) diff --git a/v2dl/option.py b/v2dl/option.py new file mode 100644 index 0000000..7ff1916 --- /dev/null +++ b/v2dl/option.py @@ -0,0 +1,77 @@ +import argparse +import logging + + +def parse_arguments(): + formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=36) + parser = argparse.ArgumentParser(description="V2PH scraper.", formatter_class=formatter) + + input_group = parser.add_mutually_exclusive_group(required=True) + input_group.add_argument("url", nargs="?", help="URL to scrape") + input_group.add_argument( + "-i", + "--input-file", + metavar="PATH", + help="Path to txt file containing URL list to be downloaded", + ) + input_group.add_argument("-a", "--account", action="store_true", help="Manage account") + input_group.add_argument("--version", action="store_true", help="Show package version") + + parser.add_argument( + "--bot", + dest="bot_type", + default="drission", + type=str, + choices=["selenium", "drission"], + required=False, + help="Type of bot to use (default: drission)", + ) + + parser.add_argument( + "--chrome-args", + type=str, + help="Override Chrome arguments (example: --chrome-args='--arg1//--arg2//--arg3')", + ) + parser.add_argument( + "--user-agent", + type=str, + help="Override user-agent (example: --user-agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64)...')", + ) + + parser.add_argument("--dry-run", action="store_true", help="Dry run without downloading") + parser.add_argument("--no-skip", action="store_true", help="Do not skip downloaded files") + parser.add_argument("--terminate", action="store_true", help="Terminate chrome after scraping") + parser.add_argument( + "--use-default-chrome-profile", + action="store_true", + help="Use default chrome profile. Using default profile with an operating chrome is not valid", + ) + + log_group = parser.add_mutually_exclusive_group() + log_group.add_argument("-q", "--quiet", action="store_true", help="Quiet mode") + log_group.add_argument("-v", "--verbose", action="store_true", help="Verbose mode") + log_group.add_argument( + "--log-level", default=None, type=int, choices=range(1, 6), help="Set log level (1~5)" + ) + + args = parser.parse_args() + + if args.quiet: + log_level = logging.ERROR + elif args.verbose: + log_level = logging.DEBUG + elif args.log_level is not None: + log_level_mapping = { + 1: logging.DEBUG, + 2: logging.INFO, + 3: logging.WARNING, + 4: logging.WARNING, + 5: logging.CRITICAL, + } + log_level = log_level_mapping.get(args.log_level, logging.INFO) + else: + log_level = logging.INFO + + args.chrome_args = args.chrome_args.split("//") if args.chrome_args else None + + return args, log_level diff --git a/v2dl/utils/security_utils.py b/v2dl/utils/security_utils.py index 1d55def..9965428 100644 --- a/v2dl/utils/security_utils.py +++ b/v2dl/utils/security_utils.py @@ -52,6 +52,7 @@ def encrypt_master_key(self, master_key: bytes) -> tuple[bytes, bytes, bytes]: # Derive the encryption key using scrypt import time + t = time.time() derived_key = argon2id.kdf( self.KEY_BYTES, diff --git a/v2dl/utils/utils.py b/v2dl/utils/utils.py index 8b42687..8cbc080 100644 --- a/v2dl/utils/utils.py +++ b/v2dl/utils/utils.py @@ -1,6 +1,7 @@ import logging import os import re +import sys import threading import time from dataclasses import dataclass @@ -329,3 +330,11 @@ def get_image_extension(url: str) -> str: else: # 如果沒找到,返回預設值 return "jpg" + + +def check_input_file(input_path: str): + if input_path and not os.path.isfile(input_path): + logging.error("Input file %s does not exist.", input_path) + sys.exit(1) + else: + logging.info("Input file %s exists and is accessible.", input_path) diff --git a/v2dl/v2dl.py b/v2dl/v2dl.py index 552bac3..9808991 100644 --- a/v2dl/v2dl.py +++ b/v2dl/v2dl.py @@ -2,12 +2,13 @@ import sys from .account_cli import cli -from .config import Config, ConfigManager, RuntimeConfig, check_input_file, parse_arguments +from .config import Config, ConfigManager, RuntimeConfig from .const import DEFAULT_CONFIG from .error import ScrapeError from .logger import setup_logging +from .option import parse_arguments from .scrapper import ScrapeHandler -from .utils.utils import ThreadingService +from .utils.utils import ThreadingService, check_input_file from .version import __version__ from .web_bot import get_bot