added lshw and lscpu cmds

Author: alexandraBara

nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py

@@ -26,7 +26,7 @@
 from typing import Optional
 
 from nodescraper.base import InBandDataCollector
-from nodescraper.connection.inband.inband import CommandArtifact
+from nodescraper.connection.inband.inband import CommandArtifact, TextFileArtifact
 from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily
 from nodescraper.models import TaskResult
 
@@ -41,6 +41,8 @@ class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel,
     CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'"
     CMD_GPU_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'VGA\\|Display\\|3D' | wc -l"
     CMD_VF_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'Virtual Function' | wc -l"
+    CMD_LSCPU_LINUX = "/usr/bin/lscpu"
+    CMD_LSHW_LINUX = "lshw"
 
     CMD_CPU_COUNT_WINDOWS = (
         'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"'
@@ -86,6 +88,12 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
 
             # Count AMD Virtual Functions
             vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX.format(vendorid_ep=vendor_id))
+
+            # Collect lscpu output
+            lscpu_res = self._run_sut_cmd(self.CMD_LSCPU_LINUX, log_artifact=False)
+
+            # Collect lshw output
+            lshw_res = self._run_sut_cmd(self.CMD_LSHW_LINUX, sudo=True, log_artifact=False)
         else:
             cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_WINDOWS)
             gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_WINDOWS)
@@ -112,14 +120,42 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
                 category=EventCategory.SW_DRIVER,
             )
 
+        if self.system_info.os_family == OSFamily.LINUX:
+            if lscpu_res.exit_code == 0 and lscpu_res.stdout:
+                device_enum.lscpu_output = lscpu_res.stdout
+                self._log_event(
+                    category=EventCategory.PLATFORM,
+                    description="Collected lscpu output",
+                    priority=EventPriority.INFO,
+                )
+            else:
+                self._warning(description="Cannot collect lscpu output", command=lscpu_res)
+
+            if lshw_res.exit_code == 0 and lshw_res.stdout:
+                device_enum.lshw_output = lshw_res.stdout
+                self.result.artifacts.append(
+                    TextFileArtifact(filename="lshw.txt", contents=lshw_res.stdout)
+                )
+                self._log_event(
+                    category=EventCategory.PLATFORM,
+                    description="Collected lshw output",
+                    priority=EventPriority.INFO,
+                )
+            else:
+                self._warning(description="Cannot collect lshw output", command=lshw_res)
+
         if device_enum.cpu_count or device_enum.gpu_count or device_enum.vf_count:
+            log_data = device_enum.model_dump(
+                exclude_none=True,
+                exclude={"lscpu_output", "lshw_output", "task_name", "task_type", "parent"},
+            )
             self._log_event(
                 category=EventCategory.PLATFORM,
                 description=f"Counted {device_enum.cpu_count} CPUs, {device_enum.gpu_count} GPUs, {device_enum.vf_count} VFs",
-                data=device_enum.model_dump(exclude_none=True),
+                data=log_data,
                 priority=EventPriority.INFO,
             )
-            self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}"
+            self.result.message = f"Device Enumeration: {log_data}"
             self.result.status = ExecutionStatus.OK
             return self.result, device_enum
         else:

nodescraper/plugins/inband/device_enumeration/deviceenumdata.py

@@ -32,3 +32,5 @@ class DeviceEnumerationDataModel(DataModel):
     cpu_count: Optional[int] = None
     gpu_count: Optional[int] = None
     vf_count: Optional[int] = None
+    lscpu_output: Optional[str] = None
+    lshw_output: Optional[str] = None

test/unit/plugin/test_device_enumeration_collector.py

@@ -51,6 +51,9 @@ def test_collect_linux(system_info, device_enumeration_collector):
     """Test linux typical output"""
     system_info.os_family = OSFamily.LINUX
 
+    lscpu_output = "Architecture:        x86_64\nCPU(s):              64\nSocket(s):           2"
+    lshw_output = "*-cpu\n  product: AMD EPYC 1234 64-Core Processor"
+
     device_enumeration_collector._run_sut_cmd = MagicMock(
         side_effect=[
             MagicMock(
@@ -71,12 +74,30 @@ def test_collect_linux(system_info, device_enumeration_collector):
                 stderr="",
                 command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l",
             ),
+            MagicMock(
+                exit_code=0,
+                stdout=lscpu_output,
+                stderr="",
+                command="/usr/bin/lscpu",
+            ),
+            MagicMock(
+                exit_code=0,
+                stdout=lshw_output,
+                stderr="",
+                command="lshw",
+            ),
         ]
     )
 
     result, data = device_enumeration_collector.collect_data()
     assert result.status == ExecutionStatus.OK
-    assert data == DeviceEnumerationDataModel(cpu_count=2, gpu_count=8, vf_count=0)
+    assert data == DeviceEnumerationDataModel(
+        cpu_count=2, gpu_count=8, vf_count=0, lscpu_output=lscpu_output, lshw_output=lshw_output
+    )
+    assert (
+        len([a for a in result.artifacts if hasattr(a, "filename") and a.filename == "lshw.txt"])
+        == 1
+    )
 
 
 def test_collect_windows(system_info, device_enumeration_collector):
@@ -135,6 +156,18 @@ def test_collect_error(system_info, device_enumeration_collector):
                 stderr="command failed",
                 command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l",
             ),
+            MagicMock(
+                exit_code=1,
+                stdout="",
+                stderr="command failed",
+                command="/usr/bin/lscpu",
+            ),
+            MagicMock(
+                exit_code=1,
+                stdout="",
+                stderr="command failed",
+                command="lshw",
+            ),
         ]
     )