From b8335188ce4b09caffd0e0d446c4d5213787fc19 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 17 Apr 2021 17:49:04 -0600
Subject: [PATCH] Add wk.kit.tools

For code related to downloading, finding, and running tools on the kit.
---
 scripts/wk/kit/__init__.py |  2 +
 scripts/wk/kit/tools.py    | 79 ++++++++++++++++++++++++++++++++++++++
 scripts/wk/repairs/win.py  | 25 +++++++-----
 3 files changed, 96 insertions(+), 10 deletions(-)
 create mode 100644 scripts/wk/kit/tools.py

diff --git a/scripts/wk/kit/__init__.py b/scripts/wk/kit/__init__.py
index 0869d9c6..4641fae7 100644
--- a/scripts/wk/kit/__init__.py
+++ b/scripts/wk/kit/__init__.py
@@ -3,5 +3,7 @@
 
 import platform
 
+from wk.kit import tools
+
 if platform.system() == 'Linux':
   from wk.kit import ufd
diff --git a/scripts/wk/kit/tools.py b/scripts/wk/kit/tools.py
new file mode 100644
index 00000000..9c7ce4ba
--- /dev/null
+++ b/scripts/wk/kit/tools.py
@@ -0,0 +1,79 @@
+"""WizardKit: Tool Functions"""
+# vim: sts=2 sw=2 ts=2
+
+import pathlib
+import sys
+
+from wk.exe import popen_program, run_program
+
+
+# STATIC VARIABLES
+ARCH = '64' if sys.maxsize > 2**32 else '32'
+
+
+# "GLOBAL" VARIABLES
+CACHED_DIRS = {}
+
+
+# Functions
+def find_kit_dir(name=None):
+  """Find folder in kit, returns pathlib.Path.
+
+  Search is performed in the script's path and then recursively upwards.
+  If name is given then search for that instead."""
+  cur_path = pathlib.Path(__file__).resolve().parent
+  search = name if name else '.bin'
+
+  # Search
+  if name in CACHED_DIRS:
+    return CACHED_DIRS[name]
+  while not cur_path.match(cur_path.anchor):
+    if cur_path.joinpath(search).exists():
+      break
+    cur_path = cur_path.parent
+
+  # Check
+  if cur_path.match(cur_path.anchor):
+    raise FileNotFoundError(f'Failed to find kit dir, {name=}')
+  if name:
+    cur_path = cur_path.joinpath(name)
+
+  # Done
+  CACHED_DIRS[name] = cur_path
+  return cur_path
+
+
+def get_tool_path(folder, name):
+  """Get tool path, returns pathlib.Path"""
+  bin_dir = find_kit_dir('.bin')
+
+  # "Search"
+  tool_path = bin_dir.joinpath(f'{folder}/{name}{ARCH}.exe')
+  if not tool_path.exists():
+    tool_path = tool_path.with_stem(name)
+
+  # Missing?
+  if not tool_path.exists():
+    raise FileNotFoundError(f'Failed to find tool, {folder=}, {name=}')
+
+  # Done
+  return tool_path
+
+
+def run_tool(folder, name, *args, popen=False):
+  """Run tool from kit."""
+  cmd = [get_tool_path(folder, name), *args]
+  proc = None
+
+  # Run
+  if popen:
+    proc = popen_program(cmd)
+  else:
+    proc = run_program(cmd, check=False)
+
+  # Done
+  return proc
+
+
+if __name__ == '__main__':
+  print("This file is not meant to be called directly.")
diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py
index 8abf080a..6cba5133 100644
--- a/scripts/wk/repairs/win.py
+++ b/scripts/wk/repairs/win.py
@@ -7,16 +7,21 @@ import os
 import platform
 import sys
 
-from wk.cfg.main import KIT_NAME_FULL
-from wk.exe import get_procs, run_program, popen_program, wait_for_procs
-from wk.log import format_log_path, update_log_path
-from wk.os.win import * # pylint: disable=wildcard-import
-from wk.std import (
-  GenericError, GenericWarning, TryAndPrint,
+from wk.cfg.main  import KIT_NAME_FULL
+from wk.exe       import get_procs, run_program, popen_program, wait_for_procs
+from wk.kit.tools import run_tool
+from wk.log       import format_log_path, update_log_path
+from wk.os.win    import reg_delete_value, reg_read_value, reg_set_value
+from wk.std       import (
+  GenericError,
+  GenericWarning,
+  TryAndPrint,
   abort,
-  clear_screen, print_info,
-  pause, sleep,
+  clear_screen,
+  pause,
+  print_info,
   set_title,
+  sleep,
   )
 
 
@@ -43,7 +48,7 @@ def end_session():
   run_program(cmd)
 
   # Disable Autologon
-  # TODO: Run Autologon
+  run_tool('Sysinternals', 'Autologon')
   reg_set_value(
     'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon',
     'AutoAdminLogon', '0', 'SZ',
@@ -111,7 +116,7 @@ def init_session():
   # One-time tasks
   # TODO: Backup Registry
   # TODO: Enable and create restore point
-  # TODO: Run Autologon
+  run_tool('Sysinternals', 'Autologon')
   # TODO: Disable Windows updates
   # TODO: Reset Windows updates
   reboot()