Fix some typing warnings

nodescraper/plugins/inband/dimm/dimm_collector.py

@@ -23,27 +23,37 @@
 # SOFTWARE.
 #
 ###############################################################################
-import re
-from typing import Optional
+import csv
+from typing import ClassVar, Optional
+
+from typing_extensions import override
 
 from nodescraper.base import InBandDataCollector
-from nodescraper.connection.inband import TextFileArtifact
+from nodescraper.connection.inband import CommandArtifact, TextFileArtifact
 from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily
 from nodescraper.models import TaskResult
 
 from .collector_args import DimmCollectorArgs
-from .dimmdata import DimmDataModel
+from .dimmdata import DimmDataModel, DimmInfo
 
 
 class DimmCollector(InBandDataCollector[DimmDataModel, DimmCollectorArgs]):
     """Collect data on installed DIMMs"""
 
-    DATA_MODEL = DimmDataModel
+    DATA_MODEL: type[DimmDataModel] = DimmDataModel
+
+    # Both platforms dump every field the firmware exposes and decode it here,
+    # rather than filtering on the SUT, so that a shell quirk on any given host
+    # cannot silently drop modules from the inventory.
+    CMD_WINDOWS: ClassVar[str] = "wmic memorychip get /format:csv"
+    CMD_DMIDECODE: ClassVar[str] = "dmidecode -q --type 17"
+    CMD_DMIDECODE_FULL: ClassVar[str] = "dmidecode"
 
-    CMD_WINDOWS = "wmic memorychip get Capacity"
-    CMD = """sh -c 'dmidecode -t 17 | tr -s " " | grep -v "Volatile\\|None\\|Module" | grep Size' 2>/dev/null"""
-    CMD_DMIDECODE_FULL = "dmidecode"
+    # Section title of a DMI type 17 record, which is all that identifies one
+    # under -q since that hides the handle lines.
+    MEMORY_DEVICE: ClassVar[str] = "Memory Device"
 
+    @override
     def collect_data(
         self,
         args: Optional[DimmCollectorArgs] = None,
@@ -52,103 +62,175 @@ def collect_data(
         if args is None:
             args = DimmCollectorArgs()
 
-        dimm_str = None
         if self.system_info.os_family == OSFamily.WINDOWS:
-            res = self._run_sut_cmd(self.CMD_WINDOWS)
-            if res.exit_code == 0:
-                capacities = {}
-                total = 0
-                for line in res.stdout.splitlines():
-                    value = line.strip()
-                    if value.isdigit():
-                        value = int(value)
-                        total += value
-                        if value not in capacities:
-                            capacities[value] = 1
-                        else:
-                            capacities[value] += 1
-                dimm_str = f"{total / 1024 / 1024:.2f}GB @ "
-                for capacity, count in capacities.items():
-                    dimm_str += f"{count} x {capacity / 1024 / 1024:.2f}GB "
+            records = self._collect_windows_records()
         else:
             if args.skip_sudo:
                 self.result.message = "Skipping sudo plugin"
                 self.result.status = ExecutionStatus.NOT_RAN
                 return self.result, None
 
-            # Collect full dmidecode output as artifact
-            dmidecode_full_res = self._run_sut_cmd(self.CMD_DMIDECODE_FULL, sudo=True)
-            if dmidecode_full_res.exit_code == 0 and dmidecode_full_res.stdout:
-                self.result.artifacts.append(
-                    TextFileArtifact(filename="dmidecode.txt", contents=dmidecode_full_res.stdout)
-                )
-            else:
-                self._log_event(
-                    category=EventCategory.OS,
-                    description="Could not collect full dmidecode output",
-                    data={
-                        "command": dmidecode_full_res.command,
-                        "exit_code": dmidecode_full_res.exit_code,
-                        "stderr": dmidecode_full_res.stderr,
-                    },
-                    priority=EventPriority.WARNING,
-                )
-
-            res = self._run_sut_cmd(self.CMD, sudo=True)
-            if res.exit_code == 0:
-                total = 0
-                topology = {}
-                size = ""
-                dimm_size_pattern = re.compile(r"Size:\s+(\d+)\s+([A-Za-z]+)")
-                matches = dimm_size_pattern.findall(res.stdout)
-                if matches:
-                    for match in matches:
-                        size = match[1]
-                        total += int(match[0])
-                        key = match[0] + match[1]
-                        if not topology.get(key, None):
-                            topology[key] = 1
-                        else:
-                            topology[key] += 1
-                topology["total"] = total
-                topology["size"] = size
-                total_gb = topology.pop("total")
-                size = topology.pop("size")
-                if total_gb == 0:
-                    dimm_str = "0 GB"
-                else:
-                    dimm_entries = [f"{v} x {k}" for k, v in topology.items()]
-                    dimm_str = f"{total_gb}{size} @ {' '.join(dimm_entries)}"
-        if res.exit_code != 0:
-            self._log_event(
-                category=EventCategory.OS,
-                description="Error checking dimms",
-                data={
-                    "command": res.command,
-                    "exit_code": res.exit_code,
-                    "stderr": res.stderr,
-                },
-                priority=EventPriority.ERROR,
-                console_log=True,
-            )
+            records = self._collect_dmidecode_records()
 
-        if dimm_str:
-            dimm_data = DimmDataModel(dimms=dimm_str)
-            self._log_event(
-                category=EventCategory.IO,
-                description="Installed DIMM check",
-                data=dimm_data.model_dump(),
-                priority=EventPriority.INFO,
-            )
-            self.result.message = f"DIMM: {dimm_str}"
-        else:
-            dimm_data = None
+        # Records for empty slots, and any the firmware reports too incompletely
+        # to decode, come back as None and are dropped from the inventory.
+        dimms = [dimm for dimm in map(DimmInfo.from_record, records) if dimm]
+
+        if not dimms:
             self._log_event(
                 category=EventCategory.IO,
                 description="DIMM info not found",
                 priority=EventPriority.CRITICAL,
             )
             self.result.message = "DIMM info not found"
             self.result.status = ExecutionStatus.ERROR
+            return self.result, None
+
+        dimm_data = DimmDataModel(dimms=dimms)
+        self.result.message = f"DIMM: {dimm_data}"
 
         return self.result, dimm_data
+
+    def _collect_dmidecode_records(self) -> list[dict[str, str]]:
+        """Dump the DMI tables and pull the memory device records out of them.
+
+        A full dump is collected first and kept as an artifact, then the type 17
+        records are read from a quiet, targeted dump. The full dump is used as a
+        fallback when the targeted one fails, since it carries the same records.
+
+        Returns:
+            list[dict[str, str]]: one raw record per memory device slot.
+        """
+        full_dump = self._run_dmidecode(self.CMD_DMIDECODE_FULL)
+        if full_dump:
+            self.result.artifacts.append(
+                TextFileArtifact(filename="dmidecode.txt", contents=full_dump)
+            )
+        else:
+            self._log_event(
+                category=EventCategory.OS,
+                description="Could not collect full dmidecode output",
+                priority=EventPriority.WARNING,
+            )
+
+        dump = self._run_dmidecode(self.CMD_DMIDECODE) or full_dump
+        if not dump:
+            return []
+
+        return self._parse_dmidecode(dump)
+
+    def _run_dmidecode(self, command: str) -> Optional[str]:
+        """Run a dmidecode command and return its output.
+
+        Args:
+            command (str): dmidecode command to run.
+
+        Returns:
+            Optional[str]: raw stdout, or None if the command produced none.
+        """
+        res = self._run_sut_cmd(command, sudo=True)
+
+        # Hosts that already run as root often ship no sudo binary at all, so a
+        # sudo specific failure is worth one retry without it.
+        if res.exit_code != 0 and "sudo" in str(res.stderr).lower():
+            res = self._run_sut_cmd(command, sudo=False)
+
+        if res.exit_code != 0 or not res.stdout:
+            self._log_cmd_error(res)
+            return None
+
+        return res.stdout
+
+    def _collect_windows_records(self) -> list[dict[str, str]]:
+        """Dump every Win32_PhysicalMemory property and split it into records.
+
+        Returns:
+            list[dict[str, str]]: one raw record per memory device slot.
+        """
+        res = self._run_sut_cmd(self.CMD_WINDOWS)
+        if res.exit_code != 0 or not res.stdout:
+            self._log_cmd_error(res)
+            return []
+
+        self.result.artifacts.append(
+            TextFileArtifact(filename="memorychip.csv", contents=res.stdout)
+        )
+
+        return self._parse_wmic_csv(res.stdout)
+
+    def _log_cmd_error(self, res: CommandArtifact) -> None:
+        """Log a failed memory query without aborting the rest of the collection.
+
+        Args:
+            res (CommandArtifact): result of the failed command.
+        """
+        self._log_event(
+            category=EventCategory.OS,
+            description="Error checking dimms",
+            priority=EventPriority.ERROR,
+            console_log=True,
+        )
+
+    @classmethod
+    def _parse_dmidecode(cls, dump: str) -> list[dict[str, str]]:
+        """Split a raw dmidecode dump into its memory device records.
+
+        Only DMI type 17 records are kept, and only their top level fields, so
+        that the "Volatile Size", "Cache Size" and "Logical Size" fields nested
+        under NVDIMM records are never mistaken for a module capacity.
+
+        Args:
+            dump (str): raw dmidecode output.
+
+        Returns:
+            list[dict[str, str]]: one raw record per memory device slot.
+        """
+        records: list[dict[str, str]] = []
+        fields: Optional[dict[str, str]] = None
+
+        for line in dump.splitlines():
+            field = line.strip()
+            if not field:
+                # A blank line closes the current record.
+                fields = None
+            elif field.startswith("Handle "):
+                # Without -q every record opens with a handle line naming its
+                # DMI type, which is the most precise way to spot type 17.
+                fields = None
+                if "DMI type 17," in field:
+                    fields = {}
+                    records.append(fields)
+            elif field == cls.MEMORY_DEVICE:
+                # Under -q the handle lines are hidden, so the section title is
+                # what opens the record.
+                if fields is None:
+                    fields = {}
+                    records.append(fields)
+            elif fields is not None:
+                key, separator, value = field.partition(":")
+                key = key.strip()
+                # Nested fields repeat names such as "Size", so the first value
+                # seen for a key wins.
+                if separator and key not in fields:
+                    fields[key] = value.strip()
+
+        return records
+
+    @staticmethod
+    def _parse_wmic_csv(stdout: str) -> list[dict[str, str]]:
+        """Split `wmic memorychip get /format:csv` output into records.
+
+        wmic emits a header row naming each property followed by one row per
+        module, so the header is used to key the fields and any property that a
+        given Windows build does not report is simply absent from the record.
+
+        Args:
+            stdout (str): raw wmic output.
+
+        Returns:
+            list[dict[str, str]]: one raw record per memory device slot.
+        """
+        # wmic pads its output with blank lines that confuse the csv reader.
+        rows = [line for line in stdout.splitlines() if line.strip()]
+
+        return list(csv.DictReader(rows))

nodescraper/plugins/inband/dimm/dimm_plugin.py

@@ -23,7 +23,10 @@
 # SOFTWARE.
 #
 ###############################################################################
+from typing import Optional
+
 from nodescraper.base import InBandDataPlugin
+from nodescraper.interfaces.dataplugin import CollectorArgsClasses, CollectorClasses
 
 from .collector_args import DimmCollectorArgs
 from .dimm_collector import DimmCollector
@@ -33,8 +36,8 @@
 class DimmPlugin(InBandDataPlugin[DimmDataModel, DimmCollectorArgs, None]):
     """Plugin for collection and analysis of DIMM data"""
 
-    DATA_MODEL = DimmDataModel
+    DATA_MODEL: type[DimmDataModel] = DimmDataModel
 
-    COLLECTOR = DimmCollector
+    COLLECTOR: Optional[CollectorClasses] = DimmCollector
 
-    COLLECTOR_ARGS = DimmCollectorArgs
+    COLLECTOR_ARGS: Optional[CollectorArgsClasses] = DimmCollectorArgs