From f219f686f6b534febd99e1a7dff63e7d23b32ad8 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Fri, 27 Oct 2023 08:54:46 +0200
Subject: [PATCH] Add selectUSB and fix error when pressing ESC

---
 internal/pages/02_select_gpu.go     | 22 ++++++--
 internal/pages/03_select_usbctrl.go | 82 +++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 4 deletions(-)
 create mode 100644 internal/pages/03_select_usbctrl.go

diff --git a/internal/pages/02_select_gpu.go b/internal/pages/02_select_gpu.go
index 5b05661..9bf77fc 100644
--- a/internal/pages/02_select_gpu.go
+++ b/internal/pages/02_select_gpu.go
@@ -2,6 +2,7 @@ package pages
 
 import (
 	"fmt"
+	"os"
 
 	lsiommu "github.com/HikariKnight/quickpassthrough/internal/lsiommu"
 	"github.com/HikariKnight/quickpassthrough/pkg/command"
@@ -19,11 +20,20 @@ func SelectGPU() {
 	// Generate a list of choices based on the GPUs and get the users selection
 	choice := menu.GenIOMMUMenu("Select a GPU to view the IOMMU groups of", gpus)
 
-	// View the selected GPU
-	ViewGPU(choice)
+	// Parse the choice
+	switch choice {
+	case "back":
+		Welcome()
+	case "":
+		// If ESC is pressed
+		fmt.Println("")
+		os.Exit(0)
+	default:
+		viewGPU(choice)
+	}
 }
 
-func ViewGPU(id string, ext ...int) {
+func viewGPU(id string, ext ...int) {
 	// Clear the screen
 	command.Clear()
 
@@ -63,10 +73,14 @@ func ViewGPU(id string, ext ...int) {
 	switch choice {
 	case "ext":
 		// Run an extended relative search
-		ViewGPU(id, 1)
+		viewGPU(id, 1)
 
 	case "n":
 		// Go back to selecting a gpu
 		SelectGPU()
+
+	case "y":
+		// Go to the select a usb controller
+		selectUSB()
 	}
 }
diff --git a/internal/pages/03_select_usbctrl.go b/internal/pages/03_select_usbctrl.go
new file mode 100644
index 0000000..555534d
--- /dev/null
+++ b/internal/pages/03_select_usbctrl.go
@@ -0,0 +1,82 @@
+package pages
+
+import (
+	"fmt"
+	"os"
+
+	lsiommu "github.com/HikariKnight/quickpassthrough/internal/lsiommu"
+	"github.com/HikariKnight/quickpassthrough/pkg/command"
+	"github.com/HikariKnight/quickpassthrough/pkg/menu"
+	"github.com/gookit/color"
+)
+
+func selectUSB() {
+	// Clear the screen
+	command.Clear()
+
+	// Get the users GPUs
+	usbs := lsiommu.GetIOMMU("-u", "-F", "vendor:,prod_name,optional_revision:,device_id")
+
+	// Generate a list of choices based on the GPUs and get the users selection
+	choice := menu.GenIOMMUMenu("Select a USB to view the IOMMU groups of", usbs, 1)
+
+	// Parse the choice
+	switch choice {
+	case "back":
+		SelectGPU()
+	case "":
+		// If ESC is pressed
+		fmt.Println("")
+		os.Exit(0)
+	default:
+		// View the selected GPU
+		viewUSB(choice)
+	}
+}
+
+func viewUSB(id string, ext ...int) {
+	// Clear the screen
+	command.Clear()
+
+	// Set mode to relative
+	mode := "-r"
+
+	// Set mode to relative extended
+	if len(ext) > 0 {
+		mode = "-rr"
+	}
+
+	// Get the IOMMU listings for USB controllers
+	group := lsiommu.GetIOMMU("-u", mode, "-i", id, "-F", "vendor:,prod_name,optional_revision:,device_id")
+
+	// Write a title
+	color.Bold.Println("This list should only show the USB controller")
+
+	// Print all the usb controllers
+	for _, v := range group {
+		fmt.Println(v)
+	}
+
+	// Add a new line for tidyness
+	fmt.Println("")
+
+	// Make an empty string
+	var choice string
+
+	// Ask if we shall use the devices for passthrough
+	if len(ext) == 0 {
+		choice = menu.YesNo("Use all listed devices for passthrough?")
+	} else {
+		choice = menu.YesNoEXT("Use all listed devices for passthrough?")
+	}
+
+	// Parse the choice
+	switch choice {
+	case "n":
+		// Go back to selecting a gpu
+		selectUSB()
+
+	case "y":
+		// Go to the select a usb controller
+	}
+}