proxy-god/proxy_chain_manager/chain_map.py

"""World-map visualization for proxy chains (Pillow + equirectangular projection)."""
from __future__ import annotations

import ipaddress
import logging
import math
import socket
from dataclasses import dataclass
from typing import Literal
from urllib.parse import urlparse

from PIL import Image, ImageDraw, ImageFont

try:
    _RESAMPLE = Image.Resampling.LANCZOS
except AttributeError:
    _RESAMPLE = Image.LANCZOS  # type: ignore[attr-defined]

from .config import normalize_proxy_url, redact_proxy_url
from .exit_intel import ExitIntel, fetch_ip_geo

log = logging.getLogger(__name__)

ChainStatus = Literal["healthy", "connecting", "dead", "idle"]

# Theme colors (match app.py palette)
_BG = "#070713"
_GRID = "#141432"
_LAND = "#12182e"
_LAND_EDGE = "#1e2a4a"
_LINE_HEALTHY = "#00e5ff"
_LINE_CONNECTING = "#ffd000"
_LINE_DEAD = "#ff2a6d"
_LINE_IDLE = "#3052ff"
_YOU = "#00ff9c"
_HOP = "#00e5ff"
_EXIT = "#ffd000"
_LABEL = "#eaf2ff"
_LABEL_DIM = "#7a8aab"
_UNKNOWN = "#2a2f4a"


@dataclass(frozen=True)
class MapPoint:
    role: Literal["you", "hop", "exit"]
    label: str
    lat: float
    lon: float
    detail: str = ""
    hop_index: int = -1


def extract_hop_host(proxy_url: str) -> str | None:
    """Return hostname or IP from a proxy URL."""
    p = urlparse(normalize_proxy_url(proxy_url))
    host = (p.hostname or "").strip()
    return host or None


def resolve_host_ip(host: str, timeout_seconds: float = 4.0) -> str | None:
    """Resolve host to IPv4/IPv6 string. Never raises."""
    h = (host or "").strip()
    if not h:
        return None
    try:
        ipaddress.ip_address(h)
        return h
    except ValueError:
        pass
    try:
        prev = socket.getdefaulttimeout()
        socket.setdefaulttimeout(max(1.0, float(timeout_seconds)))
        try:
            infos = socket.getaddrinfo(h, None, type=socket.SOCK_STREAM)
        finally:
            socket.setdefaulttimeout(prev)
        for info in infos:
            addr = info[4][0]
            try:
                ipaddress.ip_address(addr)
                return addr
            except ValueError:
                continue
    except Exception as e:
        log.debug("resolve_host_ip(%s): %s", h, e)
    return None


def _loc_label(intel: ExitIntel) -> str:
    bits = [b for b in (intel.city, intel.region, intel.country) if b]
    return ", ".join(bits) if bits else intel.country_code or intel.ip or "—"


def _intel_to_point(
    role: Literal["you", "hop", "exit"],
    label: str,
    intel: ExitIntel,
    *,
    hop_index: int = -1,
) -> MapPoint | None:
    if not intel.ok or intel.lat is None or intel.lon is None:
        return None
    return MapPoint(
        role=role,
        label=label,
        lat=float(intel.lat),
        lon=float(intel.lon),
        detail=_loc_label(intel),
        hop_index=hop_index,
    )


def build_chain_map_points(
    hops: list[str],
    *,
    direct_ip: str = "",
    exit_ip: str = "",
    exit_intel: ExitIntel | None = None,
    timeout_seconds: float = 8.0,
) -> list[MapPoint]:
    """Resolve geo for YOU → hop₁ → … → EXIT. Never raises."""
    points: list[MapPoint] = []

    origin_ip = (direct_ip or "").strip()
    if origin_ip and origin_ip != "—":
        you = _intel_to_point("you", "YOU", fetch_ip_geo(origin_ip, timeout_seconds))
        if you:
            points.append(you)

    for i, hop in enumerate(hops):
        host = extract_hop_host(hop)
        if not host:
            continue
        ip = resolve_host_ip(host, timeout_seconds=min(4.0, timeout_seconds))
        if not ip:
            continue
        short = redact_proxy_url(hop)
        if len(short) > 22:
            short = short[:20] + "…"
        intel = fetch_ip_geo(ip, timeout_seconds)
        pt = _intel_to_point("hop", f"H{i + 1}", intel, hop_index=i)
        if pt:
            points.append(pt)

    exit = (exit_ip or "").strip()
    if exit and exit != "—":
        if (
            exit_intel
            and exit_intel.ok
            and exit_intel.ip == exit
            and exit_intel.lat is not None
            and exit_intel.lon is not None
        ):
            pt = _intel_to_point("exit", "EXIT", exit_intel)
        else:
            pt = _intel_to_point("exit", "EXIT", fetch_ip_geo(exit, timeout_seconds))
        if pt:
            points.append(pt)

    return points


def _project(lat: float, lon: float, width: int, height: int, margin: int) -> tuple[float, float]:
    inner_w = max(1, width - margin * 2)
    inner_h = max(1, height - margin * 2)
    x = margin + (float(lon) + 180.0) / 360.0 * inner_w
    y = margin + (90.0 - float(lat)) / 180.0 * inner_h
    return x, y


def _great_circle_points(
    lat1: float,
    lon1: float,
    lat2: float,
    lon2: float,
    steps: int = 24,
) -> list[tuple[float, float]]:
    """Interpolate along a great circle for curved hop arcs."""
    phi1, lam1 = math.radians(lat1), math.radians(lon1)
    phi2, lam2 = math.radians(lat2), math.radians(lon2)
    d = 2 * math.asin(
        math.sqrt(
            math.sin((phi2 - phi1) / 2) ** 2
            + math.cos(phi1) * math.cos(phi2) * math.sin((lam2 - lam1) / 2) ** 2
        )
    )
    if d < 1e-9:
        return [(lat1, lon1), (lat2, lon2)]
    out: list[tuple[float, float]] = []
    for i in range(steps + 1):
        f = i / steps
        a = math.sin((1 - f) * d) / math.sin(d)
        b = math.sin(f * d) / math.sin(d)
        x = a * math.cos(phi1) * math.cos(lam1) + b * math.cos(phi2) * math.cos(lam2)
        y = a * math.cos(phi1) * math.sin(lam1) + b * math.cos(phi2) * math.sin(lam2)
        z = a * math.sin(phi1) + b * math.sin(phi2)
        out.append((math.degrees(math.atan2(z, math.sqrt(x * x + y * y))), math.degrees(math.atan2(y, x))))
    return out


def _line_color(status: ChainStatus) -> str:
    return {
        "healthy": _LINE_HEALTHY,
        "connecting": _LINE_CONNECTING,
        "dead": _LINE_DEAD,
        "idle": _LINE_IDLE,
    }.get(status, _LINE_IDLE)


def _node_color(role: str) -> str:
    return {"you": _YOU, "hop": _HOP, "exit": _EXIT}.get(role, _HOP)


def _load_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    for name in ("segoeui.ttf", "arial.ttf", "DejaVuSans.ttf"):
        try:
            return ImageFont.truetype(name, size)
        except OSError:
            continue
    return ImageFont.load_default()


def _draw_landmasses(draw: ImageDraw.ImageDraw, width: int, height: int, margin: int) -> None:
    """Minimal continent silhouettes — stylized, not survey-accurate."""
    blobs: list[list[tuple[float, float]]] = [
        # North America
        [(-168, 72), (-140, 75), (-100, 72), (-80, 50), (-75, 25), (-95, 15), (-110, 20), (-125, 48), (-168, 72)],
        # South America
        [(-82, 12), (-72, -5), (-55, -35), (-65, -55), (-75, -18), (-82, 12)],
        # Europe / Africa
        [(-10, 72), (30, 70), (40, 55), (35, 30), (20, 5), (-5, 5), (-18, 28), (-10, 72)],
        [(15, 35), (35, 32), (50, 12), (42, -5), (18, -35), (15, 35)],
        # Asia
        [(40, 72), (100, 75), (140, 55), (130, 35), (110, 10), (80, 8), (60, 25), (40, 45), (40, 72)],
        # Australia
        [(115, -12), (135, -12), (150, -25), (145, -38), (115, -38), (115, -12)],
    ]
    for poly in blobs:
        pts = [_project(lat, lon, width, height, margin) for lat, lon in poly]
        draw.polygon(pts, fill=_LAND, outline=_LAND_EDGE)


def render_chain_map(
    points: list[MapPoint],
    *,
    width: int = 960,
    height: int = 220,
    status: ChainStatus = "idle",
    margin: int = 18,
) -> Image.Image:
    """Render chain hops on a dark equirectangular world map."""
    img = Image.new("RGB", (max(320, width), max(120, height)), _BG)
    draw = ImageDraw.Draw(img)
    _draw_landmasses(draw, width, height, margin)

    # Lat/lon grid
    for lat in range(-60, 91, 30):
        y = _project(lat, 0, width, height, margin)[1]
        draw.line([(margin, y), (width - margin, y)], fill=_GRID, width=1)
    for lon in range(-150, 181, 30):
        x = _project(0, lon, width, height, margin)[0]
        draw.line([(x, margin), (x, height - margin)], fill=_GRID, width=1)

    if len(points) < 2:
        font = _load_font(13)
        msg = "Start the chain to plot hops on the map." if not points else "Need 2+ geo points to draw path."
        draw.text((margin, height // 2 - 8), msg, fill=_LABEL_DIM, font=font)
        return img

    line_c = _line_color(status)
    projected = [_project(p.lat, p.lon, width, height, margin) for p in points]

    # Glow underlay
    for i in range(len(projected) - 1):
        arc = _great_circle_points(points[i].lat, points[i].lon, points[i + 1].lat, points[i + 1].lon)
        arc_xy = [_project(lat, lon, width, height, margin) for lat, lon in arc]
        draw.line(arc_xy, fill=line_c, width=5)
    for i in range(len(projected) - 1):
        arc = _great_circle_points(points[i].lat, points[i].lon, points[i + 1].lat, points[i + 1].lon)
        arc_xy = [_project(lat, lon, width, height, margin) for lat, lon in arc]
        draw.line(arc_xy, fill=line_c, width=2)

    label_font = _load_font(11)
    detail_font = _load_font(9)
    for pt, (x, y) in zip(points, projected):
        color = _node_color(pt.role)
        r = 7 if pt.role == "exit" else (6 if pt.role == "you" else 5)
        draw.ellipse((x - r - 2, y - r - 2, x + r + 2, y + r + 2), fill=color)
        draw.ellipse((x - r, y - r, x + r, y + r), fill=_BG, outline=color, width=2)
        tx, ty = x + 10, y - 14
        draw.text((tx + 1, ty + 1), pt.label, fill="#000000", font=label_font)
        draw.text((tx, ty), pt.label, fill=_LABEL, font=label_font)
        if pt.detail:
            draw.text((tx, ty + 13), pt.detail[:28], fill=_LABEL_DIM, font=detail_font)

    return img


def resize_map_display(pil_img: Image.Image, width: int, height: int) -> Image.Image:
    """Resize rendered map for the GUI widget."""
    w = max(320, int(width))
    h = max(120, int(height))
    return pil_img.resize((w, h), _RESAMPLE)