WizardKit/scripts/wk/hw/system.py

"""WizardKit: System object and functions"""
# vim: sts=2 sw=2 ts=2

import logging
import plistlib
import re

from dataclasses import dataclass, field
from typing import Any

from wk.cfg.hw import KNOWN_RAM_VENDOR_IDS
from wk.exe import get_json_from_command, run_program
from wk.hw.test import Test
from wk.std import (
  PLATFORM,
  bytes_to_string,
  color_string,
  string_to_bytes,
  )


# STATIC VARIABLES
LOG = logging.getLogger(__name__)


@dataclass(slots=True)
class System:
  """Object for tracking system specific hardware data."""
  cpu_description: str =        field(init=False)
  ram_dimms: list[str] =        field(init=False, default_factory=list)
  ram_total: str =              field(init=False, default='Unknown')
  raw_details: dict[Any, Any] = field(init=False, default_factory=dict)
  tests: list[Test] =           field(init=False, default_factory=list)

  def __post_init__(self) -> None:
    self.get_cpu_details()
    self.set_cpu_description()
    self.get_ram_details()

  def generate_report(self) -> list[str]:
    """Generate CPU & RAM report, returns list."""
    report = []
    report.append(color_string('Device', 'BLUE'))
    report.append(f'  {self.cpu_description}')

    # Include RAM details
    report.append(color_string('RAM', 'BLUE'))
    report.append(f'  {self.ram_total} ({", ".join(self.ram_dimms)})')

    # Tests
    for test in self.tests:
      report.extend(test.report)

    return report

  def get_cpu_details(self) -> None:
    """Get CPU details using OS specific methods."""
    cmd = ['lscpu', '--json']

    # Bail early
    if PLATFORM != 'Linux':
      # Only Linux is supported ATM
      return

    # Parse details
    json_data = get_json_from_command(cmd)
    for line in json_data.get('lscpu', [{}]):
      _field = line.get('field', '').replace(':', '')
      _data = line.get('data', '')
      if not (_field or _data):
        # Skip
        continue
      self.raw_details[_field] = _data

  def get_ram_details(self) -> None:
    """Get RAM details using OS specific methods."""
    if PLATFORM == 'Darwin':
      dimm_list = get_ram_list_macos()
    elif PLATFORM == 'Linux':
      dimm_list = get_ram_list_linux()

    details = {'Total': 0}
    for dimm_details in dimm_list:
      size, manufacturer = dimm_details
      if size <= 0:
        # Skip empty DIMMs
        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))
    self.ram_dimms = [
      f'{count}x {desc}' for desc, count in sorted(details.items())
      ]

  def set_cpu_description(self) -> None:
    """Set CPU description."""
    self.cpu_description = self.raw_details.get('Model name', 'Unknown CPU')

    # macOS
    if PLATFORM == 'Darwin':
      cmd = 'sysctl -n machdep.cpu.brand_string'.split()
      proc = run_program(cmd, check=False)
      self.cpu_description = re.sub(r'\s+', ' ', proc.stdout.strip())


def get_ram_list_linux() -> list[list]:
  """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: ', '')
      try:
        size = string_to_bytes(size, assume_binary=True)
      except ValueError:
        # Assuming empty module
        size = 0
    elif line.startswith('Manufacturer:'):
      manufacturer = line.replace('Manufacturer: ', '')
      dimm_list.append([size, manufacturer])

  # Save details
  return dimm_list


def get_ram_list_macos() -> list[list]:
  """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,
      f'Unknown ({manufacturer})')
    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.")