links/routermon/receiver.py

"""
UDP syslog receiver for routermon.

Architecture:
  - One background thread owns the UDP socket and enqueues parsed events.
  - A separate writer thread drains the queue and bulk-saves to the DB every
    FLUSH_INTERVAL seconds (or when the queue reaches FLUSH_SIZE).
  - Only one Gunicorn worker binds the socket (EADDRINUSE is silently ignored
    in the others — matches the project's existing in-process APScheduler style).

dnsmasq reply correlation:
  - A small in-memory dict `_pending` maps `domain` → (pk, timestamp) for
    outstanding query rows awaiting a reply.  Entries expire after REPLY_TTL
    seconds so stale pending records don't accumulate.
"""
import errno
import logging
import queue
import socket
import threading
import time
from datetime import timedelta

from django.utils import timezone

logger = logging.getLogger(__name__)

FLUSH_INTERVAL = 2       # seconds between writer flushes
FLUSH_SIZE = 100         # flush immediately when queue reaches this size
REPLY_TTL = 10           # seconds to keep a query in _pending for reply matching
MAX_QUEUE = 10_000       # drop oldest items if queue exceeds this (backpressure)

_receiver_thread: threading.Thread | None = None
_writer_thread: threading.Thread | None = None
_sock: socket.socket | None = None
_queue: queue.Queue = queue.Queue(maxsize=MAX_QUEUE)

# in-memory pending dict: domain → (query_pk, created_at)
_pending: dict = {}
_pending_lock = threading.Lock()


# ── Public API ────────────────────────────────────────────────────────────────

def start_receiver(port: int):
    """
    Attempt to start the UDP receiver on `port`.  Safe to call from every
    Gunicorn worker — only the first to bind succeeds; others silently skip.
    """
    global _receiver_thread, _writer_thread, _sock

    if _receiver_thread and _receiver_thread.is_alive():
        return  # already running in this process

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # No SO_REUSEADDR: we intentionally want only one worker to bind.
        # Other workers receive EADDRINUSE and skip cleanly.
        sock.bind(('0.0.0.0', port))
        sock.settimeout(1.0)
        _sock = sock
        print(f'routermon: UDP syslog receiver bound on port {port}', flush=True)
        logger.info('routermon: UDP syslog receiver bound on port %d', port)
    except OSError as exc:
        if exc.errno == errno.EADDRINUSE:
            logger.debug('routermon: port %d already bound (another worker), skipping', port)
        else:
            import sys
            print(f'routermon: failed to bind UDP port {port}: {exc}', file=sys.stderr, flush=True)
            logger.error('routermon: failed to bind UDP port %d: %s', port, exc)
        return

    _writer_thread = threading.Thread(target=_writer_loop, daemon=True, name='routermon-writer')
    _writer_thread.start()

    _receiver_thread = threading.Thread(
        target=_receiver_loop, args=(sock,), daemon=True, name='routermon-receiver'
    )
    _receiver_thread.start()


def stop_receiver():
    global _sock
    if _sock:
        try:
            _sock.close()
        except Exception:
            pass
        _sock = None


# ── Socket receiver loop ──────────────────────────────────────────────────────

def _receiver_loop(sock: socket.socket):
    from .parser import parse_line

    logger.info('routermon: receiver loop started')
    while True:
        try:
            data, _ = sock.recvfrom(4096)
        except socket.timeout:
            continue
        except OSError:
            # Socket closed (shutdown)
            break
        try:
            line = data.decode('utf-8', errors='replace').rstrip('\n\r')
            result = parse_line(line)
            if result is not None:
                _enqueue(result)
        except Exception as exc:
            logger.debug('routermon: parse error: %s', exc)

    logger.info('routermon: receiver loop stopped')


def _enqueue(result):
    try:
        _queue.put_nowait(result)
    except queue.Full:
        # Drop oldest item to make room (keep most recent)
        try:
            _queue.get_nowait()
            _queue.put_nowait(result)
        except Exception:
            pass


# ── Writer loop ───────────────────────────────────────────────────────────────

def _writer_loop():
    logger.info('routermon: writer loop started')
    batch = []
    last_flush = time.monotonic()

    while True:
        # Collect items until flush threshold or timeout
        deadline = last_flush + FLUSH_INTERVAL
        while time.monotonic() < deadline and len(batch) < FLUSH_SIZE:
            try:
                item = _queue.get(timeout=max(0.1, deadline - time.monotonic()))
                batch.append(item)
            except queue.Empty:
                break

        if batch:
            try:
                _flush(batch)
            except Exception as exc:
                logger.error('routermon: flush error: %s', exc, exc_info=True)
            batch = []

        last_flush = time.monotonic()

    logger.info('routermon: writer loop stopped')


def _flush(batch: list):
    """Persist a batch of ParseResult objects to the database."""
    from .models import RouterMonSettings, DnsQuery, DhcpLease

    settings = RouterMonSettings.get()
    if not settings.enabled:
        return

    excluded = _build_excluded_set(settings.excluded_clients)
    now = timezone.now()

    # --- DHCP leases (upsert immediately) ---
    dhcp_events = [r.dhcp for r in batch if r.dhcp]
    for ev in dhcp_events:
        DhcpLease.objects.update_or_create(
            ip=ev.ip,
            defaults={'mac': ev.mac, 'hostname': ev.hostname},
        )

    # --- DNS queries ---
    hostname_cache = {
        lease.ip: lease.hostname
        for lease in DhcpLease.objects.filter(
            ip__in={r.query.client_ip for r in batch if r.query}
        )
    }

    query_objs = []
    for r in batch:
        if not r.query:
            continue
        q = r.query
        if q.client_ip in excluded:
            continue
        obj = DnsQuery(
            timestamp=now,
            client_ip=q.client_ip,
            client_name=hostname_cache.get(q.client_ip, ''),
            domain=q.domain,
            query_type=q.query_type,
        )
        query_objs.append(obj)

    if query_objs:
        created = DnsQuery.objects.bulk_create(query_objs)
        # Register created queries in pending dict for reply correlation
        with _pending_lock:
            for obj in created:
                _pending[obj.domain] = (obj.pk, now)

    # --- DNS replies: correlate with pending queries ---
    reply_events = [r.reply for r in batch if r.reply]
    if reply_events:
        _apply_replies(reply_events, now)

    # Expire old pending entries
    _expire_pending(now)

    # Touch last_received_at once per flush (coalesced)
    RouterMonSettings.objects.filter(pk=1).update(last_received_at=now)

    # Trigger geo enrichment in background (lazy: only for rows with resolved_ip that have no geo yet)
    pks_needing_geo = list(
        DnsQuery.objects.filter(
            resolved_ip__isnull=False, country='', timestamp__gte=now - timedelta(minutes=5)
        ).values_list('pk', flat=True)[:200]
    )
    if pks_needing_geo:
        threading.Thread(
            target=_enrich_geo, args=(pks_needing_geo,), daemon=True
        ).start()


def _apply_replies(reply_events, now):
    from .models import DnsQuery

    with _pending_lock:
        for ev in reply_events:
            entry = _pending.pop(ev.domain, None)
            if entry is None:
                continue
            pk, _ = entry
            updates = {}
            if ev.is_nxdomain:
                updates['is_nxdomain'] = True
            if ev.resolved_ip:
                updates['resolved_ip'] = ev.resolved_ip
            if updates:
                DnsQuery.objects.filter(pk=pk).update(**updates)


def _expire_pending(now):
    cutoff = now - timedelta(seconds=REPLY_TTL)
    with _pending_lock:
        expired = [domain for domain, (_, ts) in _pending.items() if ts < cutoff]
        for domain in expired:
            del _pending[domain]


def _enrich_geo(pks: list[int]):
    """Geo-enrich DnsQuery rows by resolved_ip using nginxmon's IPGeoCache."""
    try:
        from nginxmon.models import IPGeoCache
        from routermon.models import DnsQuery
        import ipaddress

        rows = list(DnsQuery.objects.filter(pk__in=pks).values('pk', 'resolved_ip'))
        ips = list({r['resolved_ip'] for r in rows if r['resolved_ip']})
        if not ips:
            return

        # Fetch missing IPs from ip-api.com
        cached = {c.ip: c for c in IPGeoCache.objects.filter(ip__in=ips)}
        missing = [ip for ip in ips if ip not in cached]
        if missing:
            from nginxmon.geo import _lookup_batch
            api_results = _lookup_batch(missing)
            for ip, data in api_results.items():
                obj, _ = IPGeoCache.objects.get_or_create(ip=ip)
                obj.country = data.get('country', '')
                obj.country_code = data.get('countryCode', '')
                obj.region = data.get('regionName', '')
                obj.city = data.get('city', '')
                obj.lat = data.get('lat')
                obj.lon = data.get('lon')
                obj.isp = data.get('isp', '')
                obj.save()
                cached[ip] = obj

        for row in rows:
            geo = cached.get(row['resolved_ip'])
            if geo and geo.country:
                DnsQuery.objects.filter(pk=row['pk'], country='').update(
                    country=geo.country,
                    country_code=geo.country_code,
                    city=geo.city,
                )
    except Exception as exc:
        logger.warning('routermon: geo enrichment error: %s', exc)


def _build_excluded_set(text: str) -> set:
    result = set()
    for line in (text or '').splitlines():
        entry = line.strip()
        if entry and not entry.startswith('#'):
            result.add(entry)
    return result