wmts-exfiltrator/pkg/thread.py

import shutil
from pathlib import Path

import requests
from tqdm import tqdm

from pkg.proxies import PROXIES
from .image import is_png


def del_path(p: Path):
    if p.is_file() or p.is_symlink():
        p.unlink(missing_ok=True)
    else:
        shutil.rmtree(p)


def download_tile(task):
    row, col, base_url, r_headers, output, use_proxy = task
    try:
        output_path: Path = output / f"{row}_{col}.png"
        if output_path.exists():
            if not is_png(output_path):
                # Delete the file and try again.
                del_path(output_path)
                tqdm.write(f'cannot identify image file: "{output_path}", deleting and retrying...')
            else:
                return row, col, 'exist'
        tile_url = f"{base_url}/{row}/{col}".replace('//', '/').replace(':/', '://')
        response = requests.get(tile_url, headers=r_headers, proxies=PROXIES if use_proxy else None, timeout=60)
        if response.status_code == 200:
            if not response.headers.get('Content-Type') == 'image/png':
                raise Exception(f'Response gave Content-Type: {response.headers.get("Content-Type")}')
            with open(output_path, "wb") as f:
                f.write(response.content)
            return row, col, 'success'
        else:
            print(f"Failed to download tile {row}_{col}")
    except Exception as e:
        # traceback.print_exc()
        tqdm.write(f'Exception on {(row, col)} - {e.__class__.__name__}: {e}')
        return row, col, 'failure'