Added macOS support for CpuRam() object.

scripts/wk/obj.py

@@ -3,6 +3,8 @@
 
 import logging
 import pathlib
+import platform
+import plistlib
 import re
 
 from collections import OrderedDict
@@ -13,6 +15,16 @@ from wk.std import bytes_to_string, color_string, string_to_bytes
 # STATIC VARIABLES
 KEY_NVME = 'nvme_smart_health_information_log'
 KEY_SMART = 'ata_smart_attributes'
+KNOWN_RAM_VENDOR_IDS = {
+  # https://github.com/hewigovens/hewigovens.github.com/wiki/Memory-vendor-code
+  '0x014F': 'Transcend',
+  '0x2C00': 'Micron',
+  '0x802C': 'Micron',
+  '0x80AD': 'Hynix',
+  '0x80CE': 'Samsung',
+  '0xAD00': 'Hynix',
+  '0xCE00': 'Samsung',
+  }
 LOG = logging.getLogger(__name__)
 REGEX_POWER_ON_TIME = re.compile(
   r'^(\d+)([Hh].*|\s+\(\d+\s+\d+\s+\d+\).*)'
@@ -22,57 +34,58 @@ REGEX_POWER_ON_TIME = re.compile(
 class CpuRam():
   """Object for tracking CPU & RAM specific data."""
   def __init__(self):
-    self.lscpu = {}
+    self.description = 'Unknown'
+    self.details = {}
+    self.ram_total = 'Unknown'
+    self.ram_dimms = []
     self.tests = OrderedDict()
+
+    # Update details
     self.get_cpu_details()
     self.get_ram_details()
-    self.name = self.lscpu.get('Model name', 'Unknown CPU')
-    self.description = self.name
 
   def get_cpu_details(self):
-    """Get CPU details from lscpu."""
+    """Get CPU details using OS specific methods."""
-    cmd = ['lscpu', '--json']
+    if platform.system() == 'Darwin':
-    json_data = get_json_from_command(cmd)
+      cmd = 'sysctl -n machdep.cpu.brand_string'.split()
-    for line in json_data.get('lscpu', [{}]):
+      proc = run_program(cmd, check=False)
-      _field = line.get('field', '').replace(':', '')
+      self.description = re.sub(r'\s+', ' ', proc.stdout.strip())
-      _data = line.get('data', '')
+    elif platform.system() == 'Linux':
-      if not (_field or _data):
+      cmd = ['lscpu', '--json']
-        # Skip
+      json_data = get_json_from_command(cmd)
-        continue
+      for line in json_data.get('lscpu', [{}]):
-      self.lscpu[_field] = _data
+        _field = line.get('field', '').replace(':', '')
+        _data = line.get('data', '')
+        if not (_field or _data):
+          # Skip
+          continue
+        self.details[_field] = _data
+
+      self.description = self.details.get('Model name', '')
+
+    # Replace empty description
+    if not self.description:
+      self.description = 'Unknown CPU'
 
   def get_ram_details(self):
-    """Get RAM details from dmidecode."""
+    """Get RAM details using OS specific methods."""
-    cmd = ['sudo', 'dmidecode', '--type', 'memory']
+    if platform.system() == 'Darwin':
-    manufacturer = 'UNKNOWN'
+      dimm_list = get_ram_list_macos()
+    elif platform.system() == 'Linux':
+      dimm_list = get_ram_list_linux()
+
     details = {'Total': 0}
-    size = 0
+    for dimm_details in dimm_list:
-
+      size, manufacturer = dimm_details
-    # Get DMI data
+      if size <= 0:
-    proc = run_program(cmd)
+        # Skip empty DIMMs
-    dmi_data = proc.stdout.splitlines()
+        continue
-
+      description = f'{bytes_to_string(size)} {manufacturer}'
-    # Parse data
+      details['Total'] += size
-    for line in dmi_data:
+      if description in details:
-      line = line.strip()
+        details[description] += 1
-      if line == 'Memory Device':
+      else:
-        # Reset vars
+        details[description] = 1
-        manufacturer = 'UNKNOWN'
-        size = 0
-      elif line.startswith('Size:'):
-        size = line.replace('Size: ', '')
-        size = string_to_bytes(size, assume_binary=True)
-      elif line.startswith('Manufacturer:'):
-        manufacturer = line.replace('Manufacturer: ', '')
-        if size <= 0:
-          # Skip non-populated slots
-          continue
-        description = f'{bytes_to_string(size)} {manufacturer}'
-        details['Total'] += size
-        if description in details:
-          details[description] += 1
-        else:
-          details[description] = 1
 
     # Save details
     self.ram_total = bytes_to_string(details.pop('Total', 0))
@@ -84,7 +97,7 @@ class CpuRam():
     """Generate CPU report with data from all tests."""
     report = []
     report.append(color_string('Device', 'BLUE'))
-    report.append(f'  {self.name}')
+    report.append(f'  {self.description}')
 
     # Include RAM details
     report.append(color_string('RAM', 'BLUE'))
@@ -101,7 +114,7 @@ class Disk():
   """Object for tracking disk specific data."""
   def __init__(self, path):
     self.attributes = {}
-    self.description = 'UNKNOWN'
+    self.description = 'Unknown'
     self.lsblk = {}
     self.nvme_smart_notes = {}
     self.path = pathlib.Path(path).resolve()
@@ -155,7 +168,7 @@ class Disk():
 
     # Set description
     self.description = '{size_str} ({tran}) {model} {serial}'.format(
-      size_str = bytes_to_string(self.lsblk['size'], use_binary=False),
+      size_str=bytes_to_string(self.lsblk['size'], use_binary=False),
       **self.lsblk,
       )
 
@@ -207,9 +220,9 @@ class Disk():
           # Ignoring invalid attribute
           LOG.error('Invalid SMART attribute: %s', attribute)
           continue
-        name = str(attribute.get('name', 'UNKNOWN')).replace('_', ' ').title()
+        name = str(attribute.get('name', 'Unknown')).replace('_', ' ').title()
         raw = int(attribute.get('raw', {}).get('value', -1))
-        raw_str = attribute.get('raw', {}).get('string', 'UNKNOWN')
+        raw_str = attribute.get('raw', {}).get('string', 'Unknown')
 
         # Fix power-on time
         match = REGEX_POWER_ON_TIME.match(raw_str)
@@ -221,5 +234,69 @@ class Disk():
           'name': name, 'raw': raw, 'raw_str': raw_str}
 
 
+# Functions
+def get_ram_list_linux():
+  """Get RAM list using dmidecode."""
+  cmd = ['sudo', 'dmidecode', '--type', 'memory']
+  dimm_list = []
+  manufacturer = 'Unknown'
+  size = 0
+
+  # Get DMI data
+  proc = run_program(cmd)
+  dmi_data = proc.stdout.splitlines()
+
+  # Parse data
+  for line in dmi_data:
+    line = line.strip()
+    if line == 'Memory Device':
+      # Reset vars
+      manufacturer = 'Unknown'
+      size = 0
+    elif line.startswith('Size:'):
+      size = line.replace('Size: ', '')
+      size = string_to_bytes(size, assume_binary=True)
+    elif line.startswith('Manufacturer:'):
+      manufacturer = line.replace('Manufacturer: ', '')
+      dimm_list.append([size, manufacturer])
+
+  # Save details
+  return dimm_list
+
+def get_ram_list_macos():
+  """Get RAM list using system_profiler."""
+  dimm_list = []
+
+  # Get and parse plist data
+  cmd = [
+    'system_profiler',
+    '-xml',
+    'SPMemoryDataType',
+    ]
+  proc = run_program(cmd, check=False, encoding=None, errors=None)
+  try:
+    plist_data = plistlib.loads(proc.stdout)
+  except (TypeError, ValueError):
+    # Ignore and return an empty list
+    return dimm_list
+
+  # Check DIMM data
+  dimm_details = plist_data[0].get('_items', [{}])[0].get('_items', [])
+  for dimm in dimm_details:
+    manufacturer = dimm.get('dimm_manufacturer', None)
+    manufacturer = KNOWN_RAM_VENDOR_IDS.get(manufacturer, 'Unknown')
+    size = dimm.get('dimm_size', '0 GB')
+    try:
+      size = string_to_bytes(size, assume_binary=True)
+    except ValueError:
+      # Empty DIMM?
+      LOG.error('Invalid DIMM size: %s', size)
+      continue
+    dimm_list.append([size, manufacturer])
+
+  # Save details
+  return dimm_list
+
+
 if __name__ == '__main__':
   print("This file is not meant to be called directly.")