info-user/agent-deb/opt/info-user-agent/agent.py

import psutil
import socket
import platform
import time
import os
import getpass
import json
import requests
import logging
import subprocess
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)


def load_config():
    cfg_path = Path(__file__).with_suffix("").parent / "config.json"
    if cfg_path.exists():
        with open(cfg_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return {
        "server_url": "http://127.0.0.1",
        "agent_token": "agent-secret-token",
        "interval_minutes": 60,
        "send_public_ip": False,
    }


def get_local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(1)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"


def get_public_ip():
    try:
        r = requests.get("https://api.ipify.org", timeout=5)
        return r.text.strip()
    except Exception:
        return None


def get_uptime():
    return int(time.time() - psutil.boot_time())


def get_disk_info():
    disks = []
    for part in psutil.disk_partitions(all=False):
        try:
            usage = psutil.disk_usage(part.mountpoint)
            disks.append({
                "device": part.device,
                "mountpoint": part.mountpoint,
                "total_gb": round(usage.total / (2**30), 1),
                "used_gb": round(usage.used / (2**30), 1),
                "free_gb": round(usage.free / (2**30), 1),
                "percent": usage.percent,
            })
        except Exception:
            continue
    return disks


def get_active_sessions():
    """Return active user sessions similar to the `who` command.

    Each session is a dict with keys: username, tty, login_time, origin.
    Falls back to the current process user if `who` is unavailable.
    """
    try:
        result = subprocess.run(
            ["who", "-u"],
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
        if result.returncode != 0:
            return []
        sessions = []
        for line in result.stdout.strip().splitlines():
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) < 5:
                continue
            username = parts[0]
            tty = parts[1]
            # login time is usually "YYYY-MM-DD HH:MM" (parts[2] and parts[3])
            login_time = " ".join(parts[2:4]) if len(parts) >= 4 else ""
            # The last field may be a comment in parentheses, e.g. (192.168.255.50)
            origin = ""
            if parts[-1].startswith("(") and parts[-1].endswith(")"):
                origin = parts[-1][1:-1]
            sessions.append({
                "username": username,
                "tty": tty,
                "login_time": login_time,
                "origin": origin,
            })
        return sessions
    except Exception:
        return []


def get_display_username(sessions):
    """Return a short username representation for the dashboard.

    Single user  -> username
    Multiple     -> comma-separated usernames
    No sessions  -> fallback to process user
    """
    usernames = [s["username"] for s in sessions if s.get("username")]
    if not usernames:
        return getpass.getuser()
    return ", ".join(dict.fromkeys(usernames))


def get_cpu_temp():
    try:
        temps = psutil.sensors_temperatures()
        if temps:
            for name, entries in temps.items():
                for entry in entries:
                    if entry.current:
                        return round(entry.current, 1)
    except Exception:
        pass
    return None


def get_load_avg():
    try:
        return ", ".join(str(x) for x in os.getloadavg())
    except Exception:
        return None


def get_top_processes(limit=10):
    procs = []
    for p in psutil.process_iter(["pid", "name", "cpu_percent", "memory_percent"]):
        try:
            procs.append({
                "pid": p.info["pid"],
                "name": p.info["name"],
                "cpu_percent": round(p.info["cpu_percent"] or 0, 1),
                "ram_percent": round((p.info["memory_percent"] or 0), 1),
            })
        except Exception:
            continue
    procs.sort(key=lambda x: x["cpu_percent"], reverse=True)
    return procs[:limit]


def collect():
    cfg = load_config()
    sessions = get_active_sessions()
    payload = {
        "hostname": socket.gethostname(),
        "username": get_display_username(sessions),
        "sessions": sessions,
        "os_info": f"{platform.system()} {platform.release()}",
        "local_ip": get_local_ip(),
        "public_ip": get_public_ip() if cfg.get("send_public_ip") else None,
        "uptime_sec": get_uptime(),
        "cpu_percent": psutil.cpu_percent(interval=1),
        "ram_percent": psutil.virtual_memory().percent,
        "ram_total_gb": round(psutil.virtual_memory().total / (2**30), 1),
        "disk_info": get_disk_info(),
        "cpu_temp": get_cpu_temp(),
        "load_avg": get_load_avg(),
        "processes": get_top_processes(10),
    }
    return payload, cfg


def send(payload, cfg):
    url = cfg["server_url"].rstrip("/") + "/api/heartbeat"
    headers = {"Authorization": f"Bearer {cfg['agent_token']}", "Content-Type": "application/json"}
    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=15)
        resp.raise_for_status()
        logging.info("Heartbeat sent: %s", resp.status_code)
    except Exception as e:
        logging.error("Failed to send heartbeat: %s", e)


def main():
    payload, cfg = collect()
    send(payload, cfg)


if __name__ == "__main__":
    main()