squid-manager/squid_mgr.py

"""squid_mgr.py — interface to Squid's built-in `squidclient mgr:` admin endpoints.

The Squid HTTP proxy exposes a management interface on the cache_object port
(default 3128). `squidclient mgr:<section>` returns plain-text key/value
information for sections like `info`, `5min`, `counters`, `io`, `storedir`, etc.

This module is intentionally dependency-free (stdlib only) so it can be reused
in the Flask app, the alert evaluator, and unit tests without installing extra
packages.

Public surface
--------------
fetch_mgr_info(...)         -> (ok, raw_text, error_msg)
parse_mgr_info(text)        -> dict[str, str]
fetch_and_parse(...)        -> (ok, parsed_dict, error_msg)
parse_counters(text)        -> list[(label, value)]
get_performance_summary(...) -> dict  (the high-level payload the UI consumes)

All parsers are forgiving: malformed lines are skipped silently, never raised.
This matters because Squid output varies between versions and locales.
"""
from __future__ import annotations

import re
import subprocess
from datetime import datetime, timezone


# ---------------------------------------------------------------------------
# Fetching
# ---------------------------------------------------------------------------

def fetch_mgr_info(squidclient: str = "squidclient",
                   host: str = "127.0.0.1",
                   port: int = 3128,
                   mgr: str = "info",
                   password: str = "",
                   timeout: int = 5) -> tuple[bool, str, str]:
    """Run `squidclient -h host -p port mgr:<section>` and capture stdout.

    Returns (ok, raw_text, error_msg).
    - ok is True iff the process exited with status 0 AND produced output.
    - error_msg is the empty string on success; otherwise a short diagnostic.
    - On timeout we surface a clear message so the caller can degrade gracefully.
    """
    cmd = [squidclient, "-h", str(host), "-p", str(port), f"mgr:{mgr}"]
    if password:
        cmd[2:2] = ["-w", password]  # insert -w <pw> right after the binary
    try:
        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except FileNotFoundError as exc:
        return False, "", f"squidclient not found: {exc}"
    except subprocess.TimeoutExpired:
        return False, "", f"squidclient mgr:{mgr} timeout after {timeout}s"
    except Exception as exc:  # pragma: no cover - defensive
        return False, "", f"squidclient mgr:{mgr} failed: {exc}"

    if proc.returncode != 0:
        err = (proc.stderr or proc.stdout or "").strip()
        return False, proc.stdout or "", f"exit={proc.returncode} {err[:300]}"
    if not proc.stdout.strip():
        return False, "", f"squidclient mgr:{mgr} returned empty output"
    return True, proc.stdout, ""


# ---------------------------------------------------------------------------
# Parsing helpers
# ---------------------------------------------------------------------------

# Section headers look like "Connection information:" (a key ending with a
# trailing colon). Within a section, indented lines hold `Key: value` pairs.
# `mgr:5min`/`mgr:counters` outputs are flat `key = value` (no indentation),
# so we allow either: 2+ spaces of indent with ANY separator, OR no indent
# but an '=' separator (so we don't accidentally grab plain English text).
_SECTION_RE = re.compile(r"^\s*([A-Z][A-Za-z0-9 /%()\-]+):\s*$")
_INDENTED_KV_RE = re.compile(r"^\s{2,}([A-Za-z0-9 _/%.()\-]+?)\s*[:=]\s*(.+?)\s*$")
_FLAT_KV_RE = re.compile(r"^([A-Za-z0-9_.]+)\s*=\s*(.+?)\s*$")


def _strip_inline_section(line: str) -> tuple[str | None, str | None]:
    """First-line special case: "Squid Object Cache (Version 5.7): John Doe"
    which mixes a header and a value on one line. Returns ("Squid Object Cache",
    "(Version 5.7): John Doe") so the normal KV regex can pick it up.
    """
    return None, None  # handled by caller via simpler regex


def parse_mgr_info(text: str) -> dict:
    """Parse a `squidclient mgr:info` text blob into a flat dict.

    The shape of mgr:info output (see task example):
        Squid Object Cache (Version 5.7): John Doe
            Start Time:     Mon, 14 Jul 2026 12:00:00 GMT
            Current Time:   Mon, 14 Jul 2026 13:00:00 GMT
        Connection information:
            Number of clients accessing cache: 5
        Cache information:
            Hits as % of bytes sent: 35.2%
        ...

    We collect every `Key: value` line as `Key -> value`. Section headers are
    preserved as `Section:` entries so callers can correlate them if needed.
    Values are kept as strings; numeric coercion is the job of the summary
    builder.
    """
    out: dict[str, str] = {}
    if not text:
        return out

    for raw in text.splitlines():
        line = raw.rstrip()
        if not line.strip():
            continue

        # Section header: ends with ":", no value on the same line
        m_sec = _SECTION_RE.match(line)
        if m_sec:
            out[m_sec.group(1).strip()] = ""  # marker, no value
            continue

        # Indented key/value line (2+ spaces of indent, `:` or `=` separator).
        # Used by mgr:info.
        m_kv = _INDENTED_KV_RE.match(line)
        if m_kv:
            key = m_kv.group(1).strip()
            val = m_kv.group(2).strip()
            out[key] = val
            continue

        # Flat key/value line (no indent, '=' separator, snake_case key).
        # Used by mgr:5min / mgr:counters.
        m_flat = _FLAT_KV_RE.match(line)
        if m_flat:
            key = m_flat.group(1).strip()
            val = m_flat.group(2).strip()
            out[key] = val
            continue

        # Fallback: "Squid Object Cache (Version 5.7): John Doe" — header +
        # value on the same line. Try a relaxed split on ":".
        if ":" in line:
            k, v = line.split(":", 1)
            k = k.strip()
            if k and v.strip():
                out.setdefault(k, v.strip())

    return out


# ---------------------------------------------------------------------------
# Field extractors (named accessors used by get_performance_summary)
# ---------------------------------------------------------------------------

_TIME_FIELD_RE = re.compile(
    r"(?:(?P<days>\d+)\s+days?,\s*)?(?P<hours>\d+):(?P<minutes>\d+):(?P<seconds>\d+(?:\.\d+)?)"
)


def parse_uptime_to_seconds(text: str) -> int | None:
    """Convert Squid uptime strings like "3 days, 12:34:56" or "0:05:30" to seconds."""
    if not text:
        return None
    m = _TIME_FIELD_RE.search(text)
    if not m:
        return None
    days = int(m.group("days") or 0)
    hours = int(m.group("hours") or 0)
    minutes = int(m.group("minutes") or 0)
    seconds = float(m.group("seconds") or 0)
    return int(days * 86400 + hours * 3600 + minutes * 60 + seconds)


_BYTE_UNITS = {
    "B": 1,
    "KB": 1024,
    "MB": 1024 ** 2,
    "GB": 1024 ** 3,
    "TB": 1024 ** 4,
    "K": 1024,
    "M": 1024 ** 2,
    "G": 1024 ** 3,
}

_BYTES_RE = re.compile(
    r"^\s*(?P<num>[0-9]+(?:\.[0-9]+)?)\s*(?P<unit>[A-Za-z]+)?\s*$"
)


def parse_bytes_to_int(text: str) -> int | None:
    """Convert strings like "1.2 GB" or "256 MB" or "12345" to an int byte count."""
    if text is None:
        return None
    s = text.strip()
    if not s:
        return None
    m = _BYTES_RE.match(s)
    if not m:
        # Try to pull a number out of mixed strings ("1,234,567")
        digits = re.sub(r"[^0-9.]", "", s)
        try:
            return int(float(digits)) if digits else None
        except ValueError:
            return None
    num = float(m.group("num"))
    unit = (m.group("unit") or "B").upper()
    factor = _BYTE_UNITS.get(unit)
    if factor is None:
        # Unknown unit — treat as raw count
        return int(num)
    return int(num * factor)


def parse_percent(text: str) -> float | None:
    """Convert strings like "35.2%" or "35.2" to a float percentage value."""
    if not text:
        return None
    s = text.strip().rstrip("%").strip()
    try:
        return float(s)
    except ValueError:
        return None


_VERSION_RE = re.compile(r"Version\s+(\d+(?:\.\d+)*)")


def parse_version(text: str) -> str | None:
    """Pull '5.7' out of a header like 'Squid Object Cache (Version 5.7)'."""
    if not text:
        return None
    m = _VERSION_RE.search(text)
    return m.group(1) if m else None


# ---------------------------------------------------------------------------
# Counters parsing (mgr:counters — for chart rendering)
# ---------------------------------------------------------------------------

def parse_counters(text: str) -> list[tuple[str, str]]:
    """Parse `mgr:counters` output, returning [(label, value), ...].

    mgr:counters emits blank-line-separated blocks, each with a counter name
    on the first line and `sample_time = NN, client_http.requests = 12345`
    style key=value pairs beneath. We flatten everything into (label, value)
    pairs skipping the 'sample_time' bookkeeping rows so the UI can show a
    simple table.
    """
    out: list[tuple[str, str]] = []
    if not text:
        return out
    current_block = ""
    for raw in text.splitlines():
        line = raw.rstrip()
        if not line.strip():
            current_block = ""
            continue
        if not line.startswith((" ", "\t")):
            # New block header (e.g. "client_http.requests")
            current_block = line.strip()
            out.append((current_block, ""))
            continue
        # Inside a block: "key = value"
        if "=" in line:
            k, v = line.split("=", 1)
            k = k.strip()
            v = v.strip()
            if k == "sample_time":
                continue
            label = f"{current_block}.{k}" if current_block else k
            out.append((label, v))
    return out


# ---------------------------------------------------------------------------
# One-step convenience wrappers
# ---------------------------------------------------------------------------

def fetch_and_parse(mgr: str = "info", **kwargs) -> tuple[bool, dict, str]:
    """fetch_mgr_info + parse_mgr_info in one call.

    Returns (ok, parsed_dict, error_msg).
    """
    ok, raw, err = fetch_mgr_info(mgr=mgr, **kwargs)
    if not ok:
        return False, {}, err
    return True, parse_mgr_info(raw), ""


# ---------------------------------------------------------------------------
# High-level aggregator
# ---------------------------------------------------------------------------

def _fmt_now() -> str:
    return datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S")


def _safe_first(d: dict, *keys: str) -> str:
    """Return the first non-empty value among keys, or ''."""
    for k in keys:
        v = d.get(k)
        if v not in (None, ""):
            return str(v)
    return ""


def get_performance_summary(squidclient: str = "squidclient", **kwargs) -> dict:
    """Aggregate the most important mgr:info metrics into a single dict.

    Pulls mgr:info, then optionally mgr:5min for request/byte rates. Other
    mgr:* sections (counters, storedir, io) can be fetched on demand by the
    UI; this function only does the must-have numbers so the page loads fast.

    The returned dict is the contract for both the rendered template and the
    /performance/refresh JSON endpoint.
    """
    fetched_at = _fmt_now()
    summary = {
        "squid_uptime_seconds": None,
        "squid_uptime_human": "",
        "squid_version": None,
        "squid_author": "",
        "cpu_time_seconds": None,
        "cpu_time_human": "",
        "max_rss_mb": None,
        "max_rss_human": "",
        "current_clients": None,
        "cache_hits_pct": None,
        "storage_swap_mb": None,
        "storage_swap_human": "",
        "storage_mem_mb": None,
        "storage_mem_human": "",
        "request_rate_5min": None,   # requests / minute
        "byte_rate_5min": None,      # bytes / second
        "start_time": "",
        "current_time": "",
        "fetched_ok": False,
        "fetched_at": fetched_at,
        "error": None,
        "raw_info": "",  # included so the UI can show it in a <details> block
    }

    # ---- mgr:info ----
    ok, raw, err = fetch_mgr_info(squidclient=squidclient, mgr="info", **kwargs)
    summary["raw_info"] = raw or ""
    if not ok:
        summary["error"] = err
        return summary

    info = parse_mgr_info(raw)

    # Squid Object Cache (Version 5.7): John Doe
    # parse_mgr_info's fallback split puts the (Version X.Y) literal in the
    # KEY (with the author trailing) — or in the value, depending on the
    # line. We find whichever side contains "Version" and parse from there.
    header_full = ""
    header_author = ""
    for k, v in info.items():
        if k.startswith("Squid Object Cache"):
            if "Version" in k:
                header_full = k
                header_author = (v or "").strip()
                break
            if "Version" in v:
                header_full = f"{k}: {v}" if v else k
                break

    if header_full:
        summary["squid_version"] = parse_version(header_full)
        if header_author:
            summary["squid_author"] = header_author
        elif "):" in header_full:
            summary["squid_author"] = header_full.split("):", 1)[1].strip()
        elif ":" in header_full:
            summary["squid_author"] = header_full.split(":", 1)[1].strip()

    # Uptime: from Start Time + Current Time if Squid didn't print it directly.
    start_time_s = _safe_first(info, "Start Time")
    current_time_s = _safe_first(info, "Current Time")
    summary["start_time"] = start_time_s
    summary["current_time"] = current_time_s
    uptime_s = _parse_duration_diff_seconds(start_time_s, current_time_s)
    if uptime_s is not None:
        summary["squid_uptime_seconds"] = uptime_s
        summary["squid_uptime_human"] = _seconds_to_human(uptime_s)

    # Connection information: Number of clients accessing cache
    clients = _safe_first(info, "Number of clients accessing cache")
    if clients:
        try:
            summary["current_clients"] = int(re.sub(r"[^0-9]", "", clients) or 0)
        except ValueError:
            summary["current_clients"] = None

    # Cache information: Hits as % of bytes sent
    hits_pct = _safe_first(info, "Hits as % of bytes sent")
    summary["cache_hits_pct"] = parse_percent(hits_pct)

    # Storage info
    swap = _safe_first(info, "Storage Swap size")
    summary["storage_swap_human"] = swap
    summary["storage_swap_mb"] = _bytes_to_mb(parse_bytes_to_int(swap))

    mem = _safe_first(info, "Storage Mem size")
    summary["storage_mem_human"] = mem
    summary["storage_mem_mb"] = _bytes_to_mb(parse_bytes_to_int(mem))

    # Resource usage
    cpu = _safe_first(info, "CPU Time")
    summary["cpu_time_human"] = cpu
    summary["cpu_time_seconds"] = parse_uptime_to_seconds(cpu)

    rss = _safe_first(info, "Maximum Resident Size")
    summary["max_rss_human"] = rss
    summary["max_rss_mb"] = _bytes_to_mb(parse_bytes_to_int(rss))

    # ---- mgr:5min (best-effort, don't fail the whole page if missing) ----
    ok5, raw5, _err5 = fetch_mgr_info(
        squidclient=squidclient, mgr="5min", **kwargs
    )
    if ok5:
        five = parse_mgr_info(raw5)
        sample = _safe_first(five, "sample_time")
        # Squid prints per-5min averages with names like:
        #   "client_http.requests = 1234" and "client_http.kbytes_in/out = ..."
        # Compute deltas ourselves if the same metric appears with a previous
        # sample; for an MVP we just read the absolute number — labelled
        # "5min 平均" by the UI.
        req = _safe_first(five, "client_http.requests")
        if req:
            try:
                summary["request_rate_5min"] = float(req)
            except ValueError:
                pass
        kb_in = _safe_first(five, "client_http.kbytes_in")
        kb_out = _safe_first(five, "client_http.kbytes_out")
        total_kb = 0.0
        for s in (kb_in, kb_out):
            try:
                total_kb += float(s)
            except (ValueError, TypeError):
                pass
        if total_kb > 0:
            # Squid's 5min averages are reported per-second (kbytes/sec).
            summary["byte_rate_5min"] = int(total_kb * 1024)
        _ = sample  # currently unused but kept for future deltas

    summary["fetched_ok"] = True
    return summary


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _bytes_to_mb(n: int | None) -> int | None:
    if n is None:
        return None
    return int(n / (1024 * 1024))


def _seconds_to_human(sec: int) -> str:
    """Format seconds into 'Xd HH:MM:SS' or 'HH:MM:SS'."""
    days, rem = divmod(sec, 86400)
    hours, rem = divmod(rem, 3600)
    minutes, seconds = divmod(rem, 60)
    if days:
        return f"{days}d {hours:02d}:{minutes:02d}:{seconds:02d}"
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"


# Accept RFC 822 / RFC 1123 / "YYYY-MM-DD HH:MM:SS" — best-effort.
_DATE_FORMATS = (
    "%a, %d %b %Y %H:%M:%S %Z",     # RFC 822 with named TZ (Squid mgr:info)
    "%a, %d %b %Y %H:%M:%S %z",
    "%a, %d %b %Y %H:%M:%S GMT",
    "%Y-%m-%d %H:%M:%S",
)


def _parse_duration_diff_seconds(start: str, current: str) -> int | None:
    if not start or not current:
        return None
    # Strip trailing timezone names that %Z may not recognize on all systems
    # (Squid's mgr:info uses "GMT"; Python's %Z accepts it on Linux).
    s_dt = _try_parse(start)
    c_dt = _try_parse(current)
    if s_dt is None or c_dt is None:
        return None
    delta = (c_dt - s_dt).total_seconds()
    return int(delta) if delta >= 0 else None


def _try_parse(s: str):
    s = (s or "").strip()
    if not s:
        return None
    for fmt in _DATE_FORMATS:
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            continue
    # last resort: split off a trailing TZ like "GMT" / "UTC"
    parts = s.rsplit(" ", 1)
    if len(parts) == 2 and parts[1].upper() in {"GMT", "UTC"}:
        for fmt in _DATE_FORMATS:
            try:
                return datetime.strptime(parts[0], fmt)
            except ValueError:
                continue
    return None