NexaFabric/backend/app/agent_assets/nexafabric-agent.py

#!/usr/bin/env python3
import argparse
import hashlib
import json
import os
import platform
import re
import select
import socket
import ssl
import struct
import subprocess
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


VERSION = "0.2.0"
VM_INTERFACE_RE = re.compile(r"(?:tap|fwbr|fwln|fwpr)(\d+)")
VM_INTERFACE_DETAIL_RE = re.compile(r"(?:tap|fwbr|fwln|fwpr)(\d+)i(\d+)")
ETH_P_IP = 0x0800
ETH_P_ALL = 0x0003
ETH_P_8021Q = 0x8100
ETH_P_8021AD = 0x88A8
IP_PROTOCOLS = {1: "icmp", 6: "tcp", 17: "udp"}


def read_text(path: str) -> str | None:
    try:
        return Path(path).read_text(encoding="utf-8").strip()
    except OSError:
        return None


def run_command(args: list[str], timeout: int = 5) -> tuple[int, str]:
    try:
        result = subprocess.run(args, capture_output=True, text=True, timeout=timeout, check=False)
        output = result.stdout.strip() or result.stderr.strip()
        return result.returncode, output
    except (OSError, subprocess.SubprocessError):
        return 127, ""


def collect_interfaces() -> list[dict[str, Any]]:
    interfaces = []
    for item in Path("/sys/class/net").iterdir() if Path("/sys/class/net").exists() else []:
        name = item.name
        if name == "lo":
            continue
        vm_match = VM_INTERFACE_RE.search(name)
        interfaces.append(
            {
                "name": name,
                "vmid": vm_match.group(1) if vm_match else None,
                "operstate": read_text(f"/sys/class/net/{name}/operstate") or "unknown",
                "rx_bytes": int(read_text(f"/sys/class/net/{name}/statistics/rx_bytes") or 0),
                "tx_bytes": int(read_text(f"/sys/class/net/{name}/statistics/tx_bytes") or 0),
                "rx_packets": int(read_text(f"/sys/class/net/{name}/statistics/rx_packets") or 0),
                "tx_packets": int(read_text(f"/sys/class/net/{name}/statistics/tx_packets") or 0),
            }
        )
    return interfaces


def interface_rank(name: str) -> int:
    if name.startswith("tap"):
        return 0
    if name.startswith("fwln"):
        return 1
    if name.startswith("fwpr"):
        return 2
    if name.startswith("fwbr"):
        return 3
    return 9


def collect_interface_traffic(interfaces: list[dict[str, Any]]) -> list[dict[str, Any]]:
    selected: dict[tuple[str, str], dict[str, Any]] = {}
    for interface in interfaces:
        vmid = interface.get("vmid")
        if not vmid:
            continue
        detail_match = VM_INTERFACE_DETAIL_RE.search(str(interface.get("name") or ""))
        nic = detail_match.group(2) if detail_match else "0"
        key = (str(vmid), nic)
        current = selected.get(key)
        if current and interface_rank(str(current.get("name") or "")) <= interface_rank(str(interface.get("name") or "")):
            continue
        selected[key] = interface

    traffic = []
    for (vmid, nic), interface in sorted(selected.items()):
        rx_bytes = int(interface.get("rx_bytes") or 0)
        tx_bytes = int(interface.get("tx_bytes") or 0)
        rx_packets = int(interface.get("rx_packets") or 0)
        tx_packets = int(interface.get("tx_packets") or 0)
        traffic.append(
            {
                "vmid": vmid,
                "nic": nic,
                "interface": interface.get("name"),
                "source": f"vm:{vmid}",
                "destination": "network",
                "protocol": "interface-counter",
                "rx_bytes": rx_bytes,
                "tx_bytes": tx_bytes,
                "bytes": rx_bytes + tx_bytes,
                "rx_packets": rx_packets,
                "tx_packets": tx_packets,
                "packets": rx_packets + tx_packets,
                "state": interface.get("operstate") or "unknown",
            }
        )
    return traffic


def selected_flow_interfaces(interfaces: list[dict[str, Any]]) -> list[dict[str, Any]]:
    selected: dict[tuple[str, str], dict[str, Any]] = {}
    for interface in interfaces:
        name = str(interface.get("name") or "")
        vmid = interface.get("vmid")
        if not vmid or not name.startswith(("tap", "fwln")):
            continue
        detail_match = VM_INTERFACE_DETAIL_RE.search(name)
        nic = detail_match.group(2) if detail_match else "0"
        key = (str(vmid), nic)
        current = selected.get(key)
        if current and interface_rank(str(current.get("name") or "")) <= interface_rank(name):
            continue
        selected[key] = interface
    return list(selected.values())


def ipv4_address(raw: bytes) -> str:
    return socket.inet_ntoa(raw)


def parse_packet_flow(packet: bytes) -> dict[str, Any] | None:
    if len(packet) < 34:
        return None
    offset = 12
    eth_type = struct.unpack("!H", packet[offset:offset + 2])[0]
    offset = 14
    while eth_type in {ETH_P_8021Q, ETH_P_8021AD}:
        if len(packet) < offset + 4:
            return None
        eth_type = struct.unpack("!H", packet[offset + 2:offset + 4])[0]
        offset += 4
    if eth_type != ETH_P_IP or len(packet) < offset + 20:
        return None

    version_ihl = packet[offset]
    version = version_ihl >> 4
    ihl = (version_ihl & 0x0F) * 4
    if version != 4 or ihl < 20 or len(packet) < offset + ihl:
        return None
    total_length = struct.unpack("!H", packet[offset + 2:offset + 4])[0]
    protocol_number = packet[offset + 9]
    protocol = IP_PROTOCOLS.get(protocol_number)
    if not protocol:
        return None
    source_ip = ipv4_address(packet[offset + 12:offset + 16])
    destination_ip = ipv4_address(packet[offset + 16:offset + 20])
    transport_offset = offset + ihl
    source_port = None
    destination_port = None
    if protocol in {"tcp", "udp"}:
        if len(packet) < transport_offset + 4:
            return None
        source_port, destination_port = struct.unpack("!HH", packet[transport_offset:transport_offset + 4])
    elif protocol == "icmp" and len(packet) >= transport_offset + 2:
        source_port = packet[transport_offset]
        destination_port = packet[transport_offset + 1]

    return {
        "source_ip": source_ip,
        "destination_ip": destination_ip,
        "protocol": protocol,
        "source_port": source_port,
        "destination_port": destination_port,
        "packets": 1,
        "bytes": total_length if total_length else max(len(packet) - offset, 0),
        "state": "observed",
    }


def collect_packet_flows(interfaces: list[dict[str, Any]], duration: int, limit: int = 500) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    flow_interfaces = selected_flow_interfaces(interfaces)
    sockets: dict[socket.socket, dict[str, Any]] = {}
    diagnostics: dict[str, Any] = {
        "collector": "linux-af-packet",
        "duration_seconds": duration,
        "interfaces_requested": [interface.get("name") for interface in flow_interfaces],
        "interfaces_opened": [],
        "errors": [],
    }
    for interface in flow_interfaces:
        name = str(interface.get("name") or "")
        try:
            packet_socket = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(ETH_P_ALL))
            packet_socket.bind((name, 0))
            packet_socket.setblocking(False)
            sockets[packet_socket] = interface
            diagnostics["interfaces_opened"].append(name)
        except OSError as exc:
            diagnostics["errors"].append({"interface": name, "error": str(exc)})

    if not sockets:
        return [], diagnostics

    flows: dict[tuple[object, ...], dict[str, Any]] = {}
    deadline = time.monotonic() + max(duration, 1)
    try:
        while time.monotonic() < deadline:
            timeout = min(1.0, max(deadline - time.monotonic(), 0.0))
            readable, _, _ = select.select(list(sockets), [], [], timeout)
            for packet_socket in readable:
                interface = sockets[packet_socket]
                try:
                    packet = packet_socket.recv(65535)
                except OSError:
                    continue
                flow = parse_packet_flow(packet)
                if not flow:
                    continue
                name = str(interface.get("name") or "")
                vmid = str(interface.get("vmid") or "")
                detail_match = VM_INTERFACE_DETAIL_RE.search(name)
                nic = detail_match.group(2) if detail_match else "0"
                key = (
                    vmid,
                    nic,
                    flow["source_ip"],
                    flow["destination_ip"],
                    flow["protocol"],
                    flow.get("source_port"),
                    flow.get("destination_port"),
                )
                current = flows.get(key)
                if current:
                    current["packets"] += 1
                    current["bytes"] += int(flow["bytes"])
                    continue
                flow.update(
                    {
                        "vmid": vmid,
                        "nic": nic,
                        "interface": name,
                        "collector": "linux-af-packet",
                    }
                )
                flows[key] = flow
                if len(flows) >= limit:
                    diagnostics["truncated"] = True
                    return sorted(flows.values(), key=lambda item: int(item.get("bytes") or 0), reverse=True), diagnostics
    finally:
        for packet_socket in sockets:
            packet_socket.close()

    return sorted(flows.values(), key=lambda item: int(item.get("bytes") or 0), reverse=True), diagnostics


def parse_conntrack_line(line: str) -> dict[str, Any] | None:
    parts = line.split()
    if len(parts) < 5 or parts[0] not in {"tcp", "udp", "icmp"}:
        return None
    protocol = parts[0]
    state = None
    if protocol == "tcp" and len(parts) > 3 and "=" not in parts[3]:
        state = parts[3]

    values: dict[str, list[str]] = {}
    for part in parts:
        if "=" not in part:
            continue
        key, value = part.split("=", 1)
        values.setdefault(key, []).append(value)

    src_values = values.get("src", [])
    dst_values = values.get("dst", [])
    if not src_values or not dst_values:
        return None

    packet_values = [int(value) for value in values.get("packets", []) if value.isdigit()]
    byte_values = [int(value) for value in values.get("bytes", []) if value.isdigit()]
    sport = values.get("sport", [None])[0]
    dport = values.get("dport", [None])[0]
    return {
        "source_ip": src_values[0],
        "destination_ip": dst_values[0],
        "protocol": protocol,
        "source_port": int(sport) if sport and sport.isdigit() else None,
        "destination_port": int(dport) if dport and dport.isdigit() else None,
        "packets": sum(packet_values),
        "bytes": sum(byte_values),
        "state": state,
    }


def collect_flows(limit: int = 500) -> list[dict[str, Any]]:
    code, output = run_command(["conntrack", "-L", "-o", "extended"], timeout=10)
    if code != 0 or not output:
        return []
    flows = []
    seen = set()
    for line in output.splitlines():
        flow = parse_conntrack_line(line)
        if not flow:
            continue
        key = (
            flow["source_ip"],
            flow["destination_ip"],
            flow["protocol"],
            flow.get("source_port"),
            flow.get("destination_port"),
        )
        if key in seen:
            continue
        seen.add(key)
        flows.append(flow)
        if len(flows) >= limit:
            break
    return flows


def collect_conntrack() -> dict[str, Any]:
    code, output = run_command(["conntrack", "-C"])
    if code == 0 and output.isdigit():
        return {"count": int(output), "source": "conntrack"}
    for path in ("/proc/net/nf_conntrack", "/proc/net/ip_conntrack"):
        try:
            with open(path, "r", encoding="utf-8", errors="ignore") as handle:
                return {"count": sum(1 for _ in handle), "source": path}
        except OSError:
            continue
    return {"count": None, "source": "unavailable"}


def collect_flow_diagnostics(conntrack: dict[str, Any], flow_count: int, packet_diagnostics: dict[str, Any] | None = None) -> dict[str, Any]:
    code, conntrack_path = run_command(["sh", "-c", "command -v conntrack"])
    return {
        "flow_count": flow_count,
        "packet_collector": packet_diagnostics,
        "conntrack_binary": conntrack_path if code == 0 else None,
        "conntrack_count": conntrack.get("count"),
        "conntrack_source": conntrack.get("source"),
        "bridge_nf_call_iptables": read_text("/proc/sys/net/bridge/bridge-nf-call-iptables"),
        "bridge_nf_call_ip6tables": read_text("/proc/sys/net/bridge/bridge-nf-call-ip6tables"),
        "nf_conntrack_max": read_text("/proc/sys/net/netfilter/nf_conntrack_max"),
        "nf_conntrack_count": read_text("/proc/sys/net/netfilter/nf_conntrack_count"),
    }


def collect_firewall() -> dict[str, Any]:
    status = {}
    code, output = run_command(["systemctl", "is-active", "pve-firewall"])
    status["pve_firewall"] = output if code == 0 else "unknown"
    code, output = run_command(["nft", "-j", "list", "ruleset"], timeout=10)
    if code == 0 and output:
        status["nft_ruleset_sha256"] = hashlib.sha256(output.encode("utf-8")).hexdigest()
    else:
        status["nft_ruleset_sha256"] = None
    return status


def collect_payload(config: dict[str, Any]) -> dict[str, Any]:
    uptime = read_text("/proc/uptime")
    interfaces = collect_interfaces()
    flow_limit = int(config.get("flow_limit", 500))
    packet_flows: list[dict[str, Any]] = []
    packet_diagnostics: dict[str, Any] | None = None
    if bool(config.get("packet_flow_collector", True)):
        packet_flows, packet_diagnostics = collect_packet_flows(
            interfaces,
            int(config.get("packet_flow_window_seconds", 10)),
            flow_limit,
        )
    flows = packet_flows or collect_flows(flow_limit)
    conntrack = collect_conntrack()
    return {
        "version": VERSION,
        "node_name": config.get("node_name"),
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "hostname": socket.gethostname(),
        "kernel": platform.release(),
        "uptime_seconds": float(uptime.split()[0]) if uptime else None,
        "loadavg": list(os.getloadavg()) if hasattr(os, "getloadavg") else [],
        "interfaces": interfaces,
        "interface_traffic": collect_interface_traffic(interfaces),
        "flows": flows,
        "conntrack": conntrack,
        "firewall": collect_firewall(),
        "extra": {"platform": platform.platform(), "flow_diagnostics": collect_flow_diagnostics(conntrack, len(flows), packet_diagnostics)},
    }


def normalized_api_url(config: dict[str, Any]) -> str:
    api_url = str(config["api_url"]).rstrip("/")
    while api_url.endswith("/api/v1/api/v1"):
        api_url = api_url.removesuffix("/api/v1")
    if not api_url.endswith("/api/v1"):
        api_url = f"{api_url}/api/v1"
    return api_url


def post_heartbeat(config: dict[str, Any], payload: dict[str, Any]) -> None:
    api_url = normalized_api_url(config)
    heartbeat_url = f"{api_url}/agents/heartbeat"
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        heartbeat_url,
        data=data,
        headers={
            "Authorization": f"Bearer {config['token']}",
            "Content-Type": "application/json",
            "User-Agent": f"nexafabric-agent/{VERSION}",
        },
        method="POST",
    )
    context = None
    if not bool(config.get("verify_tls", True)):
        context = ssl._create_unverified_context()
    with urllib.request.urlopen(request, timeout=15, context=context) as response:
        response.read()


def load_config(path: str) -> dict[str, Any]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def main() -> int:
    parser = argparse.ArgumentParser(description="NexaFabric Proxmox node telemetry agent")
    parser.add_argument("--config", default="/etc/nexafabric-agent/config.json")
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    config = load_config(args.config)
    interval = int(config.get("interval_seconds", 30))

    while True:
        started_at = time.monotonic()
        payload = collect_payload(config)
        try:
            post_heartbeat(config, payload)
            print(f"heartbeat ok: {payload['collected_at']}", flush=True)
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")[:500]
            print(f"heartbeat failed: HTTP {exc.code} {exc.reason} url={exc.url} body={body}", flush=True)
        except (OSError, urllib.error.URLError) as exc:
            print(f"heartbeat failed: {exc} api_url={normalized_api_url(config)}", flush=True)
        if args.once:
            return 0
        elapsed = time.monotonic() - started_at
        time.sleep(max(interval - elapsed, 1))


if __name__ == "__main__":
    raise SystemExit(main())