lingniu-vehicle-ingest/tools/go_systemd_prod_smoke.py

#!/usr/bin/env python3
"""Systemd and port smoke checks for the Go native production ECS."""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from dataclasses import asdict, dataclass


DEFAULT_HOST = "115.29.187.205"
DEFAULT_USER = "root"
DEFAULT_SERVICES = [
    "lingniu-go-gateway.service",
    "lingniu-go-history-writer.service",
    "lingniu-go-stat-writer.service",
    "lingniu-go-realtime-api.service",
]
DEFAULT_PORTS = {
    808: "gateway",
    32960: "gateway",
    20210: "realtime-api",
}
DEFAULT_SPOOL_DIR = "/opt/lingniu-go-native/spool/gateway"


@dataclass(frozen=True)
class Check:
    name: str
    status: str
    message: str


@dataclass(frozen=True)
class ServiceState:
    active: str
    enabled: str


def parse_service_states(output: str) -> dict[str, ServiceState]:
    states: dict[str, ServiceState] = {}
    for line in output.splitlines():
        parts = line.split()
        if len(parts) >= 2:
            enabled = parts[2] if len(parts) >= 3 else "enabled"
            states[parts[0]] = ServiceState(parts[1], enabled)
    return states


def service_checks(output: str, services: list[str]) -> list[Check]:
    states = parse_service_states(output)
    checks: list[Check] = []
    for service in services:
        state = states.get(service, ServiceState("missing", "missing"))
        status = "pass" if state.active == "active" and state.enabled == "enabled" else "fail"
        checks.append(Check(
            "service." + service,
            status,
            "state=" + state.active + "; enabled=" + state.enabled,
        ))
    return checks


def port_checks(output: str, ports: dict[int, str]) -> list[Check]:
    lines = output.splitlines()
    checks: list[Check] = []
    for port, process in sorted(ports.items()):
        matching = [line.strip() for line in lines if has_port(line, port)]
        if not matching:
            checks.append(Check("port." + str(port), "fail", "not listening"))
            continue
        owned = [line for line in matching if process in line]
        if owned:
            checks.append(Check("port." + str(port), "pass", owned[0]))
        else:
            checks.append(Check("port." + str(port), "fail", matching[0]))
    return checks


def has_port(line: str, port: int) -> bool:
    needle = ":" + str(port)
    return needle + " " in line or needle + "\t" in line


def spool_check(output: str) -> Check:
    values = parse_key_values(output)
    files = int(values.get("files", "-1"))
    recent = int(values.get("recent", "-1"))
    status = "pass" if files == 0 and recent == 0 else "fail"
    return Check("spool.gateway", status, output.strip() or "missing")


def parse_key_values(output: str) -> dict[str, str]:
    values: dict[str, str] = {}
    for part in output.split():
        if "=" not in part:
            continue
        key, value = part.split("=", 1)
        values[key] = value
    return values


def ssh(host: str, user: str, command: str, timeout: float) -> str:
    target = user + "@" + host if user else host
    completed = subprocess.run(
        ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", target, command],
        check=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        timeout=timeout,
    )
    return completed.stdout


def remote_command(services: list[str]) -> str:
    quoted = " ".join(services)
    return (
        "for svc in " + quoted + "; do "
        "printf '%s ' \"$svc\"; "
        "printf '%s ' \"$(systemctl is-active \"$svc\" 2>/dev/null || true)\"; "
        "systemctl is-enabled \"$svc\" 2>/dev/null || true; "
        "done; "
        "printf '\\n--PORTS--\\n'; "
        "ss -lntp; "
        "printf '\\n--SPOOL--\\n'; "
        "spool='" + DEFAULT_SPOOL_DIR + "'; "
        "if test -d \"$spool\"; then "
        "files=$(find \"$spool\" -type f | wc -l); "
        "bytes=$(du -sb \"$spool\" 2>/dev/null | awk '{print $1}'); "
        "recent=$(find \"$spool\" -type f -mmin -5 | wc -l); "
        "printf 'files=%s bytes=%s recent=%s\\n' \"$files\" \"${bytes:-0}\" \"$recent\"; "
        "else printf 'files=-1 bytes=0 recent=-1 missing=%s\\n' \"$spool\"; fi"
    )


def split_remote_output(output: str) -> tuple[str, str, str]:
    ports_marker = "\n--PORTS--\n"
    spool_marker = "\n--SPOOL--\n"
    if ports_marker not in output:
        return output, "", ""
    service_output, rest = output.split(ports_marker, 1)
    if spool_marker not in rest:
        return service_output, rest, ""
    port_output, spool_output = rest.split(spool_marker, 1)
    return service_output, port_output, spool_output


def run(args: argparse.Namespace) -> tuple[str, list[Check]]:
    output = ssh(args.host, args.user, remote_command(DEFAULT_SERVICES), args.timeout)
    service_output, port_output, spool_output = split_remote_output(output)
    checks = (
        service_checks(service_output, DEFAULT_SERVICES)
        + port_checks(port_output, DEFAULT_PORTS)
        + [spool_check(spool_output)]
    )
    status = "fail" if any(check.status == "fail" for check in checks) else "pass"
    return status, checks


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--user", default=DEFAULT_USER)
    parser.add_argument("--timeout", type=float, default=20.0)
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    try:
        status, checks = run(args)
    except Exception as exc:
        status = "fail"
        checks = [Check("systemd.ssh", "fail", str(exc))]
    print(json.dumps({
        "status": status,
        "host": args.host,
        "checks": [asdict(check) for check in checks],
    }, ensure_ascii=False, indent=2))
    return 0 if status == "pass" else 1


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