From 47b49077da5142b9a82bc514dc10a11bc8357306 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Thu, 15 Apr 2021 21:16:02 -0600
Subject: [PATCH] Show CHKDSK progress in separate pane under ConEmu

---
 scripts/wk/exe.py    |  7 ++++++-
 scripts/wk/os/win.py | 34 ++++++++++++++++------------------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/scripts/wk/exe.py b/scripts/wk/exe.py
index 87b90935..82e9dff3 100644
--- a/scripts/wk/exe.py
+++ b/scripts/wk/exe.py
@@ -130,7 +130,7 @@ def get_json_from_command(cmd, check=True, encoding='utf-8', errors='ignore'):
   return json_data
 
 
-def get_procs(name, exact=True):
+def get_procs(name, exact=True, try_again=True):
   """Get process object(s) based on name, returns list of proc objects."""
   LOG.debug('name: %s, exact: %s', name, exact)
   processes = []
@@ -141,6 +141,11 @@ def get_procs(name, exact=True):
     if re.search(regex, proc.name(), re.IGNORECASE):
       processes.append(proc)
 
+  # Try again?
+  if not processes and try_again:
+    time.sleep(1)
+    processes = get_procs(name, exact, try_again=False)
+
   # Done
   return processes
 
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index 63100b7c..919a4406 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 run_program
+from wk.exe import get_procs, run_program
 from wk.io import non_clobber_path
 from wk.log import format_log_path
 from wk.std import GenericError, GenericWarning, sleep
@@ -23,6 +23,7 @@ from wk.std import GenericError, GenericWarning, sleep
 
 # STATIC VARIABLES
 LOG = logging.getLogger(__name__)
+CONEMU = 'ConEmuPID' in os.environ
 KNOWN_DATA_TYPES = {
   'BINARY': winreg.REG_BINARY,
   'DWORD': winreg.REG_DWORD,
@@ -316,8 +317,8 @@ def reg_set_value(hive, key, name, data, data_type, option=None):
 # Repair Functions
 def run_chkdsk_offline():
   """Set filesystem 'dirty bit' to force a CHKDSK during startup."""
-  cmd = f'fsutil dirty set {os.environ.get("SYSTEMDRIVE")}'
-  proc = run_program(cmd.split(), check=False)
+  cmd = ['fsutil', 'dirty', 'set', os.environ.get('SYSTEMDRIVE', 'C:')]
+  proc = run_program(cmd, check=False)
 
   # Check result
   if proc.returncode > 0:
@@ -325,37 +326,34 @@ def run_chkdsk_offline():
 
 
 def run_chkdsk_online():
-  """Run CHKDSK in a split window.
+  """Run CHKDSK.
 
   NOTE: If run on Windows 8+ online repairs are attempted.
   """
   cmd = ['CHKDSK', os.environ.get('SYSTEMDRIVE', 'C:')]
   if OS_VERSION >= 8:
     cmd.extend(['/scan', '/perf'])
-  log_path = format_log_path(log_name='CHKDSK', tool=True)
-  err_path = log_path.with_suffix('.err')
+  if CONEMU:
+    cmd.extend(['-new_console:n', '-new_console:s33V'])
   retried = False
 
   # Run scan
-  proc = run_program(cmd, check=False)
-  if proc.returncode > 1:
+  run_program(cmd, check=False)
+  proc = get_procs('chkdsk.exe')[0]
+  return_code = proc.wait()
+  if return_code > 1:
     # Try again
     retried = True
-    proc = run_program(cmd, check=False)
+    run_program(cmd, check=False)
+    proc = get_procs('chkdsk.exe')[0]
+    return_code = proc.wait()
 
   # Check result
-  if (proc.returncode == 0 and retried) or proc.returncode == 1:
+  if (return_code == 0 and retried) or return_code == 1:
     raise GenericWarning('Repaired (or manually aborted)')
-  if proc.returncode > 1:
+  if return_code > 1:
     raise GenericError('Issue(s) detected')
 
-  # Save output
-  os.makedirs(log_path.parent, exist_ok=True)
-  with open(log_path, 'a') as _f:
-    _f.write(proc.stdout)
-  with open(err_path, 'a') as _f:
-    _f.write(proc.stderr)
-
 
 def run_sfc_scan():
   """Run SFC and save results."""