feat: toggle data/turn off detecting particular data

src/cli/flags.py

@@ -16,15 +16,18 @@ class FlagParser:
     """
 
     def __init__(self, logger, dm=None, offline=False):
-        args = sys.argv[1:]
-
+        self.args = sys.argv[1:]
         self.dm = dm
-        self.offline = offline or "--offline" in args
+        self.offline = offline or "--offline" in self.args
+        self.toggled_off = []
+
+        if "--off-data" in self.args:
+            self.off_data(self.args[self.args.index("--off-data") + 1])
 
-        if not self.dm and not list(filter(lambda x: "-h" in x.lower(), args)):
+        if not self.dm and not list(filter(lambda x: "-h" in x.lower(), self.args)):
             print(color_text(
                 "Analyzing hardware... (this might take a while, don't panic)", "red"))
-            self.dm = DeviceManager(logger, offline="--offline" in args)
+            self.dm = DeviceManager(logger, off_data=self.toggled_off, offline="--offline" in self.args)
             self.dm.info = {
                 k: v
                 for (k, v) in self.dm.info.items()
@@ -34,15 +37,15 @@ def __init__(self, logger, dm=None, offline=False):
         self.logger = logger
         self.completed = []
         self.missing = []
-        self.interactive = not "--no-interactive" in args
+        self.interactive = not "--no-interactive" in self.args
 
-        if not self.interactive or "--offline" in args:
-            for i in range(len(args)):
-                if "--no-interactive" in args[i].lower():
-                    del args[i]
+        if not self.interactive or "--offline" in self.args:
+            for i in range(len(self.args)):
+                if "--no-interactive" in self.args[i].lower():
+                    del self.args[i]
 
-                if "--offline" in args[i].lower():
-                    del args[i]
+                if "--offline" in self.args[i].lower():
+                    del self.args[i]
 
         self.flags = [
             {
@@ -87,7 +90,7 @@ def __init__(self, logger, dm=None, offline=False):
             },
         ]
 
-        self.handle(self.parse_flags(args))
+        self.handle(self.parse_flags(self.args))
 
     def help(self):
         try:
@@ -115,6 +118,10 @@ def help(self):
                     "[--plist] <path>",
                     "dumps hardware information into a plist file, inside of the specified directory",
                 ),
+                (
+                    "[--off-data] \"{data}\"",
+                    "disables detecting particular data (such as CPU, GPU, etc.) - must supply array under string."
+                ),
                 (
                     "[--no-interactive]",
                     "disables dynamic prompts for missing/invalid values (such as invalid dump types and [missing] paths)",
@@ -148,7 +155,7 @@ def help(self):
                 )
 
             print(
-                "\n\nExample:\n  <executable> -T ~/Downloads/myfolder -J ~/Downloads -X ~/Documents -P ."
+                "\n\nExample:\n  <executable> -T ~/Downloads/myfolder -J ~/Downloads -X ~/Documents -P . --off-data \"{Memory, Network}\""
             )
         except Exception as e:
             raise e
@@ -205,11 +212,40 @@ def handle(self, vals):
         self.logger.info("Successfully exited after dumping.\n\n", __file__)
         exit(0)
 
+    def off_data(self, arr):
+        orig = arr
+        arr = arr.replace("{", "").replace("}", "").split(", ")
+
+        if not arr:
+            return
+
+        repl = {
+            "cpu": "CPU",
+            "gpu": "GPU",
+            "motherboard": "Motherboard",
+            "memory": "Memory",
+            "network": "Network",
+            "audio": "Audio",
+            "input": "Input",
+            "storage": "Storage",
+        }
+
+        for val in arr:
+            val = repl[val.lower()]
+
+        self.toggled_off = arr
+
+        del self.args[self.args.index("--off-data")]
+        del self.args[self.args.index(orig)]
+
     def parse_flags(self, args):
         if list(filter(lambda x: "-h" in x.lower(), args)):
             self.help()
             exit(0)
 
+        if "--off-data" in args:
+            return self.off_data(args[args.index("--off-data") + 1])
+
         vals = []
 
         for i in range(len(args)):

src/cli/ui.py

@@ -69,6 +69,7 @@ def __init__(self, dm, logger):
         self.dm = dm
         self.logger = logger
         self.dump_dir = root
+        self.state = "menu"
 
     def handle_cmd(self, options=[]):
         cmd = input("\n\nPlease select an option: ")
@@ -102,6 +103,77 @@ def handle_cmd(self, options=[]):
             clear()
             self.create_ui()
 
+    def toggle_data(self):
+        # Sanitise the UI
+        clear()
+        title()
+
+        data_options = [
+            ("C", "CPU"),
+            ("G", "GPU"),
+            ("B", "Motherboard"),
+            ("M", "Memory"),
+            ("N", "Network"),
+            ("A", "Audio"),
+            ("I", "Input"),
+            ("S", "Storage"),
+        ]
+
+        opts = ["C", "G", "B", "M", "N", "A", "I", "S"]
+
+        for opt in data_options:
+            toggled = " " if opt[1] in self.dm.off_data else "X"
+
+            print(f"[{toggled}] ({opt[0]}) - {opt[1]}")
+
+        selected = input(
+            "Please select an option to toggle (or 'R'/'Q' to return): ")
+
+        if selected.lower() == "q" or selected.lower() == "r":
+            clear()
+            title()
+
+            print(color_text("Please wait while we adjust settings...\n", "green"))
+
+            asyncs = []
+
+            for opt in data_options:
+                if not opt[1] in self.dm.off_data and \
+                        not self.dm.info.get(opt[1]):
+                    asyncs.append(opt[1])
+
+            asyncs = ", ".join(asyncs)
+
+            print(f"Attempting to retrieve dumps for: {asyncs}...\n")
+
+            try:
+                self.dm.manager.dump()
+                self.dm.info = self.dm.manager.info
+
+                if input(color_text("[   OK   ] Successfully retrieved additional dumps.", "green") + " Press [enter] to return...") is not None:
+                    pass
+            except Exception:
+                if input(color_text("[ FAILED ] Unable to retrieve dumps.", "red") + " Press [enter] to return...") is not None:
+                    pass
+
+            if self.state == "discovery":
+                return self.discover()
+            else:
+                return self.create_ui()
+
+        if not selected.upper() in opts and \
+                input(color_text("Invalid option! Press [enter] to retry...", "red")) is not None:
+            self.toggle_data()
+
+        option = data_options[opts.index(selected.upper())][1]
+
+        if option in self.dm.off_data:
+            del self.dm.off_data[self.dm.off_data.index(option)]
+        else:
+            self.dm.off_data.append(option)
+
+        return self.toggle_data()
+
     def change_dump_dir(self):
         # Sanitise the UI
         clear()
@@ -142,7 +214,7 @@ def discover(self):
                     else self.dm.info[key] == {}
                 )
 
-                if key and not is_empty:
+                if key and not is_empty and not key in self.dm.off_data:
                     val = tree(key, self.dm.info[key])
                     print(val)
             except Exception as e:
@@ -161,6 +233,7 @@ def discover(self):
             (color_text("X. ", "yellow"), "Dump as XML"),
             (color_text("P. ", "yellow"), "Dump as Plist"),
             (color_text("C. ", "yellow"), "Change dump directory"),
+            (color_text("A. ", "yellow"), "Toggle data"),
             (color_text("Q. ", "yellow"), "Quit"),
         ]
 
@@ -171,6 +244,7 @@ def discover(self):
             ("X", "X.", self.dump_xml),
             ("P", "P.", self.dump_plist),
             ("C", "C.", self.change_dump_dir),
+            ("A", "A.", self.toggle_data),
             ("Q", "Q.", self.quit),
         ]
 
@@ -179,14 +253,16 @@ def discover(self):
 
         self.logger.info("Successfully ran 'discovery'.", __file__)
 
+        self.state = "discovery"
+
         self.handle_cmd(cmd_options)
 
     def dump_txt(self):
         return dump_txt(self.dm, self.dump_dir, self.logger)
 
     def dump_json(self):
         return dump_json(self.dm, self.dump_dir, self.logger)
-    
+
     def dump_xml(self):
         return dump_xml(self.dm, self.dump_dir, self.logger)
 
@@ -206,6 +282,7 @@ def create_ui(self):
             (color_text("X. ", "yellow"), "Dump as XML"),
             (color_text("P. ", "yellow"), "Dump as Plist"),
             (color_text("C. ", "yellow"), "Change dump directory"),
+            (color_text("A. ", "yellow"), "Toggle data"),
             ("\n\n", color_text("Q. ", "yellow"), "Quit"),
         ]
 
@@ -216,6 +293,7 @@ def create_ui(self):
             ("X", "X.", self.dump_xml),
             ("P", "P.", self.dump_plist),
             ("C", "C.", self.change_dump_dir),
+            ("A", "A.", self.toggle_data),
             ("Q", "Q.", self.quit),
         ]
 
@@ -243,6 +321,8 @@ def create_ui(self):
 
             self.logger.info("UI creation ran successfully.", __file__)
 
+            self.state = "menu"
+
             self.handle_cmd(cmd_options)
         except Exception as e:
             self.logger.critical(

src/dumps/Linux/linux.py

@@ -21,16 +21,25 @@ def __init__(self, parent):
         self.pci = parent.pci
         self.logger = parent.logger
         self.offline = parent.offline
+        self.off_data = parent.off_data
 
     def dump(self):
-        self.cpu_info()
-        self.mobo_info()
-        self.gpu_info()
-        self.mem_info()
-        self.net_info()
-        self.audio_info()
-        self.input_info()
-        self.block_info()
+        if not "CPU" in self.off_data and not self.info.get("CPU", []):
+            self.cpu_info()
+        if not "Motherboard" in self.off_data and not self.info.get("Motherboard", {}):
+            self.mobo_info()
+        if not "GPU" in self.off_data and not self.info.get("GPU", []):
+            self.gpu_info()
+        if not "Memory" in self.off_data and not self.info.get("Memory", []):
+            self.mem_info()
+        if not "Network" in self.off_data and not self.info.get("Network", []):
+            self.net_info()
+        if not "Audio" in self.off_data and not self.info.get("Audio", []):
+            self.audio_info()
+        if not "Input" in self.off_data and not self.info.get("Input", []):
+            self.input_info()
+        if not "Storage" in self.off_data and not self.info.get("Storage", []):
+            self.block_info()
 
     def cpu_info(self):
         try:

src/dumps/Windows/win.py

@@ -23,18 +23,27 @@ def __init__(self, parent):
         self.pci = parent.pci
         self.logger = parent.logger
         self.offline = parent.offline
+        self.off_data = parent.off_data
         self.cpu = {}
         self.c = wmi.WMI()
 
     def dump(self):
-        self.cpu_info()
-        self.gpu_info()
-        self.mem_info()
-        self.net_info()
-        self.audio_info()
-        self.mobo_info()
-        self.storage_info()
-        self.input_info()
+        if not "CPU" in self.off_data and not self.info.get("CPU", []):
+            self.cpu_info()
+        if not "Motherboard" in self.off_data and not self.info.get("Motherboard", {}):
+            self.mobo_info()
+        if not "GPU" in self.off_data and not self.info.get("GPU", []):
+            self.gpu_info()
+        if not "Memory" in self.off_data and not self.info.get("Memory", []):
+            self.mem_info()
+        if not "Network" in self.off_data and not self.info.get("Network", []):
+            self.net_info()
+        if not "Audio" in self.off_data and not self.info.get("Audio", []):
+            self.audio_info()
+        if not "Input" in self.off_data and not self.info.get("Input", []):
+            self.input_info()
+        if not "Storage" in self.off_data and not self.info.get("Storage", []):
+            self.storage_info()
 
     # Credits: https://github.com/flababah/cpuid.py/blob/master/example.py#L25
     def is_set(self, cpu, leaf, subleaf, reg_idx, bit):

src/dumps/macOS/mac.py

@@ -22,6 +22,7 @@ def __init__(self, parent):
         self.pci = parent.pci
         self.logger = parent.logger
         self.offline = parent.offline
+        self.off_data = parent.off_data
         self.vendor = None
         self.cpu = {}
 
@@ -31,13 +32,21 @@ def __init__(self, parent):
         }
 
     def dump(self):
-        self.cpu_info()
-        self.gpu_info()
-        self.mem_info()
-        self.net_info()
-        self.audio_info()
-        self.storage_info()
-        self.input_info()
+        if not "CPU" in self.off_data and not self.info.get("CPU", []):
+            self.cpu_info()
+        if not "GPU" in self.off_data and not self.info.get("GPU", []):
+            self.gpu_info()
+        if not "Memory" in self.off_data and not self.info.get("Memory", []):
+            self.mem_info()
+        if not "Network" in self.off_data and not self.info.get("Network", []):
+            self.net_info()
+        if not "Audio" in self.off_data and not self.info.get("Audio", []):
+            self.audio_info()
+        if not "Input" in self.off_data and not self.info.get("Input", []):
+            self.input_info()
+        if not "Storage" in self.off_data and not self.info.get("Storage", []):
+            self.storage_info()
+
 
     def cpu_info(self):
         try:

src/managers/devicemanager.py

@@ -5,7 +5,7 @@
 class DeviceManager:
     """Instance responsible for exposing all important information about the current system's hardware."""
 
-    def __init__(self, logger, offline=False):
+    def __init__(self, logger, off_data=[], offline=False):
         self.info = {
             "CPU": [],
             "Motherboard": {},
@@ -20,6 +20,7 @@ def __init__(self, logger, offline=False):
         self.platform = platform.system().lower()
         self.logger = logger
         self.offline = offline
+        self.off_data = off_data
 
         if self.platform == "darwin":
             from src.dumps.macOS.mac import MacHardwareManager