76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
import aiofiles
|
|
import aiohttp
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
from utils.validators import sanitize_filename
|
|
from utils.logger import setup_logger
|
|
|
|
logger = setup_logger('FileManager')
|
|
|
|
|
|
class FileManager:
|
|
def __init__(self, base_path: str = './Downloaded'):
|
|
self.base_path = Path(base_path)
|
|
self.base_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
def get_save_path(self, author_name: str, mode: str = None, aweme_title: str = None,
|
|
aweme_id: str = None, folderstyle: bool = True) -> Path:
|
|
safe_author = sanitize_filename(author_name)
|
|
|
|
if mode:
|
|
save_dir = self.base_path / safe_author / mode
|
|
else:
|
|
save_dir = self.base_path / safe_author
|
|
|
|
if folderstyle and aweme_title and aweme_id:
|
|
safe_title = sanitize_filename(aweme_title)
|
|
save_dir = save_dir / f"{safe_title}_{aweme_id}"
|
|
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
|
return save_dir
|
|
|
|
async def download_file(
|
|
self,
|
|
url: str,
|
|
save_path: Path,
|
|
session: aiohttp.ClientSession = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
) -> bool:
|
|
should_close = False
|
|
if session is None:
|
|
default_headers = headers or {
|
|
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
|
|
'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36',
|
|
'Referer': 'https://www.douyin.com/',
|
|
'Accept': '*/*',
|
|
}
|
|
session = aiohttp.ClientSession(headers=default_headers)
|
|
should_close = True
|
|
|
|
try:
|
|
async with session.get(
|
|
url,
|
|
timeout=aiohttp.ClientTimeout(total=300),
|
|
headers=headers,
|
|
) as response:
|
|
if response.status == 200:
|
|
async with aiofiles.open(save_path, 'wb') as f:
|
|
async for chunk in response.content.iter_chunked(8192):
|
|
await f.write(chunk)
|
|
return True
|
|
else:
|
|
logger.error(f"Download failed: {url}, status: {response.status}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Download error: {url}, error: {e}")
|
|
return False
|
|
finally:
|
|
if should_close:
|
|
await session.close()
|
|
|
|
def file_exists(self, file_path: Path) -> bool:
|
|
return file_path.exists() and file_path.stat().st_size > 0
|
|
|
|
def get_file_size(self, file_path: Path) -> int:
|
|
return file_path.stat().st_size if self.file_exists(file_path) else 0
|