light-heart-labs--dreamserver/ods/scripts/healthcheck.py

#!/usr/bin/env python3
"""ODS — universal healthcheck.

Why this exists
--------------
A lot of containers use curl/wget for HEALTHCHECK instructions, but minimal images
(distro-less, scratch-ish, python slim) frequently do not include them.

This script provides a single, dependency-free healthcheck implementation that:
  - Works with *either* HTTP(S) endpoints or raw TCP sockets
  - Supports GET fallback when HEAD is blocked
  - Allows matching on status code ranges and/or response body regex
  - Emits structured output for debugging in CI

Usage
-----
  healthcheck.py http://localhost:8080/health
  healthcheck.py tcp://localhost:5432
  healthcheck.py localhost:5432

Options
-------
  --timeout SECONDS              Overall timeout for the request/connection
  --retries N                    Retry count (with small backoff)
  --method {HEAD,GET}            HTTP method (default: HEAD, with GET fallback)
  --expect-status 200,204,3xx    Allowed HTTP status codes/ranges
  --expect-body-regex REGEX      Regex to match in response body (GET only)
  --user-agent UA                Custom user-agent
  --json                         Emit machine-readable JSON result

Exit codes
----------
  0  Healthy
  1  Unhealthy (check failed)
  2  Usage / invalid input

Notes
-----
- For HTTP checks we prefer HEAD to avoid moving large bodies, but many
  frameworks disable HEAD or route it differently. We automatically fall back
  to GET when HEAD fails with method-related errors.
- For TCP checks we just attempt to connect. This validates listening and basic
  accept() path.
"""

from __future__ import annotations

import argparse
import json
import re
import socket
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import List, Optional, Sequence, Set, Tuple


# -----------------------------
# Data model
# -----------------------------


@dataclass(frozen=True)
class Result:
    ok: bool
    target: str
    kind: str  # http|tcp
    detail: str
    status: Optional[int] = None
    elapsed_ms: Optional[int] = None

    def to_json(self) -> str:
        return json.dumps(
            {
                "ok": self.ok,
                "target": self.target,
                "kind": self.kind,
                "detail": self.detail,
                "status": self.status,
                "elapsed_ms": self.elapsed_ms,
            },
            separators=(",", ":"),
        )


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


def _parse_target(raw: str) -> Tuple[str, str]:
    """Return (kind, normalized_target)."""
    if raw.startswith("http://") or raw.startswith("https://"):
        return ("http", raw)

    if raw.startswith("tcp://"):
        return ("tcp", raw[len("tcp://") :])

    # host:port shorthand
    if ":" in raw and not raw.startswith("["):
        return ("tcp", raw)

    raise ValueError("target must be http(s) URL, tcp://host:port, or host:port")


def _parse_host_port(raw: str) -> Tuple[str, int]:
    host, port_s = raw.rsplit(":", 1)
    host = host.strip()
    if not host:
        raise ValueError("host is empty")
    try:
        port = int(port_s)
    except ValueError as exc:
        raise ValueError("port must be an integer") from exc
    if not (1 <= port <= 65535):
        raise ValueError("port out of range (1-65535)")
    return (host, port)


def _parse_expected_status(expr: str) -> Set[int]:
    """Parse '200,204,3xx,401-403' => allowed status codes set."""
    allowed: Set[int] = set()
    for part in (p.strip() for p in expr.split(",") if p.strip()):
        if part.endswith("xx") and len(part) == 3 and part[0].isdigit():
            base = int(part[0]) * 100
            allowed.update(range(base, base + 100))
            continue
        if "-" in part:
            lo_s, hi_s = (x.strip() for x in part.split("-", 1))
            lo = int(lo_s)
            hi = int(hi_s)
            if lo > hi:
                lo, hi = hi, lo
            allowed.update(range(lo, hi + 1))
            continue
        allowed.add(int(part))

    if not allowed:
        raise ValueError("--expect-status produced empty set")
    return allowed


# -----------------------------
# Check implementations
# -----------------------------


def check_tcp(host: str, port: int, timeout: float) -> Tuple[bool, str]:
    """Check TCP port is open."""
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return (True, "tcp connect ok")
    except socket.timeout:
        return (False, "tcp connect timeout")
    except ConnectionRefusedError:
        return (False, "tcp connection refused")
    except OSError as exc:
        return (False, f"tcp error: {exc}")


def _http_request(
    url: str,
    *,
    method: str,
    timeout: float,
    user_agent: str,
) -> urllib.response.addinfourl:
    req = urllib.request.Request(url, method=method)
    req.add_header("User-Agent", user_agent)
    return urllib.request.urlopen(req, timeout=timeout)  # nosec B310


def check_http(
    url: str,
    *,
    method: str,
    timeout: float,
    allowed_status: Optional[Set[int]],
    body_regex: Optional[re.Pattern[str]],
    user_agent: str,
) -> Tuple[bool, str, Optional[int]]:
    """Check HTTP endpoint matches expected status and optional body regex."""

    # If a body regex is provided, we must use GET.
    if body_regex is not None:
        method = "GET"

    try_methods: List[str]
    if method.upper() == "HEAD":
        # Prefer HEAD, fallback to GET if HEAD isn't supported.
        try_methods = ["HEAD", "GET"]
    else:
        try_methods = [method.upper()]

    last_err: Optional[str] = None

    for m in try_methods:
        try:
            with _http_request(url, method=m, timeout=timeout, user_agent=user_agent) as resp:
                status = getattr(resp, "status", None)

                # Status validation
                if status is None:
                    return (False, f"http {m}: missing status", None)

                if allowed_status is not None and status not in allowed_status:
                    return (False, f"http {m}: unexpected status {status}", status)

                if body_regex is not None:
                    # Limit body read to avoid memory blowups in bad configs.
                    body = resp.read(1024 * 1024)  # 1 MiB cap
                    try:
                        text = body.decode("utf-8", errors="replace")
                    except Exception:
                        text = str(body)
                    if not body_regex.search(text):
                        return (False, f"http {m}: body regex did not match", status)

                return (True, f"http {m}: ok", status)

        except urllib.error.HTTPError as exc:
            # HTTPError is a valid response with status code; treat via status checks.
            status = getattr(exc, "code", None)
            if allowed_status is not None and status in allowed_status:
                return (True, f"http {m}: ok (error status allowed)", int(status) if status is not None else None)
            last_err = f"http {m}: HTTPError {status}"

            # If HEAD is rejected, allow retry with GET.
            if m == "HEAD" and status in (400, 404, 405, 501):
                continue

            return (False, last_err, int(status) if status is not None else None)

        except urllib.error.URLError as exc:
            last_err = f"http {m}: URLError {exc.reason}"
            continue
        except socket.timeout:
            last_err = f"http {m}: timeout"
            continue

    return (False, last_err or "http: request failed", None)


# -----------------------------
# Retry wrapper
# -----------------------------


def with_retries(fn, *, retries: int, base_sleep: float = 0.15):
    last = None
    for attempt in range(retries + 1):
        if attempt > 0:
            time.sleep(base_sleep * (1.6 ** (attempt - 1)))
        last = fn()
        ok = last[0]
        if ok:
            return last
    return last


# -----------------------------
# CLI
# -----------------------------


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(prog="healthcheck.py", add_help=True)
    p.add_argument("target", help="http(s) URL, tcp://host:port, or host:port")
    p.add_argument("--timeout", type=float, default=5.0, help="Timeout seconds (default: 5)")
    p.add_argument("--retries", type=int, default=1, help="Retry count (default: 1)")
    p.add_argument("--method", default="HEAD", choices=["HEAD", "GET"], help="HTTP method")
    p.add_argument(
        "--expect-status",
        default=None,
        help="Allowed HTTP statuses, e.g. '200,204,3xx,401-403' (default: 200)",
    )
    p.add_argument(
        "--expect-body-regex",
        default=None,
        help="Regex to match in response body (forces GET).",
    )
    p.add_argument(
        "--user-agent",
        default="ODS-Healthcheck/1.0",
        help="User-Agent header",
    )
    p.add_argument("--json", action="store_true", help="Emit JSON result")
    return p.parse_args(argv)


def main(argv: Sequence[str]) -> int:
    try:
        args = parse_args(argv)
    except SystemExit as exc:
        if isinstance(exc.code, int):
            return exc.code
        return 1

    try:
        kind, norm = _parse_target(args.target)
    except ValueError as exc:
        res = Result(ok=False, target=args.target, kind="unknown", detail=str(exc))
        if args.json:
            print(res.to_json())
        else:
            print(f"[FAIL] {res.detail}")
        return 2

    if args.timeout <= 0:
        res = Result(ok=False, target=args.target, kind=kind, detail="--timeout must be > 0")
        if args.json:
            print(res.to_json())
        else:
            print("[FAIL] --timeout must be > 0")
        return 2

    if args.retries < 0 or args.retries > 50:
        res = Result(ok=False, target=args.target, kind=kind, detail="--retries out of range (0-50)")
        if args.json:
            print(res.to_json())
        else:
            print("[FAIL] --retries out of range")
        return 2

    allowed_status: Optional[Set[int]]
    if kind == "http":
        if args.expect_status is None:
            allowed_status = {200}
        else:
            try:
                allowed_status = _parse_expected_status(args.expect_status)
            except Exception as exc:
                res = Result(ok=False, target=args.target, kind=kind, detail=f"invalid --expect-status: {exc}")
                if args.json:
                    print(res.to_json())
                else:
                    print(f"[FAIL] {res.detail}")
                return 2
    else:
        allowed_status = None

    body_re: Optional[re.Pattern[str]] = None
    if args.expect_body_regex:
        try:
            body_re = re.compile(args.expect_body_regex)
        except re.error as exc:
            res = Result(ok=False, target=args.target, kind=kind, detail=f"invalid regex: {exc}")
            if args.json:
                print(res.to_json())
            else:
                print(f"[FAIL] {res.detail}")
            return 2

    start = time.perf_counter()

    if kind == "tcp":
        try:
            host, port = _parse_host_port(norm)
        except ValueError as exc:
            res = Result(ok=False, target=args.target, kind=kind, detail=str(exc))
            if args.json:
                print(res.to_json())
            else:
                print(f"[FAIL] {res.detail}")
            return 2

        ok, detail = with_retries(lambda: check_tcp(host, port, args.timeout), retries=args.retries)
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        res = Result(ok=bool(ok), target=args.target, kind=kind, detail=str(detail), elapsed_ms=elapsed_ms)

    else:
        ok, detail, status = with_retries(
            lambda: check_http(
                norm,
                method=args.method,
                timeout=args.timeout,
                allowed_status=allowed_status,
                body_regex=body_re,
                user_agent=args.user_agent,
            ),
            retries=args.retries,
        )
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        res = Result(
            ok=bool(ok),
            target=args.target,
            kind=kind,
            detail=str(detail),
            status=int(status) if status is not None else None,
            elapsed_ms=elapsed_ms,
        )

    if args.json:
        print(res.to_json())
    else:
        if res.ok:
            status_part = f" status={res.status}" if res.status is not None else ""
            print(f"[PASS] {res.kind} {res.target}{status_part} ({res.elapsed_ms}ms)")
        else:
            status_part = f" status={res.status}" if res.status is not None else ""
            print(f"[FAIL] {res.kind} {res.target}{status_part} ({res.elapsed_ms}ms): {res.detail}")

    return 0 if res.ok else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))