#!/usr/bin/env python3
"""Local auth_request backend: allow overseas + Shanghai; block other mainland CN."""

import ipaddress
import os
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Optional
from urllib.parse import parse_qs, urlparse

import geoip2.database
from geoip2.errors import AddressNotFoundError

DB_PATH = os.environ.get("GEOIP_CITY_DB", "/usr/share/GeoIP/GeoLite2-City.mmdb")
LISTEN = os.environ.get("WYRS_GEO_LISTEN", "127.0.0.1")
PORT = int(os.environ.get("WYRS_GEO_PORT", "18099"))

_reader = None  # type: Optional[geoip2.database.Reader]


def get_reader():
    global _reader
    if _reader is None:
        _reader = geoip2.database.Reader(DB_PATH)
    return _reader


def client_ip(handler):
    # Prefer query ?ip= for local probes; production uses X-Real-IP from nginx.
    qs = parse_qs(urlparse(handler.path).query)
    if "ip" in qs and qs["ip"]:
        return qs["ip"][0].strip()
    xri = handler.headers.get("X-Real-IP", "").strip()
    if xri:
        return xri.split(",")[0].strip()
    return handler.client_address[0]


def is_trusted_local(ip):
    try:
        addr = ipaddress.ip_address(ip)
    except ValueError:
        return True
    return bool(
        addr.is_private
        or addr.is_loopback
        or addr.is_link_local
        or addr.is_reserved
    )


def is_shanghai_mainland(ip):
    """Return True if Shanghai, False if other mainland CN, None if not mainland / unknown."""
    try:
        rec = get_reader().city(ip)
    except AddressNotFoundError:
        return None
    except Exception as exc:  # keep site up if DB quirks
        sys.stderr.write("geo lookup error ip=%s: %s\n" % (ip, exc))
        return None

    cc = (rec.country.iso_code or "").upper()
    if cc != "CN":
        return None

    for sub in rec.subdivisions:
        iso = (sub.iso_code or "").upper()
        name = (sub.name or "").strip().lower()
        if iso == "SH" or name == "shanghai":
            return True

    city = (rec.city.name or "").strip().lower()
    if city == "shanghai":
        return True
    return False


def allowed(ip):
    if is_trusted_local(ip):
        return True
    decision = is_shanghai_mainland(ip)
    # None = overseas / HK-MO-TW / unknown → allow
    # True = Shanghai → allow
    # False = other mainland → deny
    return decision is not False


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt: str, *args) -> None:  # noqa: A003
        sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args))

    def do_GET(self) -> None:  # noqa: N802
        path = self.path.split("?", 1)[0]
        if path not in ("/check", "/healthz"):
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        if path == "/healthz":
            body = b"ok\n"
            self.send_response(200)
            self.send_header("Content-Type", "text/plain; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        ip = client_ip(self)
        if allowed(ip):
            self.send_response(204)
            self.send_header("Content-Length", "0")
            self.send_header("X-Wyrs-Geo", "allow")
            self.end_headers()
            return

        self.send_response(403)
        self.send_header("Content-Length", "0")
        self.send_header("X-Wyrs-Geo", "deny-cn")
        self.end_headers()


def main() -> None:
    if not os.path.isfile(DB_PATH):
        print(f"missing GeoIP DB: {DB_PATH}", file=sys.stderr)
        sys.exit(1)
    get_reader()
    httpd = ThreadingHTTPServer((LISTEN, PORT), Handler)
    print(f"wyrs-geo-guard listening on {LISTEN}:{PORT} db={DB_PATH}", flush=True)
    httpd.serve_forever()


if __name__ == "__main__":
    main()