2626from typing import Optional
2727
2828from nodescraper .base import InBandDataCollector
29- from nodescraper .connection .inband .inband import CommandArtifact
29+ from nodescraper .connection .inband .inband import CommandArtifact , TextFileArtifact
3030from nodescraper .enums import EventCategory , EventPriority , ExecutionStatus , OSFamily
3131from nodescraper .models import TaskResult
3232
@@ -38,9 +38,10 @@ class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel,
3838
3939 DATA_MODEL = DeviceEnumerationDataModel
4040
41- CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'"
4241 CMD_GPU_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'VGA\\ |Display\\ |3D' | wc -l"
4342 CMD_VF_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'Virtual Function' | wc -l"
43+ CMD_LSCPU_LINUX = "lscpu"
44+ CMD_LSHW_LINUX = "lshw"
4445
4546 CMD_CPU_COUNT_WINDOWS = (
4647 'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"'
@@ -61,9 +62,8 @@ def _warning(
6162 description = description ,
6263 data = {
6364 "command" : command .command ,
64- "stdout" : command .stdout ,
65- "stderr" : command .stderr ,
6665 "exit_code" : command .exit_code ,
66+ "stderr" : command .stderr ,
6767 },
6868 priority = EventPriority .WARNING ,
6969 )
@@ -75,8 +75,7 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
7575 On Windows, use WMI and hyper-v cmdlets
7676 """
7777 if self .system_info .os_family == OSFamily .LINUX :
78- # Count CPU sockets
79- cpu_count_res = self ._run_sut_cmd (self .CMD_CPU_COUNT_LINUX )
78+ lscpu_res = self ._run_sut_cmd (self .CMD_LSCPU_LINUX , log_artifact = False )
8079
8180 # Count all AMD GPUs
8281 vendor_id = format (self .system_info .vendorid_ep , "x" )
@@ -86,17 +85,42 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
8685
8786 # Count AMD Virtual Functions
8887 vf_count_res = self ._run_sut_cmd (self .CMD_VF_COUNT_LINUX .format (vendorid_ep = vendor_id ))
88+
89+ # Collect lshw output
90+ lshw_res = self ._run_sut_cmd (self .CMD_LSHW_LINUX , sudo = True , log_artifact = False )
8991 else :
9092 cpu_count_res = self ._run_sut_cmd (self .CMD_CPU_COUNT_WINDOWS )
9193 gpu_count_res = self ._run_sut_cmd (self .CMD_GPU_COUNT_WINDOWS )
9294 vf_count_res = self ._run_sut_cmd (self .CMD_VF_COUNT_WINDOWS )
9395
9496 device_enum = DeviceEnumerationDataModel ()
9597
96- if cpu_count_res .exit_code == 0 :
97- device_enum .cpu_count = int (cpu_count_res .stdout )
98+ if self .system_info .os_family == OSFamily .LINUX :
99+ if lscpu_res .exit_code == 0 and lscpu_res .stdout :
100+ # Extract socket count from lscpu output
101+ for line in lscpu_res .stdout .splitlines ():
102+ if line .startswith ("Socket(s):" ):
103+ try :
104+ device_enum .cpu_count = int (line .split (":" )[1 ].strip ())
105+ break
106+ except (ValueError , IndexError ):
107+ self ._warning (
108+ description = "Cannot parse CPU count from lscpu output" ,
109+ command = lscpu_res ,
110+ )
111+ device_enum .lscpu_output = lscpu_res .stdout
112+ self ._log_event (
113+ category = EventCategory .PLATFORM ,
114+ description = "Collected lscpu output" ,
115+ priority = EventPriority .INFO ,
116+ )
117+ else :
118+ self ._warning (description = "Cannot collect lscpu output" , command = lscpu_res )
98119 else :
99- self ._warning (description = "Cannot determine CPU count" , command = cpu_count_res )
120+ if cpu_count_res .exit_code == 0 :
121+ device_enum .cpu_count = int (cpu_count_res .stdout )
122+ else :
123+ self ._warning (description = "Cannot determine CPU count" , command = cpu_count_res )
100124
101125 if gpu_count_res .exit_code == 0 :
102126 device_enum .gpu_count = int (gpu_count_res .stdout )
@@ -112,14 +136,33 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
112136 category = EventCategory .SW_DRIVER ,
113137 )
114138
139+ # Collect lshw output on Linux
140+ if self .system_info .os_family == OSFamily .LINUX :
141+ if lshw_res .exit_code == 0 and lshw_res .stdout :
142+ device_enum .lshw_output = lshw_res .stdout
143+ self .result .artifacts .append (
144+ TextFileArtifact (filename = "lshw.txt" , contents = lshw_res .stdout )
145+ )
146+ self ._log_event (
147+ category = EventCategory .PLATFORM ,
148+ description = "Collected lshw output" ,
149+ priority = EventPriority .INFO ,
150+ )
151+ else :
152+ self ._warning (description = "Cannot collect lshw output" , command = lshw_res )
153+
115154 if device_enum .cpu_count or device_enum .gpu_count or device_enum .vf_count :
155+ log_data = device_enum .model_dump (
156+ exclude_none = True ,
157+ exclude = {"lscpu_output" , "lshw_output" , "task_name" , "task_type" , "parent" },
158+ )
116159 self ._log_event (
117160 category = EventCategory .PLATFORM ,
118161 description = f"Counted { device_enum .cpu_count } { device_enum .gpu_count } { device_enum .vf_count } ,
119- data = device_enum . model_dump ( exclude_none = True ) ,
162+ data = log_data ,
120163 priority = EventPriority .INFO ,
121164 )
122- self .result .message = f"Device Enumeration: { device_enum . model_dump ( exclude_none = True ) }
165+ self .result .message = f"Device Enumeration: { log_data }
123166 self .result .status = ExecutionStatus .OK
124167 return self .result , device_enum
125168 else :
0 commit comments