From b6b4843d0c9ea53b9befeda12b79647f553e4957 Mon Sep 17 00:00:00 2001
From: Xpl0itU <24777100+Xpl0itU@users.noreply.github.com>
Date: Wed, 19 Jul 2023 11:53:40 +0200
Subject: [PATCH] Add cancel button

---
 cmd/WiiUDownloader/mainwindow.go | 38 +++++++++++++++++++++++++++++---
 downloader.go                    | 36 +++++++++++++++++++++++++-----
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/cmd/WiiUDownloader/mainwindow.go b/cmd/WiiUDownloader/mainwindow.go
index 3071571..3f26deb 100644
--- a/cmd/WiiUDownloader/mainwindow.go
+++ b/cmd/WiiUDownloader/mainwindow.go
@@ -42,7 +42,7 @@ func NewMainWindow(entries []wiiudownloader.TitleEntry) *MainWindow {
 	}
 
 	win.SetTitle("WiiUDownloaderGo")
-	win.SetDefaultSize(400, 300)
+	win.SetDefaultSize(716, 400)
 	win.Connect("destroy", func() {
 		gtk.MainQuit()
 	})
@@ -359,7 +359,32 @@ func (mw *MainWindow) onAddToQueueClicked() {
 	}
 }
 
+func (mw *MainWindow) updateTitlesInQueue() {
+	store, err := mw.treeView.GetModel()
+	if err != nil {
+		log.Fatal("Unable to get tree view model:", err)
+	}
+
+	storeRef := store.(*gtk.ListStore)
+
+	iter, _ := storeRef.GetIterFirst()
+	for iter != nil {
+		tid, _ := storeRef.GetValue(iter, TITLE_ID_COLUMN)
+		if tid != nil {
+			if tidStr, err := tid.GetString(); err == nil {
+				tidNum, _ := strconv.ParseUint(tidStr, 16, 64)
+				isInQueue := mw.isTitleInQueue(wiiudownloader.TitleEntry{TitleID: tidNum})
+				storeRef.SetValue(iter, IN_QUEUE_COLUMN, isInQueue)
+			}
+		}
+		if !storeRef.IterNext(iter) {
+			break
+		}
+	}
+}
+
 func (mw *MainWindow) onDownloadQueueClicked() {
+	queueCancelled := false
 	var wg sync.WaitGroup
 
 	selectedPath, err := dialog.Directory().Title("Select a path to save the games to").Browse()
@@ -374,14 +399,21 @@ func (mw *MainWindow) onDownloadQueueClicked() {
 			defer wg.Done()
 
 			tidStr := fmt.Sprintf("%016x", title.TitleID)
-			wiiudownloader.DownloadTitle(tidStr, fmt.Sprintf("%s/%s [%s]", selectedPath, title.Name, tidStr), mw.decryptContents, progressWindow)
+			if err := wiiudownloader.DownloadTitle(tidStr, fmt.Sprintf("%s/%s [%s]", selectedPath, title.Name, tidStr), mw.decryptContents, progressWindow); err != nil {
+				queueCancelled = true
+			}
 			mw.removeFromQueue(tidStr)
 		}(title, selectedPath, &mw.progressWindow)
 
+		if queueCancelled {
+			break
+		}
+
 		wg.Wait()
 	}
+	mw.titleQueue = []wiiudownloader.TitleEntry{} // Clear the queue
 	mw.progressWindow.Window.Close()
-	mw.onAddToQueueClicked()
+	mw.updateTitlesInQueue()
 }
 
 func Main() {
diff --git a/downloader.go b/downloader.go
index 515bd7e..3a2c2dc 100644
--- a/downloader.go
+++ b/downloader.go
@@ -21,6 +21,8 @@ type ProgressWindow struct {
 	label        *gtk.Label
 	bar          *gtk.ProgressBar
 	percentLabel *gtk.Label
+	cancelButton *gtk.Button
+	cancelled    bool
 }
 
 func CreateProgressWindow(parent *gtk.Window) (ProgressWindow, error) {
@@ -56,16 +58,33 @@ func CreateProgressWindow(parent *gtk.Window) (ProgressWindow, error) {
 	}
 	box.PackStart(percentLabel, false, false, 0)
 
+	cancelButton, err := gtk.ButtonNewWithLabel("Cancel")
+	if err != nil {
+		return ProgressWindow{}, err
+	}
+
+	bottomhBox, err := gtk.BoxNew(gtk.ORIENTATION_HORIZONTAL, 5)
+	if err != nil {
+		return ProgressWindow{}, err
+	}
+	bottomhBox.PackEnd(cancelButton, false, false, 0)
+	box.PackEnd(bottomhBox, false, false, 0)
+
 	return ProgressWindow{
 		Window:       win,
 		box:          box,
 		label:        filenameLabel,
 		bar:          progressBar,
 		percentLabel: percentLabel,
+		cancelButton: cancelButton,
+		cancelled:    false,
 	}, nil
 }
 
 func downloadFile(progressWindow *ProgressWindow, client *grab.Client, url string, outputPath string) error {
+	if progressWindow.cancelled {
+		return fmt.Errorf("cancelled")
+	}
 	done := false
 	req, err := grab.NewRequest(outputPath, url)
 	if err != nil {
@@ -82,13 +101,14 @@ func downloadFile(progressWindow *ProgressWindow, client *grab.Client, url strin
 				progressWindow.bar.SetFraction(resp.Progress())
 				progressWindow.percentLabel.SetText(fmt.Sprintf("%.0f%%", 100*resp.Progress()))
 			})
+
+			if progressWindow.cancelled {
+				resp.Cancel()
+				break
+			}
 		}
 
 		*err = resp.Err()
-
-		glib.IdleAdd(func() {
-			progressWindow.Window.SetTitle("Download Complete")
-		})
 		done = true
 	}(&err)
 
@@ -106,6 +126,9 @@ func downloadFile(progressWindow *ProgressWindow, client *grab.Client, url strin
 }
 
 func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, progressWindow *ProgressWindow) error {
+	progressWindow.cancelButton.Connect("clicked", func() {
+		progressWindow.cancelled = true
+	})
 	outputDir := strings.TrimRight(outputDirectory, "/\\")
 	baseURL := fmt.Sprintf("http://ccs.cdn.c.shop.nintendowifi.net/ccs/download/%s", titleID)
 	titleKeyBytes, err := hex.DecodeString(titleID)
@@ -214,9 +237,12 @@ func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, pr
 				return err
 			}
 		}
+		if progressWindow.cancelled {
+			break
+		}
 	}
 
-	if doDecryption {
+	if doDecryption && !progressWindow.cancelled {
 		if err := decryptContents(outputDir, progressWindow); err != nil {
 			return err
 		}