tile_generator.py

import io
import logging
import os
import shutil
import time
import xml.etree.ElementTree as ET

import boto3
import requests
from mapproxy.config.loader import load_configuration
from mapproxy.request.base import Request
from mapproxy.service.wmts import WMTSRestServer
from mapproxy.util.ext.wmsparse.util import parse_datetime_range

from wmts_tile_generator.config import CONFIG, ServerMode

BASEPATH = os.path.dirname(__file__)

_logger = logging.getLogger(__name__)

if CONFIG.SERVER_MODE == ServerMode.S3:
    _s3 = boto3.client("s3")


def replace_template_vars(
    time_values: str, default_time: str, bbox_values: str
) -> None:
    with open(
        os.path.join(BASEPATH, "mapproxy_template.yaml"), encoding="utf-8"
    ) as file:
        config_template = file.read()
    config_template = config_template.replace(
        "{{WMS_AUTH_TOKEN}}", CONFIG.WMS_AUTH_TOKEN
    )
    config_template = config_template.replace("{{WMS_URL}}", CONFIG.WMS_URL)
    config_template = config_template.replace("{{LAYER_NAME}}", CONFIG.LAYER_NAME)
    config_template = config_template.replace("{{LAYER_STYLE}}", CONFIG.LAYER_STYLE)
    config_template = config_template.replace(
        "{{PROJECTION_EPSG}}", CONFIG.PROJECTION_EPSG
    )
    config_template = config_template.replace("{{TIME_VALUES}}", time_values)
    config_template = config_template.replace("{{DEFAULT_TIME}}", default_time)
    config_template = config_template.replace("{{PROJECTION_BBOX}}", bbox_values)
    output_filename = os.path.join(BASEPATH, "mapproxy.yaml")
    if os.path.exists(output_filename):
        os.remove(output_filename)
    with open(output_filename, "w", encoding="utf-8") as file:
        file.write(config_template)


def get_time_values(capabilities_text: str) -> tuple[str, str]:
    found_layer = None
    root = ET.fromstring(capabilities_text)
    capabilities = [child for child in root if "Capability" in child.tag]
    for capability in capabilities:
        layers = [child for child in capability if "Layer" in child.tag]
        for layer in layers:
            actual_layers = [child for child in layer if "Layer" in child.tag]
            for actual_layer in actual_layers:
                names = [child for child in actual_layer if "Name" in child.tag]
                for name in names:
                    if name.text == CONFIG.LAYER_NAME:
                        found_layer = actual_layer

    default_time = None
    time_values = None
    if found_layer is not None:
        dimensions = [child for child in found_layer if "Dimension" in child.tag]
        for dimension in dimensions:
            if dimension.attrib["name"] == "time":
                default_time = dimension.attrib["default"]
                time_values = dimension.text
    if default_time is not None and time_values is not None:
        return time_values, default_time
    raise LookupError("No time dimension found in WMS capabilities")


def get_bbox_values(capabilities_text: str) -> str:
    found_bbox = None
    root = ET.fromstring(capabilities_text)
    capabilities = [child for child in root if "Capability" in child.tag]
    for capability in capabilities:
        layers = [child for child in capability if "Layer" in child.tag]
        for layer in layers:
            bboxes = [child for child in layer if "BoundingBox" in child.tag]
            for bbox in bboxes:
                if bbox.attrib["CRS"] == CONFIG.PROJECTION_EPSG:
                    found_bbox = [
                        bbox.attrib["minx"],
                        bbox.attrib["miny"],
                        bbox.attrib["maxx"],
                        bbox.attrib["maxy"],
                    ]

    if found_bbox is not None:
        return f"[{','.join(found_bbox)}]"
    raise LookupError("No bbox found in WMS capabilities")


def create_wmts_capabilities_xml(service: WMTSRestServer) -> None:
    request_path = "wmts/1.0.0/WMTSCapabilities.xml"
    output_filename = os.path.join(BASEPATH, request_path)
    os.makedirs(os.path.dirname(output_filename), exist_ok=True)
    if os.path.exists(output_filename):
        os.remove(output_filename)
    with open(output_filename, "w", encoding="utf-8") as file:
        file.write(
            service.handle(
                Request(
                    {
                        "PATH_INFO": f"/{request_path}",
                        "wsgi.url_scheme": CONFIG.WEBSITE_URL_SCHEME,
                        "HTTP_HOST": CONFIG.WEBSITE_HTTP_HOST,
                    }
                )
            ).response
        )
    if CONFIG.SERVER_MODE == ServerMode.S3:
        _s3.upload_file(output_filename, CONFIG.WEBSITE_BUCKET_NAME, request_path)


def file_exists_on_s3(request_path: str) -> bool:
    try:
        _s3.head_object(Bucket=CONFIG.WEBSITE_BUCKET_NAME, Key=request_path)
    except Exception:  # pylint: disable=broad-except
        return False
    return True


def cache_image(
    service: WMTSRestServer, timestamp: str, zoom_level: int, row: int, col: int
) -> None:
    request_path = f"wmts/{CONFIG.LAYER_NAME}/{timestamp}/{CONFIG.PROJECTION_EPSG}/{zoom_level:0>2}/{row}/{col}.png"
    output_filename = os.path.join(BASEPATH, request_path)
    if CONFIG.SERVER_MODE == ServerMode.S3:
        if file_exists_on_s3(request_path):
            return
    else:
        if os.path.exists(output_filename):
            return
    os.makedirs(os.path.dirname(output_filename), exist_ok=True)
    _logger.info("caching: %s", request_path)
    response = service.handle(
        Request(
            {
                "PATH_INFO": f"/{request_path}",
                "wsgi.url_scheme": CONFIG.WEBSITE_URL_SCHEME,
                "HTTP_HOST": CONFIG.WEBSITE_HTTP_HOST,
            }
        )
    ).response
    with open(output_filename, "wb") as file:
        if isinstance(response, io.BufferedReader):
            file.write(response.read())
        else:
            file.write(response.getbuffer())
    if CONFIG.SERVER_MODE == ServerMode.S3:
        _s3.upload_file(output_filename, CONFIG.WEBSITE_BUCKET_NAME, request_path)
    _logger.info("cached: %s", request_path)


def generate() -> None:
    capabilities = requests.get(
        f"{CONFIG.WMS_URL}service=WMS&request=GetCapabilities",
        headers={"Authorization": CONFIG.WMS_AUTH_TOKEN},
        timeout=60,
    ).text
    time_values, default_time = get_time_values(capabilities)
    bbox_values = get_bbox_values(capabilities)
    replace_template_vars(time_values, default_time, bbox_values)
    configuration = load_configuration(os.path.join(BASEPATH, "mapproxy.yaml"))
    services = configuration.configured_services()
    timestamps_to_cache = sorted(parse_datetime_range(time_values))[-3:]
    for timestamp in timestamps_to_cache:
        for zoom_level in [0, 1, 2]:
            for row in range(2**zoom_level):
                for col in range(2**zoom_level):
                    cache_image(services[0], timestamp, zoom_level, row, col)
                    time.sleep(0.2)

    create_wmts_capabilities_xml(services[0])

    shutil.rmtree(os.path.join(BASEPATH, "cache_data"), ignore_errors=True)
    if CONFIG.SERVER_MODE == ServerMode.S3:
        shutil.rmtree(os.path.join(BASEPATH, "wmts"), ignore_errors=True)

    _logger.info("sleeping...")
    time.sleep(10)