tg-admin-bot/system_checks.py

import subprocess
import os


def _cmd(cmd: str) -> str:
    """
    Safe shell command runner
    """
    try:
        return subprocess.getoutput(cmd)
    except Exception as e:
        return f"ERROR: {e}"


# ---------- SECURITY ----------

def security() -> str:
    out = _cmd("sshd -T | grep -i '^permitrootlogin'")

    if not out or "ERROR:" in out:
        return "🔐 Security\n\n⚠️ permitrootlogin not found"

    lines = ["🔐 Security\n"]

    if "no" in out.lower():
        lines.append("🟢 Root login disabled")
    else:
        lines.append("🔴 Root login ENABLED")

    pass_auth = _cmd("sshd -T | grep -i '^passwordauthentication'")
    if pass_auth and "ERROR:" not in pass_auth:
        lines.append("🔴 Password auth enabled" if "yes" in pass_auth.lower() else "🟢 Password auth disabled")

    pubkey_auth = _cmd("sshd -T | grep -i '^pubkeyauthentication'")
    if pubkey_auth and "ERROR:" not in pubkey_auth:
        lines.append("🟢 Pubkey auth enabled" if "yes" in pubkey_auth.lower() else "🔴 Pubkey auth disabled")

    sec_updates = _cmd("apt list --upgradable 2>/dev/null | grep -i security | wc -l")
    if sec_updates and "ERROR:" not in sec_updates:
        try:
            count = int(sec_updates.strip())
            lines.append(f"🔔 Security updates: {count}")
        except ValueError:
            pass

    time_info = _cmd("timedatectl")
    if time_info and "ERROR:" not in time_info:
        tz = None
        ntp = None
        synced = None
        for line in time_info.splitlines():
            if "Time zone:" in line:
                tz = line.split("Time zone:", 1)[1].strip()
            if "NTP service:" in line:
                ntp = line.split("NTP service:", 1)[1].strip()
            if "System clock synchronized:" in line:
                synced = line.split("System clock synchronized:", 1)[1].strip()
        if tz:
            lines.append(f"🕒 Time zone: {tz}")
        if ntp:
            lines.append(f"🔧 NTP service: {ntp}")
        if synced:
            lines.append(f"⏱ Clock synced: {synced}")

    return "\n".join(lines)


# ---------- DISKS ----------

def list_disks() -> list[str]:
    out = _cmd("lsblk -dn -o NAME,TYPE")

    disks = []
    for line in out.splitlines():
        parts = line.split()
        if len(parts) != 2:
            continue
        name, typ = parts
        if typ == "disk":
            disks.append(f"/dev/{name}")

    return disks


def smart_health(dev: str) -> str:
    out = _cmd(f"smartctl -H {dev}")

    if not out or "ERROR:" in out:
        return "⚠️ ERROR"

    if "PASSED" in out:
        return "🟢 PASSED"
    if "FAILED" in out:
        return "🔴 FAILED"

    return "⚠️ UNKNOWN"


def disk_temperature(dev: str) -> str:
    out = _cmd(f"smartctl -A {dev}")

    if not out or "ERROR:" in out:
        return "n/a"

    # NVMe
    for line in out.splitlines():
        if "Temperature:" in line and "Celsius" in line:
            try:
                temp = int("".join(filter(str.isdigit, line)))
                return f"{temp}°C"
            except Exception:
                pass

    # SATA attributes
    for line in out.splitlines():
        if line.strip().startswith(("194", "190")):
            parts = line.split()
            for p in parts[::-1]:
                if p.isdigit():
                    return f"{p}°C"

    return "n/a"


def smart_last_test(dev: str) -> str:
    out = _cmd(f"smartctl -l selftest {dev}")
    if not out or "ERROR:" in out:
        return "n/a"

    for line in out.splitlines():
        if "No self-tests have been logged" in line:
            return "no tests"
        if line.lstrip().startswith("#"):
            return line.strip()

    return "n/a"


def disks() -> str:
    disks = list_disks()

    if not disks:
        return "💽 Disks\n\n❌ No disks found"

    lines = ["💽 Disks (SMART)\n"]

    for d in disks:
        health = smart_health(d)
        temp = disk_temperature(d)

        icon = "🟢"
        if temp != "n/a":
            t = int(temp.replace("°C", ""))
            if t > 50:
                icon = "🔴"
            elif t > 40:
                icon = "🟡"

        lines.append(f"{icon} {d} — {health}, 🌡 {temp}")

    return "\n".join(lines)


def hardware() -> str:
    cpu_model = "n/a"
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                if line.lower().startswith("model name"):
                    cpu_model = line.split(":", 1)[1].strip()
                    break
    except Exception:
        pass

    mem_total = "n/a"
    swap_total = "n/a"
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    mem_kb = int(line.split()[1])
                    mem_total = f"{mem_kb / (1024**2):.2f} GiB"
                if line.startswith("SwapTotal:"):
                    swap_kb = int(line.split()[1])
                    swap_total = f"{swap_kb / (1024**2):.2f} GiB"
    except Exception:
        pass

    cores = os.cpu_count() or "n/a"
    uname = os.uname()

    lines = [
        "🧱 Hardware",
        "",
        f"🧠 CPU: {cpu_model}",
        f"🧩 Cores: {cores}",
        f"💾 RAM: {mem_total}",
        f"🌀 Swap: {swap_total}",
        f"🧬 Arch: {uname.machine}",
        f"🐧 Kernel: {uname.release}",
    ]
    gpu_lines = _gpu_info()
    if gpu_lines:
        lines.append("")
        lines.extend(gpu_lines)
    return "\n".join(lines)


def _gpu_info() -> list[str]:
    # 1) NVIDIA: use nvidia-smi if available for model + memory
    smi = _cmd("nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader")
    if smi and "ERROR:" not in smi and "not found" not in smi.lower():
        lines = ["🎮 GPU (NVIDIA)"]
        for line in smi.splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 2:
                name = parts[0]
                mem = parts[1]
                drv = parts[2] if len(parts) > 2 else "n/a"
                lines.append(f"• {name} | {mem} | driver {drv}")
        return lines

    # 2) Generic: lspci (VGA/3D/Display)
    lspci = _cmd("lspci -mm | egrep -i 'vga|3d|display'")
    if lspci and "ERROR:" not in lspci and "not found" not in lspci.lower():
        lines = ["🎮 GPU"]
        for line in lspci.splitlines():
            # Format: "00:02.0" "VGA compatible controller" "Intel Corporation" "..."
            parts = [p.strip().strip('"') for p in line.split('"') if p.strip()]
            if len(parts) >= 4:
                vendor = parts[2]
                model = parts[3]
                lines.append(f"• {vendor} {model}")
            elif line.strip():
                lines.append(f"• {line.strip()}")
        # Try AMD VRAM from sysfs if present
        vram = _cmd("cat /sys/class/drm/card*/device/mem_info_vram_total 2>/dev/null | head -n 1")
        if vram and vram.strip().isdigit():
            bytes_val = int(vram.strip())
            lines.append(f"• VRAM: {bytes_val / (1024**3):.2f} GiB")
        return lines

    return []