From 30e43b9814f6947e7d9503e9f0b12e97e008099a Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Wed, 29 Sep 2021 19:56:49 -0600
Subject: [PATCH] Add storage status sections

---
 scripts/auto_setup.py     |  4 +--
 scripts/get_raw_disks.ps1 |  3 +++
 scripts/wk/os/win.py      | 55 +++++++++++++++++++++++++++++++++++++--
 scripts/wk/setup/win.py   | 19 ++++++++++++++
 4 files changed, 77 insertions(+), 4 deletions(-)
 create mode 100644 scripts/get_raw_disks.ps1

diff --git a/scripts/auto_setup.py b/scripts/auto_setup.py
index 7d2c1f42..3728c53d 100644
--- a/scripts/auto_setup.py
+++ b/scripts/auto_setup.py
@@ -138,10 +138,10 @@ BASE_MENUS = {
       ),
     'System Summary': (
       MenuEntry('Operating System',         'auto_show_os_name'),
-      MenuEntry('Windows Activation',       'auto_show_activation'),
+      MenuEntry('Windows Activation',       'auto_show_os_activation'),
       MenuEntry('Secure Boot',              'auto_show_secure_boot_status'),
       MenuEntry('Installed RAM',            'auto_show_installed_ram'),
-      MenuEntry('Storage Volumes',          no_op),
+      MenuEntry('Storage Status',           'auto_show_storage_status'),
       MenuEntry('Virus Protection',         no_op),
       MenuEntry('Partitions 4K Aligned',    no_op),
       ),
diff --git a/scripts/get_raw_disks.ps1 b/scripts/get_raw_disks.ps1
new file mode 100644
index 00000000..afdb9300
--- /dev/null
+++ b/scripts/get_raw_disks.ps1
@@ -0,0 +1,3 @@
+# Wizard Kit: Get RAW disks
+
+Get-Disk | Where-Object {$_.PartitionStyle -eq "RAW"} | Select FriendlyName,Size,PartitionStyle | ConvertTo-JSON
\ No newline at end of file
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index b5efeaa7..e3a7bc7c 100644
--- a/scripts/wk/os/win.py
+++ b/scripts/wk/os/win.py
@@ -23,8 +23,15 @@ from wk.cfg.windows_builds import (
   OUTDATED_BUILD_NUMBERS,
   WINDOWS_BUILDS,
   )
-from wk.exe import run_program
-from wk.std import GenericError, GenericWarning, bytes_to_string, sleep
+from wk.exe import get_json_from_command, run_program
+from wk.kit.tools import find_kit_dir
+from wk.std import (
+  GenericError,
+  GenericWarning,
+  bytes_to_string,
+  color_string,
+  sleep,
+  )
 
 
 # STATIC VARIABLES
@@ -204,6 +211,50 @@ def get_os_name(as_list=False, check=True):
   return [display_name] if as_list else display_name
 
 
+def get_raw_disks():
+  """Get all disks without a partiton table, returns list."""
+  script_path = find_kit_dir('Scripts').joinpath('get_raw_disks.ps1')
+  cmd = ['PowerShell', '-ExecutionPolicy', 'Bypass', '-File', script_path]
+  json_data = get_json_from_command(cmd)
+  raw_disks = []
+
+  # Fix JSON if only one disk was detected
+  if isinstance(json_data, dict):
+    json_data = [json_data]
+
+  # Parse JSON
+  for disk in json_data:
+    size_str = bytes_to_string(int(disk["Size"]), use_binary=False)
+    raw_disks.append(f'{disk["FriendlyName"]} ({size_str})')
+
+  # Done
+  return raw_disks
+
+
+def get_volume_usage(use_colors=False):
+  """Get space usage info for all fixed volumes, returns list."""
+  report = []
+  for disk in psutil.disk_partitions():
+    if 'fixed' not in disk.opts:
+      continue
+    total, _, free, percent = psutil.disk_usage(disk.device)
+    color = None
+    if percent > 85:
+      color = 'RED'
+    elif percent > 75:
+      color = 'YELLOW'
+    display_str = (
+      f'{free/total:>5.2f}% Free'
+      f'  ({bytes_to_string(free, 2):>10} / {bytes_to_string(total, 2):>10})'
+      )
+    if use_colors:
+      display_str = color_string(display_str, color)
+    report.append(f'{disk.device}  {display_str}')
+
+  # Done
+  return report
+
+
 def show_alert_box(message, title=None):
   """Show Windows alert box with message."""
   title = title if title else f'{KIT_NAME_FULL} Warning'
diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py
index acda54b1..a1c75c6d 100644
--- a/scripts/wk/setup/win.py
+++ b/scripts/wk/setup/win.py
@@ -47,6 +47,8 @@ if platform.system() == 'Windows':
     get_installed_ram,
     get_os_activation,
     get_os_name,
+    get_raw_disks,
+    get_volume_usage,
     is_secure_boot_enabled,
     reg_read_value,
     reg_set_value,
@@ -72,6 +74,8 @@ else:
   get_installed_ram = no_op
   get_os_activation = no_op
   get_os_name = no_op
+  get_raw_disks = no_op
+  get_volume_usage = no_op
   is_secure_boot_enabled = no_op
   reg_read_value = no_op
   reg_set_value = no_op
@@ -590,6 +594,11 @@ def auto_show_secure_boot_status():
     )
 
 
+def auto_show_storage_status():
+  """Display storage status."""
+  TRY_PRINT.run('Storage Status...', get_storage_status)
+
+
 def auto_windows_temp_fix():
   """Restore default ACLs for Windows\\Temp."""
   TRY_PRINT.run(r'Windows\Temp fix...', fix_windows_temp)
@@ -924,6 +933,16 @@ def get_firefox_default_profile(profiles_ini):
   return default_profile
 
 
+def get_storage_status():
+  """Get storage status for fixed disks, returns list."""
+  report = get_volume_usage(use_colors=True)
+  for disk in get_raw_disks():
+    report.append(color_string(f'Uninitialized Disk: {disk}', 'RED'))
+
+  # Done
+  return report
+
+
 # Tool Functions
 def export_aida64_report():
   """Export AIDA64 report."""