| 1 | import json |
| 2 | import logging |
| 3 | import os |
| 4 | import sys |
| 5 | from copy import deepcopy |
| 6 | from pathlib import Path |
| 7 | from typing import Any, Dict, List, Optional |
| 8 | |
| 9 | import yaml |
| 10 | |
| 11 | from utils.cookie_utils import parse_cookie_header, sanitize_cookies |
| 12 | |
| 13 | from .default_config import DEFAULT_CONFIG |
| 14 | |
| 15 | logger = logging.getLogger("ConfigLoader") |
| 16 | |
| 17 | |
| 18 | class ConfigLoader: |
| 19 | def __init__(self, config_path: Optional[str] = None): |
| 20 | self.config_path = config_path |
| 21 | self.config = self._load_config() |
| 22 | |
| 23 | def _load_config(self) -> Dict[str, Any]: |
| 24 | config = deepcopy(DEFAULT_CONFIG) |
| 25 | override_sources: List[Dict[str, Any]] = [] |
| 26 | |
| 27 | if self.config_path and os.path.exists(self.config_path): |
| 28 | with open(self.config_path, "r", encoding="utf-8") as f: |
| 29 | file_config = yaml.safe_load(f) or {} |
| 30 | config = self._merge_config(config, file_config) |
| 31 | override_sources.append(file_config) |
| 32 | |
| 33 | env_config = self._load_env_config() |
| 34 | if env_config: |
| 35 | config = self._merge_config(config, env_config) |
| 36 | override_sources.append(env_config) |
| 37 | |
| 38 | return self._normalize_mix_aliases(config, override_sources) |
| 39 | |
| 40 | def _merge_config(self, base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: |
| 41 | result = base.copy() |
| 42 | for key, value in override.items(): |
| 43 | if key in result and isinstance(result[key], dict) and isinstance(value, dict): |
| 44 | result[key] = self._merge_config(result[key], value) |
| 45 | else: |
| 46 | result[key] = value |
| 47 | return result |
| 48 | |
| 49 | def _load_env_config(self) -> Dict[str, Any]: |
| 50 | env_config = {} |
| 51 | if os.getenv("DOUYIN_COOKIE"): |
| 52 | env_config["cookie"] = os.getenv("DOUYIN_COOKIE") |
| 53 | if os.getenv("DOUYIN_PATH"): |
| 54 | env_config["path"] = os.getenv("DOUYIN_PATH") |
| 55 | if os.getenv("DOUYIN_THREAD"): |
| 56 | try: |
| 57 | env_config["thread"] = int(os.getenv("DOUYIN_THREAD")) |
| 58 | except (TypeError, ValueError): |
| 59 | logger.warning( |
| 60 | "Invalid DOUYIN_THREAD value: %s, ignoring", |
| 61 | os.getenv("DOUYIN_THREAD"), |
| 62 | ) |
| 63 | if os.getenv("DOUYIN_PROXY"): |
| 64 | env_config["proxy"] = os.getenv("DOUYIN_PROXY") |
| 65 | return env_config |
| 66 | |
| 67 | def _normalize_mix_aliases( |
| 68 | self, config: Dict[str, Any], override_sources: List[Dict[str, Any]] |
| 69 | ) -> Dict[str, Any]: |
| 70 | # canonical key 为 mix，allmix 作为兼容别名保留并同步 |
| 71 | normalization_rules = ( |
| 72 | ("number", 0), |
| 73 | ("increase", False), |
| 74 | ) |
| 75 | for section, default_value in normalization_rules: |
| 76 | section_config = config.get(section) |
| 77 | if not isinstance(section_config, dict): |
| 78 | section_config = {} |
| 79 | config[section] = section_config |
| 80 | |
| 81 | mix_value = section_config.get("mix") |
| 82 | allmix_value = section_config.get("allmix") |
| 83 | section_default = DEFAULT_CONFIG.get(section, {}) |
| 84 | default_mix_value = ( |
| 85 | section_default.get("mix", default_value) |
| 86 | if isinstance(section_default, dict) |
| 87 | else default_value |
| 88 | ) |
| 89 | default_allmix_value = ( |
| 90 | section_default.get("allmix", default_value) |
| 91 | if isinstance(section_default, dict) |
| 92 | else default_value |
| 93 | ) |
| 94 | |
| 95 | mix_is_default = mix_value == default_mix_value |
| 96 | allmix_is_default = allmix_value == default_allmix_value |
| 97 | |
| 98 | mix_explicit = self._is_key_explicit_in_sources(override_sources, section, "mix") |
| 99 | allmix_explicit = self._is_key_explicit_in_sources(override_sources, section, "allmix") |
| 100 | |
| 101 | if mix_explicit: |
| 102 | canonical_value = mix_value |
| 103 | if allmix_explicit and allmix_value != mix_value: |
| 104 | logger.warning( |
| 105 | "mix/allmix conflict detected in %s: mix=%s, allmix=%s; using mix=%s", |
| 106 | section, |
| 107 | mix_value, |
| 108 | allmix_value, |
| 109 | mix_value, |
| 110 | ) |
| 111 | elif allmix_explicit: |
| 112 | canonical_value = allmix_value |
| 113 | elif not mix_is_default: |
| 114 | canonical_value = mix_value |
| 115 | if not allmix_is_default and allmix_value != mix_value: |
| 116 | logger.warning( |
| 117 | "mix/allmix conflict detected in %s: mix=%s, allmix=%s; using mix=%s", |
| 118 | section, |
| 119 | mix_value, |
| 120 | allmix_value, |
| 121 | mix_value, |
| 122 | ) |
| 123 | elif not allmix_is_default: |
| 124 | canonical_value = allmix_value |
| 125 | else: |
| 126 | canonical_value = default_value |
| 127 | |
| 128 | section_config["mix"] = canonical_value |
| 129 | section_config["allmix"] = canonical_value |
| 130 | |
| 131 | return config |
| 132 | |
| 133 | @staticmethod |
| 134 | def _is_key_explicit_in_sources(sources: List[Dict[str, Any]], section: str, key: str) -> bool: |
| 135 | for source in sources: |
| 136 | if not isinstance(source, dict): |
| 137 | continue |
| 138 | section_value = source.get(section) |
| 139 | if isinstance(section_value, dict) and key in section_value: |
| 140 | return True |
| 141 | return False |
| 142 | |
| 143 | def update(self, **kwargs): |
| 144 | for key, value in kwargs.items(): |
| 145 | if key in self.config: |
| 146 | if isinstance(self.config[key], dict) and isinstance(value, dict): |
| 147 | self.config[key].update(value) |
| 148 | else: |
| 149 | self.config[key] = value |
| 150 | else: |
| 151 | self.config[key] = value |
| 152 | |
| 153 | # Keys that the desktop Settings UI lets the user edit. ``save()`` writes |
| 154 | # these back to the YAML config so changes survive a sidecar restart. |
| 155 | # Kept explicit rather than dumping everything so we don't accidentally |
| 156 | # persist runtime/secret values (cookies, links, etc.) that should stay |
| 157 | # out of the on-disk config, and so fields the user added manually to |
| 158 | # their config.yml are left untouched. |
| 159 | _UI_PERSISTED_KEYS = ( |
| 160 | "path", |
| 161 | "thread", |
| 162 | "rate_limit", |
| 163 | "cover", |
| 164 | "music", |
| 165 | "avatar", |
| 166 | "json", |
| 167 | "download_pinned", |
| 168 | "proxy", |
| 169 | "retry_times", |
| 170 | "folderstyle", |
| 171 | "filename_template", |
| 172 | "folder_template", |
| 173 | "comments", |
| 174 | "live", |
| 175 | "transcript", |
| 176 | "notifications", |
| 177 | ) |
| 178 | |
| 179 | def save(self) -> bool: |
| 180 | """Persist UI-editable keys back to ``self.config_path``. |
| 181 | |
| 182 | Returns True when a file was written, False when no config path is |
| 183 | set (e.g. ``ConfigLoader(None)`` in unit tests). Any existing keys |
| 184 | the user put in the YAML file manually are preserved by merging |
| 185 | the UI keys on top of the previously loaded file contents. |
| 186 | |
| 187 | I/O errors are logged and surfaced via the return value rather than |
| 188 | raised, so a read-only disk doesn't take down the HTTP handler that |
| 189 | called us. |
| 190 | """ |
| 191 | if not self.config_path: |
| 192 | return False |
| 193 | target = Path(self.config_path) |
| 194 | try: |
| 195 | target.parent.mkdir(parents=True, exist_ok=True) |
| 196 | except OSError as exc: |
| 197 | logger.warning("Cannot create config directory %s: %s", target.parent, exc) |
| 198 | return False |
| 199 | |
| 200 | existing: Dict[str, Any] = {} |
| 201 | if target.exists(): |
| 202 | try: |
| 203 | with open(target, "r", encoding="utf-8") as handle: |
| 204 | loaded = yaml.safe_load(handle) |
| 205 | if isinstance(loaded, dict): |
| 206 | existing = loaded |
| 207 | except (yaml.YAMLError, OSError) as exc: |
| 208 | logger.warning( |
| 209 | "Failed to read existing config %s for merge: %s; " |
| 210 | "falling back to UI-only snapshot", |
| 211 | target, |
| 212 | exc, |
| 213 | ) |
| 214 | existing = {} |
| 215 | |
| 216 | for key in self._UI_PERSISTED_KEYS: |
| 217 | if key in self.config: |
| 218 | value = self.config[key] |
| 219 | # Defensive copy so a later ``config.update`` on the same |
| 220 | # process doesn't mutate what we just serialised. |
| 221 | if isinstance(value, dict): |
| 222 | value = deepcopy(value) |
| 223 | elif isinstance(value, list): |
| 224 | value = list(value) |
| 225 | existing[key] = value |
| 226 | |
| 227 | try: |
| 228 | with open(target, "w", encoding="utf-8") as handle: |
| 229 | yaml.safe_dump(existing, handle, allow_unicode=True, sort_keys=False) |
| 230 | except OSError as exc: |
| 231 | logger.warning("Failed to write config %s: %s", target, exc) |
| 232 | return False |
| 233 | |
| 234 | # ``transcript.api_key`` (Requirement 5) and other potentially |
| 235 | # sensitive fields land inside this file as plaintext. On POSIX |
| 236 | # we tighten permissions to 0o600 (owner read/write only) right |
| 237 | # after the write so a co-located malicious process or another |
| 238 | # local user can't ``cat`` the key. On Windows we skip — POSIX |
| 239 | # mode bits aren't meaningful, ACLs are the right knob, and |
| 240 | # changing them belongs to a separate hardening task. |
| 241 | if sys.platform != "win32": |
| 242 | try: |
| 243 | os.chmod(target, 0o600) |
| 244 | except OSError as exc: |
| 245 | logger.warning( |
| 246 | "settings_chmod_failed: path=%s error=%r", target, exc |
| 247 | ) |
| 248 | return True |
| 249 | |
| 250 | def get(self, key: str, default: Any = None) -> Any: |
| 251 | return self.config.get(key, default) |
| 252 | |
| 253 | def get_cookies(self) -> Dict[str, str]: |
| 254 | cookies_config = self.config.get("cookies") or self.config.get("cookie") |
| 255 | |
| 256 | if isinstance(cookies_config, str): |
| 257 | if cookies_config.strip().lower() == "auto": |
| 258 | return self._load_auto_cookies() |
| 259 | return self._parse_cookie_string(cookies_config) |
| 260 | elif isinstance(cookies_config, dict): |
| 261 | return sanitize_cookies(cookies_config) |
| 262 | if self._auto_cookie_enabled(): |
| 263 | return self._load_auto_cookies() |
| 264 | return {} |
| 265 | |
| 266 | def _parse_cookie_string(self, cookie_str: str) -> Dict[str, str]: |
| 267 | return sanitize_cookies(parse_cookie_header(cookie_str)) |
| 268 | |
| 269 | def _auto_cookie_enabled(self) -> bool: |
| 270 | raw_value = self.config.get("auto_cookie") |
| 271 | if isinstance(raw_value, str): |
| 272 | return raw_value.strip().lower() in {"1", "true", "yes", "on"} |
| 273 | return bool(raw_value) |
| 274 | |
| 275 | def _load_auto_cookies(self) -> Dict[str, str]: |
| 276 | for path in self._candidate_auto_cookie_paths(): |
| 277 | cookies = self._load_cookie_file(path) |
| 278 | if cookies is None: |
| 279 | continue |
| 280 | if cookies: |
| 281 | logger.info("Loaded auto cookies from %s", path) |
| 282 | return cookies |
| 283 | return {} |
| 284 | |
| 285 | def _candidate_auto_cookie_paths(self) -> List[Path]: |
| 286 | config_dir = ( |
| 287 | Path(self.config_path).resolve().parent if self.config_path else Path.cwd().resolve() |
| 288 | ) |
| 289 | search_roots = [ |
| 290 | config_dir, |
| 291 | config_dir.parent, |
| 292 | Path.cwd().resolve(), |
| 293 | ] |
| 294 | candidates: List[Path] = [] |
| 295 | for root in search_roots: |
| 296 | candidates.extend( |
| 297 | [ |
| 298 | root / "config" / "cookies.json", |
| 299 | root / ".cookies.json", |
| 300 | ] |
| 301 | ) |
| 302 | |
| 303 | unique: List[Path] = [] |
| 304 | seen: set[str] = set() |
| 305 | for candidate in candidates: |
| 306 | resolved = str(candidate.resolve()) |
| 307 | if resolved in seen: |
| 308 | continue |
| 309 | seen.add(resolved) |
| 310 | unique.append(candidate) |
| 311 | return unique |
| 312 | |
| 313 | @staticmethod |
| 314 | def _load_cookie_file(path: Path) -> Optional[Dict[str, str]]: |
| 315 | if not path.exists(): |
| 316 | return None |
| 317 | try: |
| 318 | raw_data = json.loads(path.read_text(encoding="utf-8")) |
| 319 | except Exception as exc: |
| 320 | logger.warning("Failed to load auto cookie file %s: %s", path, exc) |
| 321 | return {} |
| 322 | |
| 323 | if raw_data is None: |
| 324 | return {} |
| 325 | if not isinstance(raw_data, dict): |
| 326 | logger.warning("Auto cookie file %s is not a JSON object", path) |
| 327 | return {} |
| 328 | return sanitize_cookies(raw_data) |
| 329 | |
| 330 | def get_links(self) -> List[str]: |
| 331 | links = self.config.get("link", []) |
| 332 | if isinstance(links, str): |
| 333 | return [links] |
| 334 | return links |
| 335 | |
| 336 | def validate(self) -> bool: |
| 337 | if not self.get_links(): |
| 338 | return False |
| 339 | if not self.config.get("path"): |
| 340 | return False |
| 341 | |
| 342 | thread = self.config.get("thread") |
| 343 | if thread is not None: |
| 344 | try: |
| 345 | thread_val = int(thread) |
| 346 | if thread_val < 1: |
| 347 | raise ValueError |
| 348 | self.config["thread"] = thread_val |
| 349 | except (TypeError, ValueError): |
| 350 | logger.warning("Invalid thread value: %s, using default 5", thread) |
| 351 | self.config["thread"] = 5 |
| 352 | |
| 353 | retry_times = self.config.get("retry_times") |
| 354 | if retry_times is not None: |
| 355 | try: |
| 356 | retry_val = int(retry_times) |
| 357 | if retry_val < 0: |
| 358 | raise ValueError |
| 359 | self.config["retry_times"] = retry_val |
| 360 | except (TypeError, ValueError): |
| 361 | logger.warning("Invalid retry_times value: %s, using default 3", retry_times) |
| 362 | self.config["retry_times"] = 3 |
| 363 | |
| 364 | for field in ("start_time", "end_time"): |
| 365 | value = self.config.get(field) |
| 366 | if value and isinstance(value, str): |
| 367 | from datetime import datetime |
| 368 | |
| 369 | try: |
| 370 | datetime.strptime(value, "%Y-%m-%d") |
| 371 | except ValueError: |
| 372 | logger.warning( |
| 373 | "Invalid %s format: %s (expected YYYY-MM-DD), clearing", field, value |
| 374 | ) |
| 375 | self.config[field] = "" |
| 376 | |
| 377 | return True |
| 378 |