FabledSteward/plugins/syslog/receiver.py

# plugins/syslog/receiver.py
"""The syslog listener: UDP and TCP on one port, every line taken as it
arrives and stored a few seconds later.

A line is never written from inside the socket's callback. It is queued with
the time it arrived and the address it came from, and the receive task, every
few seconds, reads what is queued into a round and hands it to a Holding
(steward/core/holding.py). The NAS that stalls the storage under Postgres is
one of the devices sending, so its lines about the stall must wait out the
stall in memory, not be lost to it.

The listener follows the plugin's settings: the receive task binds it, and
binds it again when the port or a protocol is changed, so a save in Settings
takes effect within seconds with no restart (rule 25). A port that can't be
bound is said on the Syslog page, and tried again every minute.
"""
from __future__ import annotations

import asyncio
import logging
from collections import deque
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING

from sqlalchemy import func, select

from steward.core.scheduler import ScheduledTask

from .parse import MAX_MESSAGE, Message, frames, parse

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# Lines received and not yet read into a round. Past this the oldest go, and
# the page says how many: a flood is bounded in memory.
MAX_QUEUED = 20_000
# The most lines one round takes; the rest wait for the next.
ROUND_LINES = 2_000
# Rounds held while the database doesn't take them: half an hour at the
# five-second receive.
HELD_ROUNDS = 360
RECEIVE_S = 5
# How long a receive waits for its own store before the next one may start.
STORE_WAIT_S = 4.0
# How long after a failed bind it is tried again.
RETRY_BIND = timedelta(minutes=1)
# Senders known at most, added or offered: a new one past them is counted,
# not kept, so a spray of forged UDP sources can't fill the table.
MAX_SENDERS = 200
# TCP connections at once, and how long one may say nothing before it is
# closed (a sender reconnects when it next has a line).
MAX_CONNECTIONS = 64
TCP_IDLE_S = 3600.0
# The sample of a source's newest line kept on its row.
SAMPLE_CHARS = 500


@dataclass
class Listener:
    """What the listener is bound to, and what it has taken since it started.
    The page reads it as it is; nothing here is stored."""
    port: int | None = None
    udp: bool = False
    tcp: bool = False
    error: str | None = None
    tried_at: datetime | None = None
    bound_at: datetime | None = None
    received: int = 0
    dropped: int = 0          # queued lines lost to MAX_QUEUED
    unoffered: int = 0        # lines from senders past MAX_SENDERS
    connections: int = 0
    _udp_transport: object = None
    _tcp_server: object = None

    @property
    def listening(self) -> bool:
        return self._udp_transport is not None or self._tcp_server is not None


LISTENER = Listener()
# (received_at, address, raw)
QUEUE: deque[tuple[datetime, str, bytes]] = deque()


def take(address: str, raw: bytes, now: datetime | None = None) -> None:
    """Queue one received message."""
    if not raw:
        return
    LISTENER.received += 1
    QUEUE.append((now or datetime.now(timezone.utc), address, raw[:MAX_MESSAGE * 2]))
    while len(QUEUE) > MAX_QUEUED:
        QUEUE.popleft()
        LISTENER.dropped += 1


class _Udp(asyncio.DatagramProtocol):
    def datagram_received(self, data: bytes, addr) -> None:
        take(addr[0], data)


async def _tcp_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    peer = writer.get_extra_info("peername")
    address = peer[0] if peer else "?"
    if LISTENER.connections >= MAX_CONNECTIONS:
        writer.close()
        return
    LISTENER.connections += 1
    buffer = bytearray()
    try:
        while chunk := await asyncio.wait_for(reader.read(65536), TCP_IDLE_S):
            buffer += chunk
            for message in frames(buffer):
                take(address, message)
        for message in frames(buffer + b"\n"):      # the last line, unended
            take(address, message)
    except (ConnectionError, asyncio.IncompleteReadError, TimeoutError):
        pass
    finally:
        LISTENER.connections -= 1
        writer.close()


def settings(app: "Quart") -> tuple[int, bool, bool]:
    cfg = (app.config.get("PLUGINS") or {}).get("syslog") or {}
    return int(cfg.get("port", 5514)), bool(cfg.get("udp", True)), bool(cfg.get("tcp", True))


async def close() -> None:
    if LISTENER._udp_transport is not None:
        LISTENER._udp_transport.close()
    if LISTENER._tcp_server is not None:
        LISTENER._tcp_server.close()
        await LISTENER._tcp_server.wait_closed()
    LISTENER._udp_transport = LISTENER._tcp_server = None
    LISTENER.bound_at = None


async def ensure_listening(app: "Quart", now: datetime | None = None) -> None:
    """Bind the listener as the settings say, when it isn't already: at
    start, after a change of port or protocol, and a minute after a failure."""
    now = now or datetime.now(timezone.utc)
    port, udp, tcp = settings(app)
    same = (port, udp, tcp) == (LISTENER.port, LISTENER.udp, LISTENER.tcp)
    if same and LISTENER.error is None and (LISTENER.listening or not (udp or tcp)):
        return
    if same and LISTENER.error and LISTENER.tried_at and now - LISTENER.tried_at < RETRY_BIND:
        return
    await close()
    LISTENER.port, LISTENER.udp, LISTENER.tcp = port, udp, tcp
    LISTENER.error, LISTENER.tried_at = None, now
    loop = asyncio.get_running_loop()
    failed = []
    if udp:
        try:
            LISTENER._udp_transport, _ = await loop.create_datagram_endpoint(
                _Udp, local_addr=("0.0.0.0", port))
        except OSError as exc:
            failed.append(f"UDP: {exc.strerror or exc}")
    if tcp:
        try:
            LISTENER._tcp_server = await asyncio.start_server(_tcp_client, "0.0.0.0", port)
        except OSError as exc:
            failed.append(f"TCP: {exc.strerror or exc}")
    if failed:
        LISTENER.error = f"Port {port} couldn't be opened ({'; '.join(failed)})."
        logger.warning("Syslog: %s", LISTENER.error)
    if LISTENER.listening:
        LISTENER.bound_at = now
        bound = (("UDP", LISTENER._udp_transport), ("TCP", LISTENER._tcp_server))
        logger.info("Syslog: listening on %s port %d", " and ".join(p for p, on in bound if on), port)


@dataclass
class Received:
    at: datetime
    address: str
    message: Message


def read_round(limit: int = ROUND_LINES) -> list[Received]:
    """Take up to `limit` queued lines, oldest first, each read into its parts."""
    out = []
    while QUEUE and len(out) < limit:
        at, address, raw = QUEUE.popleft()
        out.append(Received(at, address, parse(raw)))
    return out


def make_receive_task(app: "Quart") -> ScheduledTask:
    async def receive() -> None:
        await _receive(app)

    return ScheduledTask(name="syslog_receive", coro_factory=receive,
                         interval_seconds=RECEIVE_S, run_on_startup=True)


async def _receive(app: "Quart") -> None:
    from steward.core.holding import holding

    await ensure_listening(app)
    held = holding(app, "Syslog", store=lambda rounds: _store(app, rounds), limit=HELD_ROUNDS)
    lines = read_round()
    if lines:
        held.hold(lines)
    await held.flush(STORE_WAIT_S)


async def _store(app: "Quart", rounds: list[list[Received]]) -> None:
    """Write held rounds, oldest first, in one transaction: each sender's
    count, newest line and when it was last heard, and the lines of the
    sources that have been added. A sender not seen before becomes an
    offered source; its lines are counted and not kept."""
    async with app.db_sessionmaker() as session:
        async with session.begin():
            await store_lines(session, [line for lines in rounds for line in lines])


async def store_lines(session, lines: list[Received]) -> None:
    from .models import SyslogLine, SyslogSource

    if not lines:
        return
    by_address: dict[str, list[Received]] = {}
    for line in lines:
        by_address.setdefault(line.address, []).append(line)
    sources = {s.address: s for s in (await session.execute(
        select(SyslogSource).where(SyslogSource.address.in_(list(by_address))))).scalars()}
    new = [a for a in by_address if a not in sources]
    if new:
        known = (await session.execute(select(func.count()).select_from(SyslogSource))).scalar_one()
        for address in new:
            if known >= MAX_SENDERS:
                LISTENER.unoffered += len(by_address.pop(address))
                continue
            first = by_address[address][0]
            source = SyslogSource(address=address, first_seen_at=first.at, last_seen_at=first.at)
            session.add(source)
            sources[address] = source
            known += 1
        await session.flush()

    for address, received in by_address.items():
        source = sources[address]
        newest = received[-1]
        source.lines_total = (source.lines_total or 0) + len(received)
        source.last_seen_at = max(source.last_seen_at, newest.at) if source.last_seen_at else newest.at
        source.last_severity = newest.message.severity
        source.last_line = newest.message.text[:SAMPLE_CHARS]
        named = next((r.message.hostname for r in reversed(received) if r.message.hostname), "")
        if named:
            source.hostname_seen = named[:255]
        if source.state == "accepted":
            session.add_all([SyslogLine(
                source_id=source.id, received_at=r.at, sent_at=r.message.sent_at,
                facility=r.message.facility, severity=r.message.severity,
                hostname=r.message.hostname, app=r.message.app, message=r.message.text,
            ) for r in received])