152 lines
4.3 KiB
Python
152 lines
4.3 KiB
Python
import asyncio
|
|
import aiohttp
|
|
import json
|
|
import pytz
|
|
import time
|
|
import httpx
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from tqdm import tqdm
|
|
|
|
CONCURRENCY = 50
|
|
RETRIES = 3
|
|
OUTPUT_DIR = Path("./downloads")
|
|
CHECKPOINT = Path("./temp/checkpoint.txt")
|
|
|
|
|
|
def load_checkpoint():
|
|
if CHECKPOINT.exists():
|
|
return set(CHECKPOINT.read_text().splitlines())
|
|
return set()
|
|
|
|
|
|
def save_to_checkpoint(path: Path):
|
|
with CHECKPOINT.open("a") as f:
|
|
f.write(str(path) + "\n")
|
|
|
|
|
|
# def utc_filename(timestamp: str, url: str) -> Path:
|
|
# dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
|
# ext = "mp4" if ".mp4" in url.lower() else "jpg"
|
|
# return OUTPUT_DIR / f"{dt.strftime('%Y-%m-%d_%H-%M-%S')}.{ext}"
|
|
|
|
def utc_filename(timestamp: str, url: str) -> Path:
|
|
# parse "YYYY-MM-DD HH:MM:SS UTC"
|
|
dt = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S UTC")
|
|
dt = dt.replace(tzinfo=pytz.utc)
|
|
ext = "mp4" if ".mp4" in url.lower() else "jpg"
|
|
out_path = os.path.join("./downloads",f"{dt.strftime('%Y-%m-%d_%H-%M-%S')}.{ext}")
|
|
print(out_path)
|
|
return Path(out_path)
|
|
|
|
|
|
async def get_cdn_url(download_link: str) -> str:
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
download_link,
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
)
|
|
response.raise_for_status()
|
|
return response.text.strip()
|
|
|
|
|
|
async def download_one(session, sem, url, path, failures, stats):
|
|
if path.exists():
|
|
return
|
|
|
|
# Already downloaded? (checkpoint resume)
|
|
if str(path) in stats["done"]:
|
|
return
|
|
|
|
async with sem:
|
|
for attempt in range(1, RETRIES + 1):
|
|
try:
|
|
cdn_url = await get_cdn_url(url)
|
|
async with session.get(cdn_url) as resp:
|
|
resp.raise_for_status()
|
|
data = await resp.read()
|
|
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_bytes(data)
|
|
|
|
# record MB
|
|
stats["mb"] += len(data) / (1024 * 1024)
|
|
|
|
# checkpoint
|
|
save_to_checkpoint(path)
|
|
stats["done"].add(str(path))
|
|
|
|
return
|
|
except Exception as e:
|
|
if attempt == RETRIES:
|
|
failures.append((url, str(e)))
|
|
await asyncio.sleep(0.3 * attempt) # small backoff
|
|
|
|
|
|
async def main():
|
|
# Load JSON
|
|
with open("./resources/json/memories_history.json", "r") as f:
|
|
memories = json.load(f)["Saved Media"]
|
|
|
|
# Prepare tasks
|
|
tasks_to_download = []
|
|
for item in memories:
|
|
url = item.get("Download Link")
|
|
ts = item.get("Date")
|
|
if not url or not ts:
|
|
continue
|
|
outpath = utc_filename(ts, url)
|
|
tasks_to_download.append((url, outpath))
|
|
|
|
stats = {
|
|
"mb": 0.0,
|
|
"done": load_checkpoint() # resume if exists
|
|
}
|
|
|
|
failures = []
|
|
sem = asyncio.Semaphore(CONCURRENCY)
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
start_time = time.time()
|
|
|
|
# tqdm based on number of files not yet done
|
|
to_process = [
|
|
(url, path) for (url, path) in tasks_to_download
|
|
if not path.exists() and str(path) not in stats["done"]
|
|
]
|
|
|
|
progress = tqdm(total=len(to_process), desc="Downloading", unit="file")
|
|
|
|
# Wrap each task so progress bar updates *after* the download completes
|
|
async def wrapped(url, path):
|
|
await download_one(session, sem, url, path, failures, stats)
|
|
progress.update(1)
|
|
|
|
# Start tasks in batches
|
|
await asyncio.gather(*(wrapped(url, path) for url, path in to_process))
|
|
progress.close()
|
|
|
|
elapsed = time.time() - start_time
|
|
|
|
# Final summary
|
|
mb_total = stats["mb"]
|
|
speed = mb_total / elapsed if elapsed > 0 else 0
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"Downloaded: {len(to_process) - len(failures)} files")
|
|
print(f"Failed: {len(failures)} files")
|
|
print(f"Data: {mb_total:.2f} MB")
|
|
print(f"Speed: {speed:.2f} MB/s")
|
|
print("=" * 60)
|
|
|
|
if failures:
|
|
print("\n⚠️ Failed downloads:")
|
|
for url, err in failures:
|
|
print(f" - {url} ({err})")
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|