From 1054794af32adc568cd5d0819c52fa8283cf8ad1 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Tue, 12 Nov 2019 19:56:39 -0700
Subject: [PATCH] Added get_disks()

* This calls either get_disks_linux() or get_disks_macos()
---
 scripts/wk/hw/diags.py | 85 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index 05b1c9fa..1ffb1162 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -6,6 +6,8 @@ import logging
 import os
 import pathlib
 import platform
+import plistlib
+import re
 import signal
 import time
 
@@ -13,6 +15,7 @@ from collections import OrderedDict
 from docopt import docopt
 
 from wk import cfg, exe, log, net, std, tmux
+from wk.hw import obj as hw_obj
 
 
 # atexit functions
@@ -63,6 +66,10 @@ MENU_SETS = {
 MENU_TOGGLES = (
   'Skip USB Benchmarks',
   )
+WK_LABEL_REGEX = re.compile(
+  fr'{cfg.main.KIT_NAME_SHORT}_(LINUX|UFD)',
+  re.IGNORECASE,
+  )
 
 
 # Classes
@@ -153,6 +160,12 @@ class State():
       )
     std.print_info('Starting Hardware Diagnostics')
 
+    # Add CPU
+    self.cpu = hw_obj.CpuRam()
+
+    # Add disks
+    self.disks = get_disks()
+
   def init_tmux(self):
     """Initialize tmux layout."""
     tmux.kill_all_panes()
@@ -286,6 +299,78 @@ def disk_surface_scan():
   std.print_warning('TODO: bb')
 
 
+def get_disks():
+  """Get disks using OS-specific methods, returns list."""
+  disks = []
+  if platform.system() == 'Darwin':
+    disks = get_disks_macos()
+  elif platform.system() == 'Linux':
+    disks = get_disks_linux()
+
+  # Done
+  return disks
+
+
+def get_disks_linux():
+  """Get disks via lsblk, returns list."""
+  cmd = ['lsblk', '--json', '--nodeps', '--paths']
+  disks = []
+
+  # Add valid disks
+  json_data = exe.get_json_from_command(cmd)
+  for disk in json_data.get('blockdevices', []):
+    disk_obj = hw_obj.Disk(disk['name'])
+    skip = False
+
+    # Skip loopback devices, optical devices, etc
+    if disk_obj.details['type'] != 'disk':
+      skip = True
+
+    # Skip WK disks
+    for label in disk_obj.get_labels():
+      if WK_LABEL_REGEX.search(label):
+        skip = True
+
+    # Add disk
+    if not skip:
+      disks.append(disk_obj)
+
+  # Done
+  return disks
+
+
+def get_disks_macos():
+  """Get disks via diskutil, returns list."""
+  cmd = ['diskutil', 'list', '-plist', 'physical']
+  disks = []
+
+  # Get info from diskutil
+  proc = exe.run_program(cmd, encoding=None, errors=None)
+  try:
+    plist_data = plistlib.loads(proc.stdout)
+  except (TypeError, ValueError):
+    # Invalid / corrupt plist data? return empty list to avoid crash
+    LOG.error('Failed to get diskutil list')
+    return disks
+
+  # Add valid disks
+  for disk in plist_data['WholeDisks']:
+    disk_obj = hw_obj.Disk(disk)
+    skip = False
+
+    # Skip WK disks
+    for label in disk_obj.get_labels():
+      if WK_LABEL_REGEX.search(label):
+        skip = True
+
+    # Add disk
+    if not skip:
+      disks.append(disk_obj)
+
+  # Done
+  return disks
+
+
 def keyboard_test():
   """Test keyboard using xev."""
   LOG.info('Keyboard Test (xev)')