feat/add-basic-telemetry #10
@@ -5,8 +5,13 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import logging as lg
|
import logging as lg
|
||||||
import os
|
import os
|
||||||
|
import platform
|
||||||
|
import re
|
||||||
|
import subprocess
|
||||||
import uuid
|
import uuid
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import psutil
|
||||||
from judas_protocol import Message
|
from judas_protocol import Message
|
||||||
|
|
||||||
from judas_client.connector import Connector
|
from judas_client.connector import Connector
|
||||||
@@ -41,12 +46,18 @@ class Client:
|
|||||||
self.logger.debug(f"ID: {self.id}")
|
self.logger.debug(f"ID: {self.id}")
|
||||||
|
|
||||||
self.connector: Connector = Connector(
|
self.connector: Connector = Connector(
|
||||||
|
client=self,
|
||||||
mac_address=self.id,
|
mac_address=self.id,
|
||||||
host=self.server_host,
|
host=self.server_host,
|
||||||
port=self.server_port,
|
port=self.server_port,
|
||||||
on_message=self.handle_message,
|
on_message=self.handle_message,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
self.logger.debug("Gathering initial telemetry data...")
|
||||||
|
self.initial_telemetry: dict[str, Any] = (
|
||||||
|
self._gather_initial_telemetry()
|
||||||
|
)
|
||||||
|
|
||||||
def _get_mac_address(self) -> str:
|
def _get_mac_address(self) -> str:
|
||||||
"""Get the MAC address of the client.
|
"""Get the MAC address of the client.
|
||||||
|
|
||||||
@@ -61,6 +72,116 @@ class Client:
|
|||||||
)
|
)
|
||||||
return mac_address
|
return mac_address
|
||||||
|
|
||||||
|
def _get_cpu_model(self) -> str:
|
||||||
|
|
||||||
|
if platform.system() == "Windows":
|
||||||
|
return platform.processor()
|
||||||
|
elif platform.system() == "Darwin":
|
||||||
|
os.environ["PATH"] = os.environ["PATH"] + os.pathsep + "/usr/sbin"
|
||||||
|
command = "sysctl -n machdep.cpu.brand_string"
|
||||||
|
return str(subprocess.check_output(command).strip())
|
||||||
|
elif platform.system() == "Linux":
|
||||||
|
command = "cat /proc/cpuinfo"
|
||||||
|
all_info = (
|
||||||
|
subprocess.check_output(command, shell=True).decode().strip()
|
||||||
|
)
|
||||||
|
for line in all_info.split("\n"):
|
||||||
|
if "model name" in line:
|
||||||
|
return re.sub(".*model name.*:", "", line, 1).strip()
|
||||||
|
return "Unknown"
|
||||||
|
|
||||||
|
def _gather_initial_telemetry(self) -> dict[str, Any]:
|
||||||
|
"""Gather initial telemetry data."""
|
||||||
|
self.logger.debug("Gathering initial telemetry data...")
|
||||||
|
inventory: dict[str, Any] = {}
|
||||||
|
|
||||||
|
# platform info
|
||||||
|
inventory["platform"] = {
|
||||||
|
"system": platform.system(),
|
||||||
|
"node": platform.node(),
|
||||||
|
"release": platform.release(),
|
||||||
|
"version": platform.version(),
|
||||||
|
"architecture": platform.architecture()[0],
|
||||||
|
}
|
||||||
|
|
||||||
|
# CPU information
|
||||||
|
cpu_frequency = psutil.cpu_freq()
|
||||||
|
inventory["cpu"] = {
|
||||||
|
"model": self._get_cpu_model(),
|
||||||
|
"physical_cores": psutil.cpu_count(logical=False),
|
||||||
|
"threads": psutil.cpu_count(logical=True),
|
||||||
|
"max_frequency_mhz": cpu_frequency.max
|
||||||
|
if cpu_frequency
|
||||||
|
else "Unknown",
|
||||||
|
"min_frequency_mhz": cpu_frequency.min
|
||||||
|
if cpu_frequency
|
||||||
|
else "Unknown",
|
||||||
|
}
|
||||||
|
|
||||||
|
# Memory information
|
||||||
|
virtual_memory = psutil.virtual_memory()
|
||||||
|
swap_memory = psutil.swap_memory()
|
||||||
|
inventory["memory"] = {
|
||||||
|
"total_ram": round(virtual_memory.total / (1024**2), 2),
|
||||||
|
"total_swap": round(swap_memory.total / (1024**2), 2),
|
||||||
|
}
|
||||||
|
|
||||||
|
# Disk information
|
||||||
|
disk_partitions = psutil.disk_partitions()
|
||||||
|
inventory["disks"] = []
|
||||||
|
for partition in disk_partitions:
|
||||||
|
try:
|
||||||
|
usage = psutil.disk_usage(partition.mountpoint)
|
||||||
|
inventory["disks"].append(
|
||||||
|
{
|
||||||
|
"device": partition.device,
|
||||||
|
"mountpoint": partition.mountpoint,
|
||||||
|
"filesystem": partition.fstype,
|
||||||
|
"total_size_gb": round(usage.total / (1024**3), 2),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
except (PermissionError, OSError) as e:
|
||||||
|
self.logger.warning(
|
||||||
|
f"Permission denied for disk partition: {partition.device}\n{e}"
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Network information
|
||||||
|
network_interfaces = psutil.net_if_addrs()
|
||||||
|
inventory["network"] = []
|
||||||
|
for interface_name, addresses in network_interfaces.items():
|
||||||
|
interface_data = {
|
||||||
|
"interface": interface_name,
|
||||||
|
"ip_addresses": [],
|
||||||
|
"mac_address": None,
|
||||||
|
}
|
||||||
|
for address in addresses:
|
||||||
|
if address.family.name in ["AF_LINK", "AF_PACKET"]:
|
||||||
|
interface_data["mac_address"] = address.address
|
||||||
|
elif address.family.name == "AF_INET":
|
||||||
|
interface_data["ip_addresses"].append(address.address)
|
||||||
|
inventory["network"].append(interface_data)
|
||||||
|
|
||||||
|
# User information
|
||||||
|
inventory["users"] = []
|
||||||
|
for user in psutil.users():
|
||||||
|
inventory["users"].append(
|
||||||
|
{
|
||||||
|
"name": user.name,
|
||||||
|
"terminal": user.terminal,
|
||||||
|
"host": user.host,
|
||||||
|
"started": user.started,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Python environment information
|
||||||
|
inventory["python"] = {
|
||||||
|
"version": platform.python_version(),
|
||||||
|
"implementation": platform.python_implementation(),
|
||||||
|
}
|
||||||
|
|
||||||
|
return inventory
|
||||||
|
|
||||||
def handle_message(self, message: Message) -> None:
|
def handle_message(self, message: Message) -> None:
|
||||||
"""Handle incoming messages.
|
"""Handle incoming messages.
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user