ServerMonitor/notify.py

"""Notification, persistence, and alert deduplication logic."""

import json
import logging
import os
import tempfile
import time
import urllib.parse
import urllib.request

logger = logging.getLogger("servermonitor.notify")


def escape_html(text):
    """Escape ``&``, ``<`` and ``>`` for Telegram ``parse_mode=HTML``."""
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def send_telegram(bot_token, chat_id, text, timeout):
    """Send an HTML message via Telegram.

    If ``bot_token`` is empty the function runs in log-only mode: it logs the
    would-be message and returns ``True``.
    """
    if not bot_token:
        logger.info("[TELEGRAM LOG-ONLY] %s", text)
        return True

    url = f"https://api.telegram.org/bot{bot_token}/sendMessage"
    data = urllib.parse.urlencode(
        {
            "chat_id": chat_id,
            "text": text,
            "parse_mode": "HTML",
            "disable_web_page_preview": "true",
        }
    ).encode("utf-8")

    req = urllib.request.Request(
        url,
        data=data,
        method="POST",
        headers={
            "Content-Type": "application/x-www-form-urlencoded",
            "User-Agent": "ServerMonitor/1.0",
        },
    )

    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            response.read()
        return True
    except Exception as exc:  # noqa: BLE001
        logger.error("Failed to send Telegram message: %s", exc)
        return False


class StateStore:
    """JSON file persistence for check/alert state.

    The file is written atomically: a temporary file is created in the same
    directory and then moved into place. If the existing file is corrupt, a
    fresh empty state is returned.
    """

    def __init__(self, path="/data/state.json"):
        self.path = path
        self._state = {}

    def load(self):
        """Load state from disk."""
        if not os.path.exists(self.path):
            self._state = {}
            return

        try:
            with open(self.path, "r", encoding="utf-8") as f:
                self._state = json.load(f)
            if not isinstance(self._state, dict):
                self._state = {}
        except Exception as exc:  # noqa: BLE001
            logger.warning("State file %s corrupt or unreadable: %s. Starting fresh.", self.path, exc)
            self._state = {}

    def save(self):
        """Persist state to disk atomically."""
        directory = os.path.dirname(os.path.abspath(self.path)) or "."
        try:
            os.makedirs(directory, exist_ok=True)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Could not create state directory %s: %s", directory, exc)
            return

        try:
            fd, tmp_path = tempfile.mkstemp(
                dir=directory, prefix="state", suffix=".json.tmp"
            )
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                json.dump(self._state, f, indent=2)
                f.write("\n")
            os.replace(tmp_path, self.path)
        except Exception as exc:  # noqa: BLE001
            logger.error("Could not save state file %s: %s", self.path, exc)

    def get(self, name):
        """Return the stored state for ``name`` or a default OK state."""
        entry = self._state.get(name)
        if not isinstance(entry, dict):
            return {"state": "OK", "last_alert_ts": 0.0, "alert_count": 0}
        return {
            "state": entry.get("state", "OK"),
            "last_alert_ts": float(entry.get("last_alert_ts", 0.0)),
            "alert_count": int(entry.get("alert_count", 0)),
        }

    def set(self, name, state, last_alert_ts, alert_count):
        """Set the stored state for ``name``."""
        self._state[name] = {
            "state": state,
            "last_alert_ts": float(last_alert_ts),
            "alert_count": int(alert_count),
        }


class AlertEngine:
    """Decide whether a check should trigger a Telegram alert.

    - First CRIT transition: alert.
    - Repeated CRIT inside the reminder window: suppress.
    - Repeated CRIT after reminder window: reminder alert.
    - OK transition after CRIT: recovery alert (exactly once).

    The caller should set the current observed state with ``set_state`` before
    calling ``should_alert``.
    """

    def __init__(self):
        self._current_states = {}

    def set_state(self, name, state):
        """Set the freshly observed state for ``name`` (``OK`` or ``CRIT``)."""
        self._current_states[name] = state

    def should_alert(self, name, now, state_store, reminder_interval):
        """Return ``(should_send, is_recovery)``.

        ``state_store`` is updated in place when an alert decision is made.
        """
        current_state = self._current_states.get(name, "OK")
        prev = state_store.get(name)
        prev_state = prev["state"]
        last_alert_ts = prev["last_alert_ts"]
        alert_count = prev["alert_count"]

        if current_state == "CRIT":
            if prev_state != "CRIT":
                state_store.set(name, "CRIT", now, alert_count + 1)
                return True, False
            if now - last_alert_ts >= reminder_interval:
                state_store.set(name, "CRIT", now, alert_count + 1)
                return True, False
            return False, False

        if current_state == "OK" and prev_state == "CRIT":
            state_store.set(name, "OK", last_alert_ts, alert_count)
            # Only send a recovery message if an alert was actually sent for
            # this check before (e.g. not after a baseline run that observed
            # the failure without alerting).
            if alert_count > 0:
                return True, True
            return False, False

        return False, False