From 9980dab27b37497dfb5fc64c7b07f6f394eef491 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 25 Jun 2023 02:21:26 -0700
Subject: [PATCH] Add initial winget support

---
 scripts/install_winget.ps1 | 37 +++++++++++++++++++++++++++++++++++++
 scripts/wk/os/win.py       | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 scripts/install_winget.ps1

diff --git a/scripts/install_winget.ps1 b/scripts/install_winget.ps1
new file mode 100644
index 00000000..72d45659
--- /dev/null
+++ b/scripts/install_winget.ps1
@@ -0,0 +1,37 @@
+# WizardKit: Install winget (if needed)
+
+#Requires -Version 3.0
+if (Test-Path Env:\DEBUG) {
+  Set-PSDebug -Trace 1
+}
+$Host.UI.RawUI.WindowTitle = "WizardKit: Winget installer"
+$Host.UI.RawUI.BackgroundColor = "black"
+$Host.UI.RawUI.ForegroundColor = "white"
+$ProgressPreference = "SilentlyContinue"
+
+# STATIC VARIABLES
+$EXIT_OK = 0
+$EXIT_INSTALLED = 1
+$EXIT_FAILED_TO_INSTALL = 2
+
+# Main
+$NeedsInstalled = $false
+try {
+  $_ = $(winget --version)
+}
+catch {
+  $NeedsInstalled = $true
+}
+
+# Install
+if (! $NeedsInstalled) {
+  exit $EXIT_INSTALLED
+  }
+try {
+  Add-AppxPackage -ErrorAction Stop -RegisterByFamilyName -MainPackage Microsoft.DesktopAppInstaller_8wekyb3d8bbwe
+}
+catch {
+  exit $EXIT_FAILED_TO_INSTALL
+}
+
+exit $EXIT_OK
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index d13e12fe..1e6ea48c 100644
--- a/scripts/wk/os/win.py
+++ b/scripts/wk/os/win.py
@@ -26,7 +26,7 @@ from wk.cfg.windows_builds import (
   OUTDATED_BUILD_NUMBERS,
   WINDOWS_BUILDS,
   )
-from wk.exe import get_json_from_command, run_program
+from wk.exe import get_json_from_command, run_program, wait_for_procs
 from wk.kit.tools import find_kit_dir
 from wk.std import (
   GenericError,
@@ -744,5 +744,39 @@ def stop_service(service_name) -> None:
     raise GenericError(f'Failed to stop service {service_name}')
 
 
+# Winget Functions
+def winget_check(raise_exceptions: bool = False) -> None:
+  """Check if winget is present, install if not."""
+  cmd = [
+    'powershell',
+    '-ExecutionPolicy', 'bypass',
+    '-File', find_kit_dir('Scripts').joinpath('install_winget.ps1'),
+  ]
+  proc = run_program(cmd, check=False)
+
+  # Raise exception if requested
+  if raise_exceptions:
+    if proc.returncode == 1:
+      raise GenericWarning('Already installed')
+    if proc.returncode == 2:
+      raise GenericError('Failed to install')
+
+
+def winget_upgrade() -> None:
+  """Upgrade all supported programs with winget, returns subprocess.Popen."""
+  cmd = ['winget', 'upgrade', '--all']
+
+  # Adjust if running inside ConEmu
+  tmp_file = fr'{os.environ.get("TMP")}\run_winget.cmd'
+  if CONEMU:
+    with open(tmp_file, 'w', encoding='utf-8') as _f:
+      _f.write('@echo off\n')
+      _f.write(" ".join(cmd))
+    cmd = ('cmd', '/c', tmp_file, '-new_console:n', '-new_console:s33V')
+  run_program(cmd, check=False, pipe=False)
+  sleep(1)
+  wait_for_procs('winget.exe')
+
+
 if __name__ == '__main__':
   print("This file is not meant to be called directly.")