diff --git a/src/judas_client/client.py b/src/judas_client/client.py index f7e41ea..d9c8bb7 100644 --- a/src/judas_client/client.py +++ b/src/judas_client/client.py @@ -5,8 +5,13 @@ from __future__ import annotations import logging as lg import os +import platform +import re +import subprocess import uuid +from typing import Any +import psutil from judas_protocol import Message from judas_client.connector import Connector @@ -41,12 +46,18 @@ class Client: self.logger.debug(f"ID: {self.id}") self.connector: Connector = Connector( + client=self, mac_address=self.id, host=self.server_host, port=self.server_port, 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: """Get the MAC address of the client. @@ -61,6 +72,116 @@ class Client: ) 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: """Handle incoming messages.