From e088f705ba9324c6c817691831ebbc236b7d4e1a Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Thu, 15 Apr 2021 23:33:11 -0600
Subject: [PATCH] Add run_dism()

---
 scripts/wk/exe.py    |  4 ++--
 scripts/wk/os/win.py | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/scripts/wk/exe.py b/scripts/wk/exe.py
index 82e9dff3..5edeb041 100644
--- a/scripts/wk/exe.py
+++ b/scripts/wk/exe.py
@@ -247,10 +247,10 @@ def wait_for_procs(name, exact=True, timeout=None):
   """Wait for all process matching name."""
   LOG.debug('name: %s, exact: %s, timeout: %s', name, exact, timeout)
   target_procs = get_procs(name, exact=exact)
-  results = psutil.wait_procs(target_procs, timeout=timeout)
+  procs = psutil.wait_procs(target_procs, timeout=timeout)
 
   # Raise exception if necessary
-  if results[1]: # Alive processes
+  if procs[1]: # Alive processes
     raise psutil.TimeoutExpired(name=name, seconds=timeout)
 
 
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index 919a4406..415679e5 100644
--- a/scripts/wk/os/win.py
+++ b/scripts/wk/os/win.py
@@ -15,7 +15,7 @@ except ImportError as err:
     raise err
 
 from wk.borrowed import acpi
-from wk.exe import get_procs, run_program
+from wk.exe import get_procs, run_program, wait_for_procs
 from wk.io import non_clobber_path
 from wk.log import format_log_path
 from wk.std import GenericError, GenericWarning, sleep
@@ -355,6 +355,41 @@ def run_chkdsk_online():
     raise GenericError('Issue(s) detected')
 
 
+def run_dism(repair=True):
+  """Run DISM to either scan or repair component store health."""
+  conemu_args = ['-new_console:n', '-new_console:s33V'] if CONEMU else []
+
+  # Bail early
+  if OS_VERSION < 8:
+    raise GenericWarning('Unsupported OS')
+
+  # Run (repair) scan
+  log_path = format_log_path(
+    log_name=f'DISM_{"Restore" if repair else "Scan"}Health', tool=True,
+    )
+  cmd = [
+    'DISM', '/Online', '/Cleanup-Image',
+    '/RestoreHealth' if repair else '/ScanHealth',
+    f'/LogPath:{log_path}',
+    *conemu_args,
+    ]
+  run_program(cmd, check=False, pipe=False)
+  wait_for_procs('dism.exe')
+
+  # Run check health
+  log_path = format_log_path(log_name='DISM_CheckHealth.log', tool=True)
+  cmd = [
+    'DISM', '/Online', '/Cleanup-Image',
+    '/CheckHealth',
+    f'/LogPath:{log_path}',
+    ]
+  proc = run_program(cmd, check=False)
+
+  # Check for errors
+  if 'no component store corruption detected' not in proc.stdout.lower():
+    raise GenericError('Issue(s) detected')
+
+
 def run_sfc_scan():
   """Run SFC and save results."""
   cmd = ['sfc', '/scannow']