wmts-exfiltrator/pkg/thread.py

47 lines
1.6 KiB
Python

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'{output_path} is not a PNG, deleting and retrying...')
else:
return row, col, 'exist'
tile_url = f"{base_url}/{row}/{col}"
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)
# Don't check the PNG after downloading, only when loading.
# if not is_png(output_path):
# del_path(output_path)
# return row, col, 'failure'
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'