refactor(pw): Improve handling of password input

commands/util/input_mgr.py

@@ -1,7 +1,6 @@
 from typing import Optional
-import getpass
+
-import time
+from commands.util.password_manager import PasswordManager
-import colorama
 
 
 def choice(default: str, options: str, msg: str) -> str:
@@ -33,18 +32,13 @@ def confirm_overwrite(msg: Optional[str] = ""):
     )
 
 
-def password(msg: str = ""):
+# Use class to store the password
-    pw = ""
+pw_manager = PasswordManager()
-    if msg != "":
+
-        pw = getpass.getpass(msg)
+
-    else:
+def password():
-        pw = getpass.getpass()
+    return pw_manager.get()
-    if pw != "":
+
-        return pw
+
-    else:
+def unlock_sudo():
-        time.sleep(1)
+    return pw_manager.validate()
-        print(
-            colorama.Fore.RED + "Error:",
-            colorama.Style.RESET_ALL + "Password cannot be empty",
-        )
-        return password(msg)

commands/util/password_manager.py

@@ -0,0 +1,70 @@
+import getpass
+import time
+import colorama
+import subprocess as sp
+
+
+class PasswordManager:
+    _pw = ""
+    _wrong_cnt = 0
+    _valid = False
+
+    def get(self, msg: str = ""):
+        """Get the user's password (uses cached password if PW is valid)
+        Otherwise prompts user
+
+        Args:
+            msg: The message to use for the password prompt
+
+        Returns:
+            The user's password
+        """
+        if self._valid or self.validate():
+            return self._pw
+
+        if msg != "":
+            self._pw = getpass.getpass(msg)
+        else:
+            self._pw = getpass.getpass()
+        if self._pw != "":
+            if not self.validate():
+                print(
+                    colorama.Fore.RED + "Error:",
+                    colorama.Style.RESET_ALL + "Invalid Password. Please try again",
+                )
+                return self.get(msg)
+            return self._pw
+        else:
+            time.sleep(1)
+            print(
+                colorama.Fore.RED + "Error:",
+                colorama.Style.RESET_ALL + "Password cannot be empty",
+            )
+            return self.get(msg)
+
+    def validate(self) -> bool:
+        """Validate that the password is correct by running sudo command
+
+        Returns:
+            True if password is valid, False otherwise
+        """
+
+        def helper():
+            if self._pw == "":
+                return False
+            try:
+                sp.run(
+                    ["sudo", "-k"], capture_output=True, text=True
+                ).check_returncode()
+                sp.run(
+                    ["sudo", "-S", "echo"],
+                    capture_output=True,
+                    input=self._pw,
+                    text=True,
+                ).check_returncode()
+            except Exception:
+                return False
+            return True
+
+        self._valid = helper()
+        return self._valid