From 12f47fadfaff4ec3a2b09dfdaaaf067e4e0b57b3 Mon Sep 17 00:00:00 2001
From: Greg Neagle <gregneagle@mac.com>
Date: Wed, 7 Aug 2024 09:31:32 -0700
Subject: [PATCH] Add support for stopRequested()

---
 code/cli/munki/shared/installer/pkg.swift    |  4 +++-
 code/cli/munki/shared/installer/rmpkgs.swift | 17 +++++++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/code/cli/munki/shared/installer/pkg.swift b/code/cli/munki/shared/installer/pkg.swift
index 1ae6cc62..5ad5f095 100644
--- a/code/cli/munki/shared/installer/pkg.swift
+++ b/code/cli/munki/shared/installer/pkg.swift
@@ -239,9 +239,11 @@ func installFromDirectory(_ directoryPath: String, options: PlistDict = [:]) asy
     // the directory.
     // Returns a tuple containing the exit code of the installer process and a boolean
     // indicating if a restart is needed
+    if stopRequested() {
+        return (0, false)
+    }
     if let items = try? FileManager.default.contentsOfDirectory(atPath: directoryPath) {
         for item in items {
-            // TODO: support user stop requests
             let itempath = (directoryPath as NSString).appendingPathComponent(item)
             if hasValidDiskImageExt(item) {
                 displayInfo("Mounting disk image \(item)")
diff --git a/code/cli/munki/shared/installer/rmpkgs.swift b/code/cli/munki/shared/installer/rmpkgs.swift
index 04af8982..2f37d489 100644
--- a/code/cli/munki/shared/installer/rmpkgs.swift
+++ b/code/cli/munki/shared/installer/rmpkgs.swift
@@ -269,7 +269,9 @@ func importFromPkgutil() async throws {
     var current = 0
     displayPercentDone(current: current, maximum: pkgCount)
     for pkg in pkglist {
-        // TODO: handle user cancellation requests
+        if stopRequested() {
+            throw UserCancelled()
+        }
         current += 1
         displayDetail("Importing \(pkg)...")
         try await getPkgDataAndAddtoDB(pkgid: pkg)
@@ -334,9 +336,7 @@ func getPathsToRemove(pkgKeys: [String]) throws -> [String] {
         SELECT path FROM paths WHERE (
             path_key IN (\(selectedPkgs)) AND path_key NOT IN (\(otherPkgs)))
     """
-    // TODO: move this output to a calling function. This is the wrong place for this
-    displayMinorStatus("Determining which filesystem items to remove")
-    // TODO: add munkistatus output
+
     let connection = try SQL3Connection(pkgDBPath())
     let query = try SQL3Statement(connection: connection, SQLString: combinedQuerySQL)
     while query.step() == SQL3Status.row {
@@ -561,6 +561,8 @@ func removePackages(
     }
     do {
         try await initReceiptDB(forcerebuild: rebuildPkgDB)
+    } catch let err as UserCancelled {
+        return -128
     } catch {
         displayError("Could not initialize receipt database: \(error)")
         return -3
@@ -575,10 +577,13 @@ func removePackages(
         displayError("Error retreiving pkg keys: \(error)")
         return -4
     }
-    // TODO: handle stop requests
-    // if stopRequested { return -128 }
+    if stopRequested() {
+        return -128
+    }
     var pathsToRemove = [String]()
     do {
+        displayMinorStatus("Determining which filesystem items to remove")
+        munkiStatusPercent(-1)
         pathsToRemove = try getPathsToRemove(pkgKeys: pkgKeys)
     } catch {
         displayError("Error getting paths to remove: \(error)")