osint-dashboard/app/ais_stream.py

"""AISStream WebSocket worker — server-side only.

aisstream.io forbids browser clients. Connect from the FastAPI/ingest
process, upsert last-known positions, and expose them via GET /api/vessels.

Idle (no crash) when AISSTREAM_API_KEY is unset. Reconnect with jittered
backoff and resend the full subscription within 3 seconds of each connect.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import random
import time

from config import AISSTREAM_API_KEY, AISSTREAM_BBOX
from keystore import get_api_key
from live_layers import transform_ais_frame, upsert_vessel

logger = logging.getLogger("osint.aisstream")

WS_URL = "wss://stream.aisstream.io/v0/stream"
FILTER_TYPES = [
    "PositionReport",
    "StandardClassBPositionReport",
    "ExtendedClassBPositionReport",
    "ShipStaticData",
]

# ── Viewport-following ─────────────────────────────────────────────────────
# The frontend POSTs its current viewport box to /api/vessels/subscribe; the
# worker retunes the AISStream subscription to it (throttled to 1/s, the
# service's subscription-update cap). ``None`` keeps the env AISSTREAM_BBOX
# default. Last writer wins; the key never reaches the browser.
_desired_boxes: list[list[list[float]]] | None = None
_bbox_guard = asyncio.Lock()


async def request_viewport_bbox(
    minlon: float, minlat: float, maxlon: float, maxlat: float
) -> None:
    """Retune the live subscription to a viewport box (lon/lat input order)."""
    global _desired_boxes
    box = [[minlat, minlon], [maxlat, maxlon]]  # AISStream wants [lat, lon] corners
    async with _bbox_guard:
        _desired_boxes = [box]


async def reset_viewport_bbox() -> None:
    """Fall back to the env AISSTREAM_BBOX default."""
    global _desired_boxes
    async with _bbox_guard:
        _desired_boxes = None


async def _take_desired_boxes() -> list[list[list[float]]] | None:
    async with _bbox_guard:
        return _desired_boxes


def _parse_boxes(raw: str) -> list[list[list[float]]]:
    """Env format: minlat,minlon,maxlat,maxlon[; ...]. AIS wants [[lat,lon],[lat,lon]]."""
    boxes = []
    for chunk in (raw or "").split(";"):
        parts = [p.strip() for p in chunk.split(",") if p.strip()]
        if len(parts) != 4:
            continue
        minlat, minlon, maxlat, maxlon = (float(p) for p in parts)
        boxes.append([[minlat, minlon], [maxlat, maxlon]])
    return boxes or [[[24.0, -125.0], [50.0, -66.0]]]


async def _resolve_key() -> str:
    return (os.getenv("AISSTREAM_API_KEY") or AISSTREAM_API_KEY
            or (await get_api_key("AISSTREAM_API_KEY")) or "").strip()


async def run_ais_worker() -> None:
    """Long-lived reconnect loop. Safe to spawn as an asyncio task."""
    try:
        import websockets
    except ImportError:
        logger.warning("websockets package not installed — AIS worker disabled")
        return

    backoff = 2.0
    while True:
        key = await _resolve_key()
        if not key:
            logger.warning(
                "AISSTREAM_API_KEY not set — AIS ingest idle. "
                "Create a free key at https://aisstream.io/account"
            )
            await asyncio.sleep(60)
            continue
        boxes = await _take_desired_boxes() or _parse_boxes(AISSTREAM_BBOX)
        try:
            async with websockets.connect(
                WS_URL,
                max_size=2 ** 22,
                ping_interval=20,
                ping_timeout=20,
                compression="deflate",
            ) as ws:
                sub = {
                    "APIKey": key,
                    "BoundingBoxes": boxes,
                    "FilterMessageTypes": FILTER_TYPES,
                }
                await ws.send(json.dumps(sub))
                logger.info("AISStream subscribed (%d bbox(es))", len(boxes))
                backoff = 2.0
                last_submit = time.monotonic()
                while True:
                    # Follow the client viewport: coalesce to the latest request
                    # and honor AISStream's 1 subscription-update/s cap.
                    desired = await _take_desired_boxes()
                    if (
                        desired is not None
                        and desired != boxes
                        and time.monotonic() - last_submit >= 1.0
                    ):
                        boxes = desired
                        sub["BoundingBoxes"] = boxes
                        await ws.send(json.dumps(sub))
                        last_submit = time.monotonic()
                        logger.info(
                            "AISStream re-subscribed to viewport (%d bbox)",
                            len(boxes),
                        )
                    try:
                        raw = await asyncio.wait_for(ws.recv(), timeout=0.25)
                    except asyncio.TimeoutError:
                        continue
                    if isinstance(raw, bytes):
                        raw = raw.decode("utf-8", errors="replace")
                    try:
                        frame = json.loads(raw)
                    except json.JSONDecodeError:
                        continue
                    marker = transform_ais_frame(frame)
                    if marker:
                        await upsert_vessel(marker)
        except asyncio.CancelledError:
            raise
        except Exception:  # noqa: BLE001
            logger.exception("AISStream disconnected")
        delay = backoff + random.uniform(0, 1.5)
        logger.info("AISStream reconnect in %.1fs", delay)
        await asyncio.sleep(delay)
        backoff = min(60.0, backoff * 1.7)