565 lines
20 KiB
Python
565 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Image Downloader Script
|
|
|
|
This script downloads images from a REST API that provides:
|
|
1. An endpoint to list all assets
|
|
2. An endpoint to download individual assets in full resolution
|
|
|
|
Usage:
|
|
python image_downloader.py --api-url <base_url> --list-endpoint <endpoint> --download-endpoint <endpoint> --output-dir <directory>
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import aiohttp
|
|
import aiofiles
|
|
import os
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from urllib.parse import urljoin, urlparse
|
|
from typing import List, Dict, Any, Optional
|
|
import time
|
|
from tqdm import tqdm
|
|
import hashlib
|
|
|
|
# Import the auth manager and asset tracker
|
|
try:
|
|
from auth_manager import AuthManager
|
|
except ImportError:
|
|
AuthManager = None
|
|
|
|
try:
|
|
from asset_tracker import AssetTracker
|
|
except ImportError:
|
|
AssetTracker = None
|
|
|
|
|
|
class ImageDownloader:
|
|
def __init__(self, api_url: str, list_endpoint: str, download_endpoint: str,
|
|
output_dir: str, max_concurrent: int = 5, timeout: int = 30, api_key: str = None,
|
|
email: str = None, password: str = None, track_assets: bool = True):
|
|
"""
|
|
Initialize the image downloader.
|
|
|
|
Args:
|
|
api_url: Base URL of the API
|
|
list_endpoint: Endpoint to get the list of assets
|
|
download_endpoint: Endpoint to download individual assets
|
|
output_dir: Directory to save downloaded images
|
|
max_concurrent: Maximum number of concurrent downloads
|
|
timeout: Request timeout in seconds
|
|
api_key: API key for authentication
|
|
email: Email for login authentication
|
|
password: Password for login authentication
|
|
track_assets: Whether to enable asset tracking to avoid re-downloads
|
|
"""
|
|
self.api_url = api_url.rstrip('/')
|
|
self.list_endpoint = list_endpoint.lstrip('/')
|
|
self.download_endpoint = download_endpoint.lstrip('/')
|
|
self.output_dir = Path(output_dir)
|
|
self.max_concurrent = max_concurrent
|
|
self.timeout = timeout
|
|
self.api_key = api_key
|
|
self.email = email
|
|
self.password = password
|
|
self.auth_manager = None
|
|
|
|
# Create output directory if it doesn't exist
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Setup logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.FileHandler(self.output_dir / 'download.log'),
|
|
logging.StreamHandler()
|
|
]
|
|
)
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
# Initialize asset tracker if enabled and available
|
|
self.asset_tracker = None
|
|
if track_assets and AssetTracker:
|
|
self.asset_tracker = AssetTracker(storage_dir=str(self.output_dir))
|
|
self.logger.info("Asset tracking enabled")
|
|
elif track_assets:
|
|
self.logger.warning("Asset tracking requested but AssetTracker not available")
|
|
else:
|
|
self.logger.info("Asset tracking disabled")
|
|
|
|
# Track download statistics
|
|
self.stats = {
|
|
'total': 0,
|
|
'successful': 0,
|
|
'failed': 0,
|
|
'skipped': 0
|
|
}
|
|
|
|
async def authenticate(self):
|
|
"""Perform login authentication if credentials are provided."""
|
|
if self.email and self.password and AuthManager:
|
|
self.logger.info("Attempting login authentication...")
|
|
self.auth_manager = AuthManager(self.api_url)
|
|
success = await self.auth_manager.login(self.email, self.password)
|
|
|
|
if success:
|
|
self.logger.info("Login authentication successful")
|
|
else:
|
|
self.logger.error("Login authentication failed")
|
|
raise Exception("Login authentication failed")
|
|
elif self.email or self.password:
|
|
self.logger.warning("Both email and password must be provided for login authentication")
|
|
raise Exception("Both email and password must be provided for login authentication")
|
|
|
|
async def get_asset_list(self, session: aiohttp.ClientSession) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch the list of assets from the API.
|
|
|
|
Args:
|
|
session: aiohttp session for making requests
|
|
|
|
Returns:
|
|
List of asset dictionaries
|
|
"""
|
|
url = urljoin(self.api_url, self.list_endpoint)
|
|
self.logger.info(f"Fetching asset list from: {url}")
|
|
|
|
try:
|
|
headers = {}
|
|
|
|
# Use API key if provided
|
|
if self.api_key:
|
|
headers['x-api-key'] = self.api_key
|
|
|
|
# Use login authentication if provided
|
|
elif self.auth_manager and self.auth_manager.is_authenticated():
|
|
headers.update(self.auth_manager.get_auth_headers())
|
|
|
|
async with session.get(url, headers=headers, timeout=self.timeout) as response:
|
|
response.raise_for_status()
|
|
data = await response.json()
|
|
|
|
# Handle different response formats
|
|
if isinstance(data, list):
|
|
assets = data
|
|
elif isinstance(data, dict):
|
|
# Common patterns for API responses
|
|
if 'data' in data:
|
|
assets = data['data']
|
|
elif 'results' in data:
|
|
assets = data['results']
|
|
elif 'items' in data:
|
|
assets = data['items']
|
|
else:
|
|
assets = [data] # Single asset
|
|
else:
|
|
raise ValueError(f"Unexpected response format: {type(data)}")
|
|
|
|
self.logger.info(f"Found {len(assets)} assets")
|
|
return assets
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to fetch asset list: {e}")
|
|
raise
|
|
|
|
def get_download_url(self, asset: Dict[str, Any]) -> str:
|
|
"""
|
|
Generate the download URL for an asset.
|
|
|
|
Args:
|
|
asset: Asset dictionary from the API
|
|
|
|
Returns:
|
|
Download URL for the asset
|
|
"""
|
|
# Try different common patterns for asset IDs
|
|
asset_id = None
|
|
|
|
# Common field names for asset identifiers
|
|
id_fields = ['id', 'asset_id', 'image_id', 'file_id', 'uuid', 'key']
|
|
for field in id_fields:
|
|
if field in asset:
|
|
asset_id = asset[field]
|
|
break
|
|
|
|
if asset_id is None:
|
|
# If no ID field found, try to use the asset itself as the ID
|
|
asset_id = str(asset)
|
|
|
|
# Build download URL with required parameters
|
|
from urllib.parse import urlencode
|
|
|
|
params = {
|
|
'key': self.api_key,
|
|
'u': asset.get('updated', '')
|
|
}
|
|
|
|
download_url = urljoin(self.api_url, f"/v1/media/{asset_id}/full?{urlencode(params)}")
|
|
return download_url
|
|
|
|
def get_filename(self, asset: Dict[str, Any], url: str) -> str:
|
|
"""
|
|
Generate a filename for the downloaded asset.
|
|
|
|
Args:
|
|
asset: Asset dictionary from the API
|
|
url: Download URL
|
|
|
|
Returns:
|
|
Filename for the asset
|
|
"""
|
|
# Try to get filename from asset metadata
|
|
if 'fileName' in asset:
|
|
filename = asset['fileName']
|
|
elif 'filename' in asset:
|
|
filename = asset['filename']
|
|
elif 'name' in asset:
|
|
filename = asset['name']
|
|
elif 'title' in asset:
|
|
filename = asset['title']
|
|
else:
|
|
# Extract filename from URL
|
|
parsed_url = urlparse(url)
|
|
filename = os.path.basename(parsed_url.path)
|
|
|
|
# If no extension, try to get it from content-type or add default
|
|
if '.' not in filename:
|
|
if 'mimeType' in asset:
|
|
ext = self._get_extension_from_mime(asset['mimeType'])
|
|
elif 'content_type' in asset:
|
|
ext = self._get_extension_from_mime(asset['content_type'])
|
|
else:
|
|
ext = '.jpg' # Default extension
|
|
filename += ext
|
|
|
|
# Sanitize filename
|
|
filename = self._sanitize_filename(filename)
|
|
|
|
# Ensure unique filename
|
|
counter = 1
|
|
original_filename = filename
|
|
while (self.output_dir / filename).exists():
|
|
name, ext = os.path.splitext(original_filename)
|
|
filename = f"{name}_{counter}{ext}"
|
|
counter += 1
|
|
|
|
return filename
|
|
|
|
def _get_extension_from_mime(self, mime_type: str) -> str:
|
|
"""Get file extension from MIME type."""
|
|
mime_to_ext = {
|
|
'image/jpeg': '.jpg',
|
|
'image/jpg': '.jpg',
|
|
'image/png': '.png',
|
|
'image/gif': '.gif',
|
|
'image/webp': '.webp',
|
|
'image/bmp': '.bmp',
|
|
'image/tiff': '.tiff',
|
|
'image/svg+xml': '.svg'
|
|
}
|
|
return mime_to_ext.get(mime_type.lower(), '.jpg')
|
|
|
|
def _sanitize_filename(self, filename: str) -> str:
|
|
"""Sanitize filename by removing invalid characters."""
|
|
# Remove or replace invalid characters
|
|
invalid_chars = '<>:"/\\|?*'
|
|
for char in invalid_chars:
|
|
filename = filename.replace(char, '_')
|
|
|
|
# Remove leading/trailing spaces and dots
|
|
filename = filename.strip('. ')
|
|
|
|
# Ensure filename is not empty
|
|
if not filename:
|
|
filename = 'image'
|
|
|
|
return filename
|
|
|
|
async def download_asset(self, session: aiohttp.ClientSession, asset: Dict[str, Any],
|
|
semaphore: asyncio.Semaphore) -> bool:
|
|
"""
|
|
Download a single asset.
|
|
|
|
Args:
|
|
session: aiohttp session for making requests
|
|
asset: Asset dictionary from the API
|
|
semaphore: Semaphore to limit concurrent downloads
|
|
|
|
Returns:
|
|
True if download was successful, False otherwise
|
|
"""
|
|
async with semaphore:
|
|
try:
|
|
download_url = self.get_download_url(asset)
|
|
filename = self.get_filename(asset, download_url)
|
|
filepath = self.output_dir / filename
|
|
|
|
# Check if file already exists and we're not tracking assets
|
|
if filepath.exists() and not self.asset_tracker:
|
|
self.logger.info(f"Skipping {filename} (already exists)")
|
|
self.stats['skipped'] += 1
|
|
return True
|
|
|
|
self.logger.info(f"Downloading {filename} from {download_url}")
|
|
|
|
async with session.get(download_url, timeout=self.timeout) as response:
|
|
response.raise_for_status()
|
|
|
|
# Get content type to verify it's an image
|
|
content_type = response.headers.get('content-type', '')
|
|
if not content_type.startswith('image/'):
|
|
self.logger.warning(f"Content type is not an image: {content_type}")
|
|
|
|
# Download the file
|
|
async with aiofiles.open(filepath, 'wb') as f:
|
|
async for chunk in response.content.iter_chunked(8192):
|
|
await f.write(chunk)
|
|
|
|
# Set file modification time to match the updated timestamp
|
|
if 'updated' in asset:
|
|
try:
|
|
from datetime import datetime
|
|
import os
|
|
# Parse the ISO timestamp
|
|
updated_time = datetime.fromisoformat(asset['updated'].replace('Z', '+00:00'))
|
|
# Set file modification time
|
|
os.utime(filepath, (updated_time.timestamp(), updated_time.timestamp()))
|
|
self.logger.info(f"Set file modification time to {asset['updated']}")
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to set file modification time: {e}")
|
|
|
|
# Mark asset as downloaded in tracker
|
|
if self.asset_tracker:
|
|
self.asset_tracker.mark_asset_downloaded(asset, filepath, True)
|
|
|
|
self.logger.info(f"Successfully downloaded {filename}")
|
|
self.stats['successful'] += 1
|
|
return True
|
|
|
|
except Exception as e:
|
|
# Mark asset as failed in tracker
|
|
if self.asset_tracker:
|
|
download_url = self.get_download_url(asset)
|
|
filename = self.get_filename(asset, download_url)
|
|
filepath = self.output_dir / filename
|
|
self.asset_tracker.mark_asset_downloaded(asset, filepath, False)
|
|
|
|
self.logger.error(f"Failed to download asset {asset.get('id', 'unknown')}: {e}")
|
|
self.stats['failed'] += 1
|
|
return False
|
|
|
|
async def download_all_assets(self, force_redownload: bool = False):
|
|
"""
|
|
Download all assets from the API.
|
|
|
|
Args:
|
|
force_redownload: If True, download all assets regardless of tracking
|
|
"""
|
|
start_time = time.time()
|
|
|
|
# Create aiohttp session with connection pooling
|
|
connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
|
|
timeout = aiohttp.ClientTimeout(total=self.timeout)
|
|
|
|
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
|
try:
|
|
# Perform authentication if needed
|
|
await self.authenticate()
|
|
|
|
# Get asset list
|
|
all_assets = await self.get_asset_list(session)
|
|
self.logger.info(f"Retrieved {len(all_assets)} total assets from API")
|
|
|
|
if not all_assets:
|
|
self.logger.warning("No assets found to download")
|
|
return
|
|
|
|
# Filter for new/modified assets if tracking is enabled
|
|
if self.asset_tracker and not force_redownload:
|
|
assets = self.asset_tracker.get_new_assets(all_assets)
|
|
self.logger.info(f"Found {len(assets)} new/modified assets to download")
|
|
if len(assets) == 0:
|
|
self.logger.info("All assets are up to date!")
|
|
return
|
|
else:
|
|
assets = all_assets
|
|
if force_redownload:
|
|
self.logger.info("Force redownload enabled - downloading all assets")
|
|
|
|
self.stats['total'] = len(assets)
|
|
|
|
# Create semaphore to limit concurrent downloads
|
|
semaphore = asyncio.Semaphore(self.max_concurrent)
|
|
|
|
# Create tasks for all downloads
|
|
tasks = [
|
|
self.download_asset(session, asset, semaphore)
|
|
for asset in assets
|
|
]
|
|
|
|
# Download all assets with progress bar
|
|
with tqdm(total=len(tasks), desc="Downloading assets") as pbar:
|
|
for coro in asyncio.as_completed(tasks):
|
|
result = await coro
|
|
pbar.update(1)
|
|
pbar.set_postfix({
|
|
'Success': self.stats['successful'],
|
|
'Failed': self.stats['failed'],
|
|
'Skipped': self.stats['skipped']
|
|
})
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error during download process: {e}")
|
|
raise
|
|
|
|
# Print final statistics
|
|
elapsed_time = time.time() - start_time
|
|
self.logger.info(f"Download completed in {elapsed_time:.2f} seconds")
|
|
self.logger.info(f"Statistics: {self.stats}")
|
|
|
|
|
|
def main():
|
|
"""Main function to run the image downloader."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Download images from a REST API",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Basic usage
|
|
python image_downloader.py --api-url "https://api.example.com" \\
|
|
--list-endpoint "/assets" \\
|
|
--download-endpoint "/download" \\
|
|
--output-dir "./images"
|
|
|
|
# With custom concurrent downloads and timeout
|
|
python image_downloader.py --api-url "https://api.example.com" \\
|
|
--list-endpoint "/assets" \\
|
|
--download-endpoint "/download" \\
|
|
--output-dir "./images" \\
|
|
--max-concurrent 10 \\
|
|
--timeout 60
|
|
"""
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--api-url',
|
|
required=True,
|
|
help='Base URL of the API (e.g., https://api.example.com)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--list-endpoint',
|
|
required=True,
|
|
help='Endpoint to get the list of assets (e.g., /assets or /images)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--download-endpoint',
|
|
required=True,
|
|
help='Endpoint to download individual assets (e.g., /download or /assets)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--output-dir',
|
|
required=True,
|
|
help='Directory to save downloaded images'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--max-concurrent',
|
|
type=int,
|
|
default=5,
|
|
help='Maximum number of concurrent downloads (default: 5)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--timeout',
|
|
type=int,
|
|
default=30,
|
|
help='Request timeout in seconds (default: 30)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--api-key',
|
|
help='API key for authentication (x-api-key header)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--email',
|
|
help='Email for login authentication'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--password',
|
|
help='Password for login authentication'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--no-tracking',
|
|
action='store_true',
|
|
help='Disable asset tracking (will re-download all assets)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--force-redownload',
|
|
action='store_true',
|
|
help='Force re-download of all assets, even if already tracked'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--show-stats',
|
|
action='store_true',
|
|
help='Show asset tracking statistics and exit'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--cleanup',
|
|
action='store_true',
|
|
help='Clean up metadata for missing files and exit'
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Handle special commands first
|
|
if args.show_stats or args.cleanup:
|
|
if AssetTracker:
|
|
tracker = AssetTracker(storage_dir=args.output_dir)
|
|
if args.cleanup:
|
|
tracker.cleanup_missing_files()
|
|
if args.show_stats:
|
|
tracker.print_stats()
|
|
else:
|
|
print("Asset tracking is not available")
|
|
return
|
|
|
|
# Create the image downloader
|
|
downloader = ImageDownloader(
|
|
api_url=args.api_url,
|
|
list_endpoint=args.list_endpoint,
|
|
download_endpoint=args.download_endpoint,
|
|
output_dir=args.output_dir,
|
|
max_concurrent=args.max_concurrent,
|
|
timeout=args.timeout,
|
|
api_key=args.api_key,
|
|
email=args.email,
|
|
password=args.password,
|
|
track_assets=not args.no_tracking
|
|
)
|
|
|
|
try:
|
|
asyncio.run(downloader.download_all_assets(force_redownload=args.force_redownload))
|
|
except KeyboardInterrupt:
|
|
print("\nDownload interrupted by user")
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|