'commit'
This commit is contained in:
HuangHai
309
DouYin/dy-downloader/core/downloader_base.py
Normal file
309
DouYin/dy-downloader/core/downloader_base.py
Normal file
@@ -0,0 +1,309 @@
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from config import ConfigLoader
|
||||
from storage import Database, FileManager, MetadataHandler
|
||||
from auth import CookieManager
|
||||
from control import QueueManager, RateLimiter, RetryHandler
|
||||
from core.api_client import DouyinAPIClient
|
||||
from utils.logger import setup_logger
|
||||
from utils.validators import sanitize_filename
|
||||
|
||||
logger = setup_logger('BaseDownloader')
|
||||
|
||||
|
||||
class DownloadResult:
|
||||
def __init__(self):
|
||||
self.total = 0
|
||||
self.success = 0
|
||||
self.failed = 0
|
||||
self.skipped = 0
|
||||
|
||||
def __str__(self):
|
||||
return f"Total: {self.total}, Success: {self.success}, Failed: {self.failed}, Skipped: {self.skipped}"
|
||||
|
||||
|
||||
class BaseDownloader(ABC):
|
||||
def __init__(
|
||||
self,
|
||||
config: ConfigLoader,
|
||||
api_client: DouyinAPIClient,
|
||||
file_manager: FileManager,
|
||||
cookie_manager: CookieManager,
|
||||
database: Optional[Database] = None,
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
retry_handler: Optional[RetryHandler] = None,
|
||||
queue_manager: Optional[QueueManager] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.api_client = api_client
|
||||
self.file_manager = file_manager
|
||||
self.cookie_manager = cookie_manager
|
||||
self.database = database
|
||||
self.rate_limiter = rate_limiter or RateLimiter()
|
||||
self.retry_handler = retry_handler or RetryHandler()
|
||||
thread_count = int(self.config.get('thread', 5) or 5)
|
||||
self.queue_manager = queue_manager or QueueManager(max_workers=thread_count)
|
||||
self.metadata_handler = MetadataHandler()
|
||||
|
||||
def _download_headers(self, user_agent: Optional[str] = None) -> Dict[str, str]:
|
||||
headers = {
|
||||
'Referer': f'{self.api_client.BASE_URL}/',
|
||||
'Origin': self.api_client.BASE_URL,
|
||||
'Accept': '*/*',
|
||||
}
|
||||
|
||||
headers['User-Agent'] = user_agent or self.api_client.headers.get('User-Agent', '')
|
||||
return headers
|
||||
|
||||
@abstractmethod
|
||||
async def download(self, parsed_url: Dict[str, Any]) -> DownloadResult:
|
||||
pass
|
||||
|
||||
async def _should_download(self, aweme_id: str) -> bool:
|
||||
if self.database:
|
||||
return not await self.database.is_downloaded(aweme_id)
|
||||
return True
|
||||
|
||||
def _filter_by_time(self, aweme_list: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
start_time = self.config.get('start_time')
|
||||
end_time = self.config.get('end_time')
|
||||
|
||||
if not start_time and not end_time:
|
||||
return aweme_list
|
||||
|
||||
filtered: List[Dict[str, Any]] = []
|
||||
for aweme in aweme_list:
|
||||
create_time = aweme.get('create_time', 0)
|
||||
|
||||
if start_time:
|
||||
from datetime import datetime
|
||||
start_ts = int(datetime.strptime(start_time, '%Y-%m-%d').timestamp())
|
||||
if create_time < start_ts:
|
||||
continue
|
||||
|
||||
if end_time:
|
||||
from datetime import datetime
|
||||
end_ts = int(datetime.strptime(end_time, '%Y-%m-%d').timestamp())
|
||||
if create_time > end_ts:
|
||||
continue
|
||||
|
||||
filtered.append(aweme)
|
||||
|
||||
return filtered
|
||||
|
||||
def _limit_count(self, aweme_list: List[Dict[str, Any]], mode: str) -> List[Dict[str, Any]]:
|
||||
number_config = self.config.get('number', {})
|
||||
limit = number_config.get(mode, 0)
|
||||
|
||||
if limit > 0:
|
||||
return aweme_list[:limit]
|
||||
return aweme_list
|
||||
|
||||
async def _download_aweme_assets(
|
||||
self,
|
||||
aweme_data: Dict[str, Any],
|
||||
author_name: str,
|
||||
mode: Optional[str] = None,
|
||||
) -> bool:
|
||||
aweme_id = aweme_data.get('aweme_id')
|
||||
if not aweme_id:
|
||||
logger.error('Missing aweme_id in aweme data')
|
||||
return False
|
||||
|
||||
desc = aweme_data.get('desc', 'no_title')
|
||||
safe_title = sanitize_filename(desc)
|
||||
|
||||
save_dir = self.file_manager.get_save_path(
|
||||
author_name=author_name,
|
||||
mode=mode,
|
||||
aweme_title=desc,
|
||||
aweme_id=aweme_id,
|
||||
folderstyle=self.config.get('folderstyle', True)
|
||||
)
|
||||
|
||||
session = await self.api_client.get_session()
|
||||
|
||||
media_type = self._detect_media_type(aweme_data)
|
||||
if media_type == 'video':
|
||||
video_info = self._build_no_watermark_url(aweme_data)
|
||||
if not video_info:
|
||||
logger.error(f'No playable video URL found for aweme {aweme_id}')
|
||||
return False
|
||||
|
||||
video_url, video_headers = video_info
|
||||
video_path = save_dir / f"{safe_title}_{aweme_id}.mp4"
|
||||
if not await self._download_with_retry(video_url, video_path, session, headers=video_headers):
|
||||
return False
|
||||
|
||||
if self.config.get('cover'):
|
||||
cover_url = self._extract_first_url(aweme_data.get('video', {}).get('cover'))
|
||||
if cover_url:
|
||||
cover_path = save_dir / f"{safe_title}_{aweme_id}_cover.jpg"
|
||||
await self._download_with_retry(
|
||||
cover_url,
|
||||
cover_path,
|
||||
session,
|
||||
headers=self._download_headers(),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
if self.config.get('music'):
|
||||
music_url = self._extract_first_url(aweme_data.get('music', {}).get('play_url'))
|
||||
if music_url:
|
||||
music_path = save_dir / f"{safe_title}_{aweme_id}_music.mp3"
|
||||
await self._download_with_retry(
|
||||
music_url,
|
||||
music_path,
|
||||
session,
|
||||
headers=self._download_headers(),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
elif media_type == 'gallery':
|
||||
image_urls = self._collect_image_urls(aweme_data)
|
||||
if not image_urls:
|
||||
logger.error(f'No images found for aweme {aweme_id}')
|
||||
return False
|
||||
|
||||
for index, image_url in enumerate(image_urls, start=1):
|
||||
suffix = Path(urlparse(image_url).path).suffix or '.jpg'
|
||||
image_path = save_dir / f"{safe_title}_{aweme_id}_{index}{suffix}"
|
||||
success = await self._download_with_retry(
|
||||
image_url,
|
||||
image_path,
|
||||
session,
|
||||
headers=self._download_headers(),
|
||||
)
|
||||
if not success:
|
||||
logger.error(f'Failed downloading image {index} for aweme {aweme_id}')
|
||||
return False
|
||||
else:
|
||||
logger.error(f"Unsupported media type for aweme {aweme_id}: {media_type}")
|
||||
return False
|
||||
|
||||
if self.config.get('avatar'):
|
||||
author = aweme_data.get('author', {})
|
||||
avatar_url = self._extract_first_url(author.get('avatar_larger'))
|
||||
if avatar_url:
|
||||
avatar_path = save_dir / 'avatar.jpg'
|
||||
await self._download_with_retry(
|
||||
avatar_url,
|
||||
avatar_path,
|
||||
session,
|
||||
headers=self._download_headers(),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
if self.config.get('json'):
|
||||
json_path = save_dir / f"{safe_title}_{aweme_id}_data.json"
|
||||
await self.metadata_handler.save_metadata(aweme_data, json_path)
|
||||
|
||||
if self.database:
|
||||
author = aweme_data.get('author', {})
|
||||
metadata_json = json.dumps(aweme_data, ensure_ascii=False)
|
||||
await self.database.add_aweme({
|
||||
'aweme_id': aweme_id,
|
||||
'aweme_type': media_type,
|
||||
'title': desc,
|
||||
'author_id': author.get('uid'),
|
||||
'author_name': author.get('nickname', author_name),
|
||||
'create_time': aweme_data.get('create_time'),
|
||||
'file_path': str(save_dir),
|
||||
'metadata': metadata_json,
|
||||
})
|
||||
|
||||
logger.info(f"Downloaded {media_type}: {desc} ({aweme_id})")
|
||||
return True
|
||||
|
||||
async def _download_with_retry(
|
||||
self,
|
||||
url: str,
|
||||
save_path: Path,
|
||||
session,
|
||||
*,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
optional: bool = False,
|
||||
) -> bool:
|
||||
async def _task():
|
||||
success = await self.file_manager.download_file(url, save_path, session, headers=headers)
|
||||
if not success:
|
||||
raise RuntimeError(f'Download failed for {url}')
|
||||
return True
|
||||
|
||||
try:
|
||||
await self.retry_handler.execute_with_retry(_task)
|
||||
return True
|
||||
except Exception as error:
|
||||
log_fn = logger.warning if optional else logger.error
|
||||
log_fn(f"Download error for {save_path.name}: {error}")
|
||||
return False
|
||||
|
||||
def _detect_media_type(self, aweme_data: Dict[str, Any]) -> str:
|
||||
if aweme_data.get('image_post_info') or aweme_data.get('images'):
|
||||
return 'gallery'
|
||||
return 'video'
|
||||
|
||||
def _build_no_watermark_url(self, aweme_data: Dict[str, Any]) -> Optional[Tuple[str, Dict[str, str]]]:
|
||||
video = aweme_data.get('video', {})
|
||||
play_addr = video.get('play_addr', {})
|
||||
url_candidates = [c for c in (play_addr.get('url_list') or []) if c]
|
||||
url_candidates.sort(key=lambda u: 0 if 'watermark=0' in u else 1)
|
||||
|
||||
fallback_candidate: Optional[Tuple[str, Dict[str, str]]] = None
|
||||
|
||||
for candidate in url_candidates:
|
||||
parsed = urlparse(candidate)
|
||||
headers = self._download_headers()
|
||||
|
||||
if parsed.netloc.endswith('douyin.com'):
|
||||
if 'X-Bogus=' not in candidate:
|
||||
signed_url, ua = self.api_client.sign_url(candidate)
|
||||
headers = self._download_headers(user_agent=ua)
|
||||
return signed_url, headers
|
||||
return candidate, headers
|
||||
|
||||
fallback_candidate = (candidate, headers)
|
||||
|
||||
if fallback_candidate:
|
||||
return fallback_candidate
|
||||
|
||||
uri = play_addr.get('uri') or video.get('vid') or video.get('download_addr', {}).get('uri')
|
||||
if uri:
|
||||
params = {
|
||||
'video_id': uri,
|
||||
'ratio': '1080p',
|
||||
'line': '0',
|
||||
'is_play_url': '1',
|
||||
'watermark': '0',
|
||||
'source': 'PackSourceEnum_PUBLISH',
|
||||
}
|
||||
signed_url, ua = self.api_client.build_signed_path('/aweme/v1/play/', params)
|
||||
return signed_url, self._download_headers(user_agent=ua)
|
||||
|
||||
return None
|
||||
|
||||
def _collect_image_urls(self, aweme_data: Dict[str, Any]) -> List[str]:
|
||||
image_urls: List[str] = []
|
||||
image_post = aweme_data.get('image_post_info', {})
|
||||
images = image_post.get('images') or aweme_data.get('images') or []
|
||||
for item in images:
|
||||
url_list = item.get('url_list') if isinstance(item, dict) else None
|
||||
if url_list:
|
||||
image_urls.append(url_list[0])
|
||||
return image_urls
|
||||
|
||||
@staticmethod
|
||||
def _extract_first_url(source: Any) -> Optional[str]:
|
||||
if isinstance(source, dict):
|
||||
url_list = source.get('url_list')
|
||||
if isinstance(url_list, list) and url_list:
|
||||
return url_list[0]
|
||||
elif isinstance(source, list) and source:
|
||||
return source[0]
|
||||
elif isinstance(source, str):
|
||||
return source
|
||||
return None
|
||||
Reference in New Issue
Block a user