src/geo.py

"""Coarse, server-side geolocation of a visitor's IP, for analytics.

Nothing here touches the browser: no navigator.geolocation, no permission
prompt, no consent dialog to dismiss. The caller hands over the IP it
already has from the request and gets back city-level detail or nothing.

`lookup` is deliberately total -- it returns {} for anything it cannot
answer rather than raising, because a failed analytics lookup must never
break an image generation.
"""

import ipaddress
import json
import os
import ssl
import urllib.error
import urllib.request

# ipwho.is: free, no API key, HTTPS, ~10k lookups/month, and it takes the
# target IP in the path so no query string or signup is involved. Picked
# over ipapi.co (1k/day and already returning 429 unauthenticated), ip-api
# .com (HTTPS is paid-only) and freeipapi.com (redirects away from its
# documented endpoint). Failures are signalled in-band as {"success": false}
# with HTTP 200, so the status code alone is not enough to trust a response.
API_URL = os.environ.get("GEO_API_URL", "https://ipwho.is").rstrip("/")

# Runs inside a user-facing request, so the budget is "cheap enough to be
# invisible". Two seconds is already generous for a single GET.
TIMEOUT = float(os.environ.get("GEO_TIMEOUT", "2"))

ENABLED = os.environ.get("GEO_LOOKUP", "on").strip().lower() != "off"

# Visitors arrive in bursts from the same address, and geolocation of an IP
# does not meaningfully change during one container's lifetime. Only hits
# are cached: a timeout or a rate-limit is transient and caching {} would
# blind the Space for as long as the process lives.
_cache: dict[str, dict] = {}
_CACHE_MAX = 4096

_HEADERS = {"User-Agent": "diffusiondb-sd15-lora-space/1.0"}

# The Space image trusts the system CA store, but a python.org interpreter on
# macOS ships none and every HTTPS call there fails verification. certifi is
# already installed as a transitive dependency of huggingface_hub, so prefer
# its bundle when it is importable; never fall back to an unverified context.
try:
    import certifi

    _SSL = ssl.create_default_context(cafile=certifi.where())
except Exception:
    _SSL = ssl.create_default_context()


def _is_public(ip: str) -> bool:
    """True only for addresses worth spending a network call on."""
    try:
        address = ipaddress.ip_address(ip)
    except ValueError:  # not an IP at all, or carries a %scope suffix
        return False
    return not (
        address.is_private
        or address.is_loopback
        or address.is_link_local
        or address.is_reserved
        or address.is_multicast
        or address.is_unspecified
    )


def _fetch(ip: str) -> dict:
    request = urllib.request.Request(f"{API_URL}/{ip}", headers=_HEADERS)
    with urllib.request.urlopen(
        request, timeout=TIMEOUT, context=_SSL
    ) as response:
        # Cap the read: a wrong GEO_API_URL could point at anything.
        return json.loads(response.read(64_000).decode("utf-8"))


def _parse(payload: dict) -> dict:
    if not isinstance(payload, dict) or payload.get("success") is False:
        return {}

    # Nested under "timezone" as an object, unlike every other field.
    timezone = payload.get("timezone")
    if isinstance(timezone, dict):
        timezone = timezone.get("id")

    result = {
        "country": payload.get("country"),
        "country_code": payload.get("country_code"),
        "region": payload.get("region"),
        "city": payload.get("city"),
        "latitude": float(payload["latitude"]),
        "longitude": float(payload["longitude"]),
        "timezone": timezone,
    }
    # An answer with no coordinates or no country is not an answer; better
    # to report "unknown" than to record a row of Nones.
    if result["country"] is None:
        return {}
    return result


def lookup(ip: str) -> dict:
    """Return {} on any failure, or a dict with these exact keys:
       'country'    e.g. 'United States'
       'country_code' e.g. 'US'
       'region'     the state/province, e.g. 'California'
       'city'       e.g. 'San Francisco'
       'latitude'   float
       'longitude'  float
       'timezone'   e.g. 'America/Los_Angeles'
    """
    # isinstance, not just falsiness: "never raises" has to hold even when a
    # caller passes whatever its framework handed it rather than a str.
    if not ENABLED or not ip or not isinstance(ip, str):
        return {}

    ip = ip.strip()
    if not _is_public(ip):
        return {}

    if ip in _cache:
        return _cache[ip]

    try:
        result = _parse(_fetch(ip))
    except Exception:
        # Timeout, DNS failure, HTTP error, truncated body, non-JSON,
        # missing or non-numeric coordinates -- all of it means the same
        # thing to the caller, and none of it is worth a traceback.
        return {}

    if result:
        if len(_cache) >= _CACHE_MAX:
            _cache.clear()  # crude, but bounded and never wrong
        _cache[ip] = result
    return result


if __name__ == "__main__":
    for probe in ("8.8.8.8", "1.1.1.1", "2606:4700:4700::1111",
                  "192.168.1.1", "127.0.0.1", "", "not-an-ip"):
        print(f"{probe or '(empty)':>22}  {lookup(probe)}")