Add delete encrypted files and fix some crashes

cmd/WiiUDownloader/mainwindow.go

@@ -22,16 +22,17 @@ const (
 )
 
 type MainWindow struct {
-	window           *gtk.Window
+	window                          *gtk.Window
-	treeView         *gtk.TreeView
+	treeView                        *gtk.TreeView
-	titles           []wiiudownloader.TitleEntry
+	titles                          []wiiudownloader.TitleEntry
-	searchEntry      *gtk.Entry
+	searchEntry                     *gtk.Entry
-	categoryButtons  []*gtk.ToggleButton
+	categoryButtons                 []*gtk.ToggleButton
-	titleQueue       []wiiudownloader.TitleEntry
+	titleQueue                      []wiiudownloader.TitleEntry
-	progressWindow   wiiudownloader.ProgressWindow
+	progressWindow                  wiiudownloader.ProgressWindow
-	addToQueueButton *gtk.Button
+	addToQueueButton                *gtk.Button
-	decryptContents  bool
+	decryptContents                 bool
-	currentRegion    uint8
+	currentRegion                   uint8
+	deleteEncryptedContentsCheckbox *gtk.CheckButton
 }
 
 func NewMainWindow(entries []wiiudownloader.TitleEntry) *MainWindow {
@@ -219,6 +220,12 @@ func (mw *MainWindow) ShowAll() {
 		log.Fatal("Unable to create button:", err)
 	}
 
+	mw.deleteEncryptedContentsCheckbox, err = gtk.CheckButtonNewWithLabel("Delete encrypted contents after decryption")
+	if err != nil {
+		log.Fatal("Unable to create button:", err)
+	}
+	mw.deleteEncryptedContentsCheckbox.SetSensitive(false)
+
 	mw.addToQueueButton.Connect("clicked", mw.onAddToQueueClicked)
 	downloadQueueButton.Connect("clicked", func() {
 		mw.progressWindow, err = wiiudownloader.CreateProgressWindow(mw.window)
@@ -231,7 +238,12 @@ func (mw *MainWindow) ShowAll() {
 	decryptContentsCheckbox.Connect("clicked", mw.onDecryptContentsClicked)
 	bottomhBox.PackStart(mw.addToQueueButton, false, false, 0)
 	bottomhBox.PackStart(downloadQueueButton, false, false, 0)
-	bottomhBox.PackStart(decryptContentsCheckbox, false, false, 0)
+
+	checkboxvBox, _ := gtk.BoxNew(gtk.ORIENTATION_VERTICAL, 0)
+	checkboxvBox.PackStart(decryptContentsCheckbox, false, false, 0)
+	checkboxvBox.PackEnd(mw.deleteEncryptedContentsCheckbox, false, false, 0)
+
+	bottomhBox.PackStart(checkboxvBox, false, false, 0)
 
 	japanButton, err := gtk.CheckButtonNewWithLabel("Japan")
 	japanButton.SetActive(true)
@@ -344,11 +356,22 @@ func (mw *MainWindow) onSelectionChanged() {
 func (mw *MainWindow) onDecryptContentsClicked() {
 	if mw.decryptContents {
 		mw.decryptContents = false
+		mw.deleteEncryptedContentsCheckbox.SetSensitive(false)
 	} else {
 		mw.decryptContents = true
+		mw.deleteEncryptedContentsCheckbox.SetSensitive(true)
 	}
 }
 
+func (mw *MainWindow) getDeleteEncryptedContents() bool {
+	if mw.deleteEncryptedContentsCheckbox.GetSensitive() {
+		if mw.deleteEncryptedContentsCheckbox.GetActive() {
+			return true
+		}
+	}
+	return false
+}
+
 func (mw *MainWindow) isTitleInQueue(title wiiudownloader.TitleEntry) bool {
 	for _, entry := range mw.titleQueue {
 		if entry.TitleID == title.TitleID {
@@ -446,7 +469,8 @@ func (mw *MainWindow) onDownloadQueueClicked() {
 			defer wg.Done()
 
 			tidStr := fmt.Sprintf("%016x", title.TitleID)
-			if err := wiiudownloader.DownloadTitle(tidStr, fmt.Sprintf("%s/%s [%s]", selectedPath, title.Name, tidStr), mw.decryptContents, progressWindow); err != nil {
+			titlePath := fmt.Sprintf("%s/%s [%s]", selectedPath, title.Name, tidStr)
+			if err := wiiudownloader.DownloadTitle(tidStr, titlePath, mw.decryptContents, progressWindow, mw.getDeleteEncryptedContents()); err != nil {
 				queueCancelled = true
 			}
 			mw.removeFromQueue(tidStr)

decryption.go

@@ -22,13 +22,13 @@ var (
 	decryptionError = false
 )
 
-func decryptContents(path string, progress *ProgressWindow) error {
+func decryptContents(path string, progress *ProgressWindow, deleteEncryptedContents bool) error {
 	wg.Add(1)
-	progressInt := 0
+	progressInt := 1
-	go runDecryption(path, &progressInt)
+	go runDecryption(path, &progressInt, deleteEncryptedContents)
+	progress.percentLabel.SetText("Decrypting...")
 	for !decryptionDone {
 		progress.bar.SetFraction(float64(progressInt) / 100)
-		progress.percentLabel.SetText(fmt.Sprintf("Decrypting... (%d%%)", progressInt))
 		time.Sleep(500 * time.Millisecond)
 	}
 
@@ -42,7 +42,7 @@ func decryptContents(path string, progress *ProgressWindow) error {
 	return nil
 }
 
-func runDecryption(path string, progress *int) {
+func runDecryption(path string, progress *int, deleteEncryptedContents bool) {
 	defer wg.Done()
 	argv := make([]*C.char, 2)
 	argv[0] = C.CString("WiiUDownloader")
@@ -50,5 +50,8 @@ func runDecryption(path string, progress *int) {
 	if int(C.cdecrypt_main(2, (**C.char)(unsafe.Pointer(&argv[0])), (*C.int)(unsafe.Pointer(progress)))) != 0 {
 		decryptionError = true
 	}
+	if deleteEncryptedContents {
+		doDeleteEncryptedContents(path)
+	}
 	decryptionDone = true
 }

downloader.go

@@ -93,7 +93,7 @@ func downloadFile(progressWindow *ProgressWindow, client *grab.Client, url strin
 
 	resp := client.Do(req)
 
-	progressWindow.label.SetText(path.Base(resp.Filename))
+	progressWindow.label.SetText(path.Base(outputPath))
 
 	go func(err *error) {
 		for !resp.IsComplete() {
@@ -125,7 +125,7 @@ func downloadFile(progressWindow *ProgressWindow, client *grab.Client, url strin
 	return nil
 }
 
-func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, progressWindow *ProgressWindow) error {
+func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, progressWindow *ProgressWindow, deleteEncryptedContents bool) error {
 	progressWindow.cancelButton.Connect("clicked", func() {
 		progressWindow.cancelled = true
 	})
@@ -164,7 +164,6 @@ func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, pr
 		if err != nil {
 			return err
 		}
-		fmt.Println(titleKey)
 		if err := generateTicket(tikPath, titleID, titleKey, titleVersion); err != nil {
 			return err
 		}
@@ -243,7 +242,7 @@ func DownloadTitle(titleID string, outputDirectory string, doDecryption bool, pr
 	}
 
 	if doDecryption && !progressWindow.cancelled {
-		if err := decryptContents(outputDir, progressWindow); err != nil {
+		if err := decryptContents(outputDir, progressWindow, deleteEncryptedContents); err != nil {
 			return err
 		}
 	}

utils.go

@@ -0,0 +1,27 @@
+package wiiudownloader
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+func isThisDecryptedFile(path string) bool {
+	return strings.Contains(path, "code") || strings.Contains(path, "content") || strings.Contains(path, "meta")
+}
+
+func doDeleteEncryptedContents(path string) error {
+	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if info.Mode().IsRegular() && !isThisDecryptedFile(filePath) {
+			if err := os.Remove(filePath); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+
+	return err
+}