Automating USGS EarthExplorer Bulk Downloads with requests

import os
import json
import time
import hashlib
import logging
from pathlib import Path
from urllib.parse import urljoin
from requests import Session
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

USGS_M2M_BASE = "https://m2m.cr.usgs.gov/api/api/json/v1.5/"
DATASET_NAME  = "landsat_ot_c2_l1"   # Landsat 8/9 Collection 2 Level-1
OUTPUT_DIR    = Path("./usgs_downloads")
OUTPUT_DIR.mkdir(exist_ok=True)


# ── Phase 1: Session + Token ──────────────────────────────────────────────────

def get_session() -> Session:
    """Return a requests.Session with connection pooling and exponential backoff.

    backoff_factor=1.5 produces delays of ~1.5 s, 3 s, 6 s, 12 s, 18 s before
    each successive retry — enough to ride out USGS 429 bursts without hammering
    the endpoint. pool_connections/pool_maxsize=10 allows concurrent workers to
    share an underlying urllib3 connection pool instead of opening a fresh TLS
    handshake per request.
    """
    session = Session()
    retry_strategy = Retry(
        total=5,
        backoff_factor=1.5,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET", "POST"],   # urllib3 >= 2.0; use method_whitelist on older
    )
    adapter = HTTPAdapter(
        max_retries=retry_strategy,
        pool_connections=10,
        pool_maxsize=10,
    )
    session.mount("https://", adapter)
    session.headers.update({"Content-Type": "application/json"})
    return session


def authenticate(session: Session, username: str, password: str) -> str:
    """Exchange USGS credentials for a time-limited X-Auth-Token.

    Tokens expire after ~120 minutes. For pipelines that run longer, catch
    HTTP 401 inside the download loop and call this function again to refresh.
    Credentials must come from environment variables or a secrets manager —
    never from source code or CLI arguments that appear in process lists.
    """
    payload = {"username": username, "password": password}
    resp = session.post(urljoin(USGS_M2M_BASE, "login"), json=payload, timeout=30)
    resp.raise_for_status()
    token = resp.json()["data"]
    logger.info("M2M authentication successful. Token valid for ~2 hours.")
    return token


# ── Phase 2: Scene Metadata Resolution ────────────────────────────────────────

def search_scenes(
    session: Session,
    token: str,
    bbox: list[float],
    start_date: str,
    end_date: str,
    max_results: int = 50,
) -> list[dict]:
    """Return scene metadata records matching a bounding box and date range.

    bbox must be [min_lon, min_lat, max_lon, max_lat] in WGS-84 decimal degrees.
    The M2M spatialFilter uses lowerLeft/upperRight dicts — not GeoJSON order.
    """
    lower_left  = {"latitude": bbox[1], "longitude": bbox[0]}
    upper_right = {"latitude": bbox[3], "longitude": bbox[2]}

    payload = {
        "datasetName": DATASET_NAME,
        "maxResults":  max_results,
        "sceneFilter": {
            "acquisitionFilter": {"start": start_date, "end": end_date},
            "spatialFilter": {
                "filterType":  "mbr",
                "lowerLeft":   lower_left,
                "upperRight":  upper_right,
            },
        },
    }
    headers = {"X-Auth-Token": token}
    resp = session.post(
        urljoin(USGS_M2M_BASE, "scene-search"),
        json=payload,
        headers=headers,
        timeout=60,
    )
    resp.raise_for_status()
    results = resp.json()["data"]["results"]
    logger.info("scene-search returned %d scenes.", len(results))
    return results


def resolve_product_id(
    session: Session,
    token: str,
    entity_id: str,
    preferred_format: str = "Bundle",
) -> str | None:
    """Enumerate available download products and return the first matching productId.

    Never hardcode productIds — they change across API versions and dataset
    releases. This helper calls download-options for a single representative
    entityId and selects the preferred bundle type.
    """
    headers = {"X-Auth-Token": token}
    payload  = {"datasetName": DATASET_NAME, "entityIds": [entity_id]}
    resp = session.post(
        urljoin(USGS_M2M_BASE, "download-options"),
        json=payload,
        headers=headers,
        timeout=30,
    )
    resp.raise_for_status()
    for option in resp.json().get("data", []):
        if preferred_format.lower() in option.get("productName", "").lower():
            return option["id"]
    return None


def get_download_urls(
    session: Session,
    token: str,
    entity_ids: list[str],
    product_id: str,
) -> dict[str, str]:
    """Resolve temporary signed download URLs via the two-step M2M protocol.

    Step 1 (download-request) enqueues the archives for preparation on USGS
    servers. Step 2 (download-retrieve) polls for readiness and returns signed
    S3 URLs. URLs are short-lived (~15 minutes); start streaming immediately.
    """
    headers = {"X-Auth-Token": token}

    # Step 1: enqueue archive preparation
    products_payload = {
        "downloads": [
            {"datasetName": DATASET_NAME, "entityId": eid, "productId": product_id}
            for eid in entity_ids
        ],
        "label": "etl-bulk-download",
    }
    resp = session.post(
        urljoin(USGS_M2M_BASE, "download-request"),
        json=products_payload,
        headers=headers,
        timeout=60,
    )
    resp.raise_for_status()

    # Step 2: retrieve signed URLs (may need a short wait for large batches)
    retrieve_payload = {"label": "etl-bulk-download"}
    for attempt in range(5):
        resp = session.post(
            urljoin(USGS_M2M_BASE, "download-retrieve"),
            json=retrieve_payload,
            headers=headers,
            timeout=60,
        )
        resp.raise_for_status()
        available = resp.json()["data"].get("available", [])
        if available:
            break
        logger.info("download-retrieve: no URLs ready yet, waiting 10 s (attempt %d/5).", attempt + 1)
        time.sleep(10)

    url_map: dict[str, str] = {}
    for item in available:
        url_map[item["entityId"]] = item["url"]
    logger.info("Resolved %d download URLs.", len(url_map))
    return url_map


# ── Phase 3: Chunked File Streaming ───────────────────────────────────────────

def _md5_file(path: Path) -> str:
    h = hashlib.md5()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def download_file(
    session: Session,
    url: str,
    dest_path: Path,
    chunk_size: int = 8192,
    expected_md5: str | None = None,
) -> None:
    """Stream a large archive to disk without loading it into memory.

    The Content-Length header is used only for progress logging; its absence
    does not prevent the download. After writing, optionally verifies MD5
    against a checksum supplied by the M2M API response metadata.
    """
    with session.get(url, stream=True, timeout=(10, 300)) as resp:
        resp.raise_for_status()
        total_bytes = int(resp.headers.get("content-length", 0))
        downloaded  = 0

        with open(dest_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=chunk_size):
                if chunk:   # filter out keep-alive empty chunks
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total_bytes:
                        pct = downloaded / total_bytes * 100
                        logger.info(
                            "%s: %.1f MB / %.1f MB (%.0f%%)",
                            dest_path.name,
                            downloaded / 1_048_576,
                            total_bytes / 1_048_576,
                            pct,
                        )

    if expected_md5 and _md5_file(dest_path) != expected_md5:
        dest_path.unlink()
        raise ValueError(f"MD5 mismatch for {dest_path.name} — file removed, retry required.")

    logger.info("Download complete: %s", dest_path)


# ── Orchestrating run ──────────────────────────────────────────────────────────

def run_pipeline(
    username: str,
    password: str,
    bbox: list[float],
    start_date: str,
    end_date: str,
) -> None:
    """Full three-phase M2M download pipeline with idempotent file checks."""
    session = get_session()
    token   = authenticate(session, username, password)

    scenes = search_scenes(session, token, bbox, start_date, end_date)
    if not scenes:
        logger.warning("No scenes found for the specified criteria.")
        return

    entity_ids = [s["entityId"] for s in scenes]

    # Discover product ID dynamically from a representative entity
    product_id = resolve_product_id(session, token, entity_ids[0])
    if not product_id:
        raise RuntimeError("No downloadable bundle found for dataset '%s'.", DATASET_NAME)

    url_map = get_download_urls(session, token, entity_ids, product_id)

    for entity_id, dl_url in url_map.items():
        dest = OUTPUT_DIR / f"{entity_id}.tar"
        if dest.exists():
            logger.info("Skipping already downloaded file: %s", dest.name)
            continue
        try:
            download_file(session, dl_url, dest)
        except Exception as exc:
            logger.error("Failed to download %s: %s", entity_id, exc)


if __name__ == "__main__":
    run_pipeline(
        username=os.environ["USGS_USERNAME"],
        password=os.environ["USGS_PASSWORD"],
        bbox=[-122.5, 37.7, -122.3, 37.9],   # San Francisco Bay Area
        start_date="2023-06-01",
        end_date="2023-06-30",
    )