From 6e557da370856ae9f62857cb7647aea06c10c14b Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Tue, 22 Oct 2019 18:49:44 -0700
Subject: [PATCH] Added CpuRAM() object.

---
 scripts/wk/__init__.py |  1 +
 scripts/wk/obj.py      | 91 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)
 create mode 100644 scripts/wk/obj.py

diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py
index 511cf8a7..589cc28b 100644
--- a/scripts/wk/__init__.py
+++ b/scripts/wk/__init__.py
@@ -10,6 +10,7 @@ from wk import io
 from wk import kit
 from wk import log
 from wk import net
+from wk import obj
 from wk import os
 from wk import std
 from wk import sw
diff --git a/scripts/wk/obj.py b/scripts/wk/obj.py
new file mode 100644
index 00000000..8507d6f0
--- /dev/null
+++ b/scripts/wk/obj.py
@@ -0,0 +1,91 @@
+"""WizardKit: Objects."""
+# vim: sts=2 sw=2 ts=2
+
+import pathlib
+
+from collections import OrderedDict
+
+from wk.exe import get_json_from_command, run_program
+from wk.std import bytes_to_string, color_string, string_to_bytes
+
+# Classes
+class CpuRam():
+  """Object for tracking CPU & RAM specific data."""
+  def __init__(self):
+    self.lscpu = {}
+    self.tests = OrderedDict()
+    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."""
+    cmd = ['lscpu', '--json']
+    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.lscpu[_field] = _data
+
+  def get_ram_details(self):
+    """Get RAM details from dmidecode."""
+    cmd = ['sudo', 'dmidecode', '--type', 'memory']
+    manufacturer = 'UNKNOWN'
+    details = {'Total': 0}
+    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: ', '')
+        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))
+    self.ram_dimms = [
+      f'{count}x {desc}' for desc, count in sorted(details.items())
+      ]
+
+  def generate_cpu_report(self):
+    """Generate CPU report with data from all tests."""
+    report = []
+    report.append(color_string('Device', 'BLUE'))
+    report.append(f'  {self.name}')
+
+    # 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.values():
+      report.extend(test.report)
+
+    return report
+
+
+if __name__ == '__main__':
+  print("This file is not meant to be called directly.")