From be8b46b1d86ddeac946b121f22e96a5a3dea2007 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Sun, 7 Jan 2024 14:58:41 +0100
Subject: [PATCH] Add nvme flag and audio flag

---
 cmd/main.go          | 16 ++++++++++++++++
 pkg/params/params.go | 12 ++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/cmd/main.go b/cmd/main.go
index b41ae44..ded202c 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -48,6 +48,22 @@ func main() {
 		iommu.PrintOutput(output, pArg)
 		os.Exit(0)
 
+	} else if pArg.Flag["nvme"] {
+		// Get all Ethernet controllers
+		output := iommu.MatchSubclass(`Non-Volatile memory controller`, pArg)
+
+		// Print the output and exit
+		iommu.PrintOutput(output, pArg)
+		os.Exit(0)
+
+	} else if pArg.Flag["audio"] {
+		// Get all Ethernet controllers
+		output := iommu.MatchSubclass(`Audio device`, pArg)
+
+		// Print the output and exit
+		iommu.PrintOutput(output, pArg)
+		os.Exit(0)
+
 	} else if len(pArg.IntList["iommu_group"]) > 0 {
 		// Get all devices in specified IOMMU groups and append it to the output
 		output := iommu.GetDevicesFromGroups(pArg.IntList["iommu_group"], pArg.FlagCounter["related"], pArg)
diff --git a/pkg/params/params.go b/pkg/params/params.go
index 064ec8e..b7395f2 100644
--- a/pkg/params/params.go
+++ b/pkg/params/params.go
@@ -68,6 +68,16 @@ func NewParams() *Params {
 		Help:     "List all SATA controllers. (use -i # to only display results from specified IOMMU group)",
 	})
 
+	nvme := parser.Flag("", "nvme", &argparse.Options{
+		Required: false,
+		Help:     "List all NVM controllers. (use -i # to only display results from specified IOMMU group)",
+	})
+
+	audio := parser.Flag("", "audio", &argparse.Options{
+		Required: false,
+		Help:     "List all Audio devices. (use -i # to only display results from specified IOMMU group)",
+	})
+
 	iommu_group := parser.IntList("i", "group", &argparse.Options{
 		Required: false,
 		Help:     "List everything in the IOMMU groups given. Supply argument multiple times to list additional groups.",
@@ -139,6 +149,8 @@ func NewParams() *Params {
 	pArg.addFlag("usb", *usb)
 	pArg.addFlag("nic", *nic)
 	pArg.addFlag("sata", *sata)
+	pArg.addFlag("nvme", *nvme)
+	pArg.addFlag("audio", *audio)
 	pArg.addFlagCounter("related", *related)
 	pArg.addStringList("ignore", *ignore)
 	pArg.addIntList("iommu_group", *iommu_group)