tai/tai_pkg/sysinfo.py

"""Host inspection probes and report assembly."""

import os
import shutil
import subprocess
import sys
import threading
import time

from tai_pkg.consts import Consts
from tai_pkg.sandbox import Sandbox
from tai_pkg.text import Text

__all__ = ["SysInfo", "SYSINFO_CHECKS"]


class SysInfo:
    """SysInfo."""

    @staticmethod
    def tmux_pane_header():
        try:
            info = subprocess.run(["tmux", "display-message", "-p", "#{session_name}:#{window_index}.#{pane_index} #{pane_current_command}"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=10)
        except (OSError, subprocess.SubprocessError):
            return ""
        if info.returncode == 0 and info.stdout.strip():
            return info.stdout.strip()
        return ""

    @staticmethod
    def sysinfo_check_os():
        try:
            uname = os.uname()
            base = "%s %s %s" % (uname.sysname, uname.release, uname.machine)
        except (AttributeError, OSError):
            base = sys.platform
        try:
            with open("/etc/os-release", "r", encoding="utf-8") as handle:
                fields = dict(line.strip().split("=", 1) for line in handle if "=" in line)
            pretty = fields.get("PRETTY_NAME", "").strip('"')
            if pretty:
                return "%s (%s)" % (pretty, base)
        except OSError:
            pass
        return base

    @staticmethod
    def sysinfo_check_python():
        return "%s (%s)" % (sys.version.split()[0], sys.executable)

    @staticmethod
    def sysinfo_check_venv():
        env_path = os.environ.get("VIRTUAL_ENV") or ""
        if env_path:
            return "yes (%s)" % env_path
        if sys.prefix != sys.base_prefix:
            return "yes (%s)" % sys.prefix
        return "no (system python at %s)" % sys.prefix

    @staticmethod
    def sysinfo_check_root():
        geteuid = getattr(os, "geteuid", None)
        if geteuid is None:
            return "unknown (no geteuid on %s)" % sys.platform
        uid = geteuid()
        return "yes (uid 0)" if uid == 0 else "no (uid %d)" % uid

    @staticmethod
    def sysinfo_check_container():
        engine = Sandbox.container_engine()
        if not engine:
            return "no engine (install podman or docker)"
        return "%s, box %s" % (engine, Sandbox.box_state(engine))

    @staticmethod
    def sysinfo_check_binaries():
        found = []
        lost = []
        for name in Consts.SYSINFO_BINARIES:
            (found if shutil.which(name) else lost).append(name)
        return "present: %s; missing: %s" % (", ".join(found) or "none", ", ".join(lost) or "none")

    @staticmethod
    def sysinfo_check_cpu():
        count = os.cpu_count() or 0
        try:
            one, five, fifteen = os.getloadavg()
            return "%d cores, load %.2f/%.2f/%.2f" % (count, one, five, fifteen)
        except OSError:
            return "%d cores" % count

    @staticmethod
    def sysinfo_check_disk():
        usage = shutil.disk_usage(os.getcwd())
        return "%.1f GiB free of %.1f GiB at %s" % (usage.free / 2 ** 30, usage.total / 2 ** 30, os.getcwd())

    @staticmethod
    def collect_sysinfo(names=None):
        wanted = [pair for pair in SYSINFO_CHECKS if names is None or pair[0] in names]
        results = {}

        def target(name, func):
            started = time.time()
            try:
                value = func()
            except Exception as exc:
                value = "error: " + Text.short_error(exc)
            results[name] = (value, int((time.time() - started) * 1000))

        workers = [threading.Thread(target=target, args=(name, func), daemon=True) for name, func in wanted]
        for worker in workers:
            worker.start()
        deadline = time.time() + Consts.SYSINFO_TIMEOUT
        for worker in workers:
            worker.join(timeout=max(0, deadline - time.time()))
        lines = []
        for name, _func in wanted:
            if name in results:
                value, elapsed = results[name]
                lines.append("%s: %s (%dms)" % (name, value, elapsed))
            else:
                lines.append("%s: timed out after %ds" % (name, Consts.SYSINFO_TIMEOUT))
        return "\n".join(lines)


SYSINFO_CHECKS = (
    ("os", SysInfo.sysinfo_check_os),
    ("python", SysInfo.sysinfo_check_python),
    ("venv", SysInfo.sysinfo_check_venv),
    ("root", SysInfo.sysinfo_check_root),
    ("container", SysInfo.sysinfo_check_container),
    ("binaries", SysInfo.sysinfo_check_binaries),
    ("cpu", SysInfo.sysinfo_check_cpu),
    ("disk", SysInfo.sysinfo_check_disk),
)