| 1 | import os |
| 2 | from pathlib import Path |
| 3 | from typing import Dict, Optional, Union |
| 4 | |
| 5 | import aiofiles |
| 6 | import aiohttp |
| 7 | |
| 8 | from utils.logger import setup_logger |
| 9 | from utils.validators import sanitize_filename |
| 10 | |
| 11 | logger = setup_logger("FileManager") |
| 12 | |
| 13 | |
| 14 | class FileManager: |
| 15 | _IMAGE_CONTENT_TYPE_SUFFIXES = { |
| 16 | "image/gif": ".gif", |
| 17 | "image/jpeg": ".jpg", |
| 18 | "image/jpg": ".jpg", |
| 19 | "image/png": ".png", |
| 20 | "image/webp": ".webp", |
| 21 | } |
| 22 | |
| 23 | # 作者目录层可选风格（与 DEFAULT_CONFIG["author_dir"]、REST SettingsPatch |
| 24 | # 的 Literal、前端下拉三处保持一致）。 |
| 25 | _AUTHOR_DIR_STYLES = ("nickname", "sec_uid", "nickname_uid") |
| 26 | |
| 27 | def __init__(self, base_path: str = "./Downloaded"): |
| 28 | self.base_path = Path(base_path) |
| 29 | self.base_path.mkdir(parents=True, exist_ok=True) |
| 30 | |
| 31 | def get_save_path( |
| 32 | self, |
| 33 | author_name: str, |
| 34 | mode: str = None, |
| 35 | aweme_title: str = None, |
| 36 | aweme_id: str = None, |
| 37 | folderstyle: bool = True, |
| 38 | download_date: str = "", |
| 39 | folder_name: Optional[str] = None, |
| 40 | *, |
| 41 | author_sec_uid: Optional[str] = None, |
| 42 | author_dir_style: str = "nickname", |
| 43 | ) -> Path: |
| 44 | """Compute (and create) the destination directory for a download. |
| 45 | |
| 46 | ``folder_name`` is the pre-rendered, already-sanitized leaf directory |
| 47 | name produced by ``utils.naming.render_template``. When provided, it |
| 48 | overrides the legacy ``{date}_{title}_{id}`` layout. When omitted we |
| 49 | fall back to the historical composition so external callers and the |
| 50 | sibling CLI project keep working unchanged. |
| 51 | |
| 52 | ``author_dir_style`` controls how the author-level directory is |
| 53 | composed (see :data:`_AUTHOR_DIR_STYLES`). Unknown values or missing |
| 54 | ``author_sec_uid`` fall back to ``nickname`` with a ``WARNING`` so |
| 55 | downloads never fail on a misconfiguration. |
| 56 | """ |
| 57 | safe_author = self._compose_author_dir(author_name, author_sec_uid, author_dir_style) |
| 58 | |
| 59 | if mode: |
| 60 | save_dir = self.base_path / safe_author / mode |
| 61 | else: |
| 62 | save_dir = self.base_path / safe_author |
| 63 | |
| 64 | if folderstyle: |
| 65 | leaf = folder_name |
| 66 | if leaf is None and aweme_title and aweme_id: |
| 67 | safe_title = sanitize_filename(aweme_title) |
| 68 | date_prefix = f"{download_date}_" if download_date else "" |
| 69 | leaf = f"{date_prefix}{safe_title}_{aweme_id}" |
| 70 | if leaf: |
| 71 | save_dir = save_dir / leaf |
| 72 | |
| 73 | save_dir.mkdir(parents=True, exist_ok=True) |
| 74 | return save_dir |
| 75 | |
| 76 | @classmethod |
| 77 | def _compose_author_dir( |
| 78 | cls, |
| 79 | author_name: str, |
| 80 | author_sec_uid: Optional[str], |
| 81 | style: str, |
| 82 | ) -> str: |
| 83 | """Build the sanitized author-level directory name per ``style``. |
| 84 | |
| 85 | Behaviour matrix (kept in lock-step with the ``author_dir`` option |
| 86 | surfaced in settings UI and ``DEFAULT_CONFIG``): |
| 87 | |
| 88 | - ``nickname`` → ``sanitize_filename(author_name)`` (legacy) |
| 89 | - ``sec_uid`` → ``sanitize_filename(author_sec_uid)``; |
| 90 | empty/None → fall back to nickname + ``logger.warning``. |
| 91 | - ``nickname_uid`` → ``sanitize_filename(f"{author_name}_{author_sec_uid}")``; |
| 92 | sec_uid missing → fall back to nickname + ``logger.warning``. |
| 93 | - Unknown style → fall back to nickname + ``logger.warning``. |
| 94 | |
| 95 | Never raises — a misconfiguration must degrade into a still-working |
| 96 | download, not a hard failure. |
| 97 | """ |
| 98 | nickname_dir = sanitize_filename(author_name) |
| 99 | sec_uid = (author_sec_uid or "").strip() |
| 100 | |
| 101 | if style not in cls._AUTHOR_DIR_STYLES: |
| 102 | logger.warning( |
| 103 | "Unknown author_dir style %r, falling back to nickname (%s)", |
| 104 | style, |
| 105 | nickname_dir, |
| 106 | ) |
| 107 | return nickname_dir |
| 108 | |
| 109 | if style == "nickname": |
| 110 | return nickname_dir |
| 111 | |
| 112 | if style == "sec_uid": |
| 113 | if not sec_uid: |
| 114 | logger.warning( |
| 115 | "author_dir=sec_uid but sec_uid is missing for %r, falling back to nickname", |
| 116 | author_name, |
| 117 | ) |
| 118 | return nickname_dir |
| 119 | return sanitize_filename(sec_uid) |
| 120 | |
| 121 | # style == "nickname_uid" |
| 122 | if not sec_uid: |
| 123 | logger.warning( |
| 124 | "author_dir=nickname_uid but sec_uid is missing for %r, falling back to nickname", |
| 125 | author_name, |
| 126 | ) |
| 127 | return nickname_dir |
| 128 | return sanitize_filename(f"{author_name}_{sec_uid}") |
| 129 | |
| 130 | async def download_file( |
| 131 | self, |
| 132 | url: str, |
| 133 | save_path: Path, |
| 134 | session: aiohttp.ClientSession = None, |
| 135 | headers: Optional[Dict[str, str]] = None, |
| 136 | proxy: Optional[str] = None, |
| 137 | *, |
| 138 | prefer_response_content_type: bool = False, |
| 139 | return_saved_path: bool = False, |
| 140 | ) -> Union[bool, Path]: |
| 141 | should_close = False |
| 142 | if session is None: |
| 143 | default_headers = headers or { |
| 144 | "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " |
| 145 | "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", |
| 146 | "Referer": "https://www.douyin.com/", |
| 147 | "Accept": "*/*", |
| 148 | } |
| 149 | session = aiohttp.ClientSession(headers=default_headers) |
| 150 | should_close = True |
| 151 | |
| 152 | final_path = save_path |
| 153 | tmp_path = save_path.with_suffix(save_path.suffix + ".tmp") |
| 154 | try: |
| 155 | async with session.get( |
| 156 | url, |
| 157 | timeout=aiohttp.ClientTimeout(total=300), |
| 158 | headers=headers, |
| 159 | proxy=proxy or None, |
| 160 | ) as response: |
| 161 | if response.status == 200: |
| 162 | final_path = self._resolve_save_path_from_content_type( |
| 163 | save_path, |
| 164 | response.headers, |
| 165 | prefer_response_content_type=prefer_response_content_type, |
| 166 | ) |
| 167 | tmp_path = final_path.with_suffix(final_path.suffix + ".tmp") |
| 168 | expected_size = response.content_length |
| 169 | written = 0 |
| 170 | async with aiofiles.open(tmp_path, "wb") as f: |
| 171 | async for chunk in response.content.iter_chunked(8192): |
| 172 | await f.write(chunk) |
| 173 | written += len(chunk) |
| 174 | if expected_size is not None and written != expected_size: |
| 175 | logger.warning( |
| 176 | "Size mismatch for %s: expected %d, got %d", |
| 177 | save_path.name, |
| 178 | expected_size, |
| 179 | written, |
| 180 | ) |
| 181 | tmp_path.unlink(missing_ok=True) |
| 182 | return False |
| 183 | os.replace(str(tmp_path), str(final_path)) |
| 184 | return final_path if return_saved_path else True |
| 185 | else: |
| 186 | logger.debug( |
| 187 | "Download failed for %s, status=%s", |
| 188 | final_path.name, |
| 189 | response.status, |
| 190 | ) |
| 191 | return False |
| 192 | except Exception as e: |
| 193 | logger.debug("Download error for %s: %s", final_path.name, e) |
| 194 | tmp_path.unlink(missing_ok=True) |
| 195 | return False |
| 196 | finally: |
| 197 | if should_close: |
| 198 | await session.close() |
| 199 | |
| 200 | @classmethod |
| 201 | def _resolve_save_path_from_content_type( |
| 202 | cls, |
| 203 | save_path: Path, |
| 204 | response_headers, |
| 205 | *, |
| 206 | prefer_response_content_type: bool = False, |
| 207 | ) -> Path: |
| 208 | if not prefer_response_content_type: |
| 209 | return save_path |
| 210 | |
| 211 | content_type = response_headers.get("Content-Type", "") if response_headers else "" |
| 212 | normalized_type = content_type.split(";", 1)[0].strip().lower() |
| 213 | suffix = cls._IMAGE_CONTENT_TYPE_SUFFIXES.get(normalized_type) |
| 214 | if not suffix: |
| 215 | return save_path |
| 216 | return save_path.with_suffix(suffix) |
| 217 | |
| 218 | def file_exists(self, file_path: Path) -> bool: |
| 219 | try: |
| 220 | return file_path.exists() and file_path.stat().st_size > 0 |
| 221 | except OSError: |
| 222 | return False |
| 223 | |
| 224 | def get_file_size(self, file_path: Path) -> int: |
| 225 | try: |
| 226 | return file_path.stat().st_size if self.file_exists(file_path) else 0 |
| 227 | except OSError: |
| 228 | return 0 |
| 229 |