From ef4a9966d19c56ee8b5dc56ec8b2415ee0d71b51 Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Thu, 3 Jul 2025 10:40:47 -0400
Subject: [PATCH] fix totp blacklist and exit

---
 src/password_manager/manager.py              | 36 ++++++++++++++-----
 src/tests/test_manager_display_totp_codes.py | 38 ++++++++++++++++++++
 src/tests/test_manager_retrieve_totp.py      |  8 +++--
 3 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py
index 473a51f..07c6156 100644
--- a/src/password_manager/manager.py
+++ b/src/password_manager/manager.py
@@ -1024,14 +1024,30 @@ class PasswordManager:
                         print(colored(f"Code: {code}", "yellow"))
                         if notes:
                             print(colored(f"Notes: {notes}", "cyan"))
-                        TotpManager.print_progress_bar(period)
-                        try:
-                            if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
-                                user_input = sys.stdin.readline().strip().lower()
-                                if user_input == "b":
-                                    break
-                        except KeyboardInterrupt:
-                            print()
+                        remaining = self.entry_manager.get_totp_time_remaining(index)
+                        exit_loop = False
+                        while remaining > 0:
+                            filled = int(20 * (period - remaining) / period)
+                            bar = "[" + "#" * filled + "-" * (20 - filled) + "]"
+                            sys.stdout.write(f"\r{bar} {remaining:2d}s")
+                            sys.stdout.flush()
+                            try:
+                                if (
+                                    sys.stdin
+                                    in select.select([sys.stdin], [], [], 1)[0]
+                                ):
+                                    user_input = sys.stdin.readline().strip().lower()
+                                    if user_input == "b":
+                                        exit_loop = True
+                                        break
+                            except KeyboardInterrupt:
+                                exit_loop = True
+                                print()
+                                break
+                            remaining -= 1
+                        sys.stdout.write("\n")
+                        sys.stdout.flush()
+                        if exit_loop:
                             break
                 except Exception as e:
                     logging.error(f"Error generating TOTP code: {e}", exc_info=True)
@@ -1240,7 +1256,9 @@ class PasswordManager:
             entries = data.get("entries", {})
             totp_list: list[tuple[str, int, int, bool]] = []
             for idx_str, entry in entries.items():
-                if entry.get("type") == EntryType.TOTP.value:
+                if entry.get("type") == EntryType.TOTP.value and not entry.get(
+                    "blacklisted", False
+                ):
                     label = entry.get("label", "")
                     period = int(entry.get("period", 30))
                     imported = "secret" in entry
diff --git a/src/tests/test_manager_display_totp_codes.py b/src/tests/test_manager_display_totp_codes.py
index 0bd562a..b462c36 100644
--- a/src/tests/test_manager_display_totp_codes.py
+++ b/src/tests/test_manager_display_totp_codes.py
@@ -56,3 +56,41 @@ def test_handle_display_totp_codes(monkeypatch, capsys):
         assert "Generated 2FA Codes" in out
         assert "[0] Example" in out
         assert "123456" in out
+
+
+def test_display_totp_codes_excludes_blacklisted(monkeypatch, capsys):
+    with TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        vault, enc_mgr = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD)
+        entry_mgr = EntryManager(vault, tmp_path)
+        backup_mgr = BackupManager(tmp_path)
+
+        pm = PasswordManager.__new__(PasswordManager)
+        pm.encryption_mode = EncryptionMode.SEED_ONLY
+        pm.encryption_manager = enc_mgr
+        pm.vault = vault
+        pm.entry_manager = entry_mgr
+        pm.backup_manager = backup_mgr
+        pm.parent_seed = TEST_SEED
+        pm.nostr_client = FakeNostrClient()
+        pm.fingerprint_dir = tmp_path
+        pm.is_dirty = False
+
+        entry_mgr.add_totp("Visible", TEST_SEED)
+        entry_mgr.add_totp("Hidden", TEST_SEED)
+        entry_mgr.modify_entry(1, blacklisted=True)
+
+        monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456")
+        monkeypatch.setattr(
+            pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30
+        )
+
+        monkeypatch.setattr(
+            "password_manager.manager.select.select",
+            lambda *a, **k: (_ for _ in ()).throw(KeyboardInterrupt()),
+        )
+
+        pm.handle_display_totp_codes()
+        out = capsys.readouterr().out
+        assert "Visible" in out
+        assert "Hidden" not in out
diff --git a/src/tests/test_manager_retrieve_totp.py b/src/tests/test_manager_retrieve_totp.py
index 09c3824..c29b6a3 100644
--- a/src/tests/test_manager_retrieve_totp.py
+++ b/src/tests/test_manager_retrieve_totp.py
@@ -42,10 +42,14 @@ def test_handle_retrieve_totp_entry(monkeypatch, capsys):
 
         monkeypatch.setattr("builtins.input", lambda *a, **k: "0")
         monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456")
-        monkeypatch.setattr(TotpManager, "print_progress_bar", lambda period: None)
+        monkeypatch.setattr(
+            pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 1
+        )
+        monkeypatch.setattr("password_manager.manager.time.sleep", lambda *a, **k: None)
+        monkeypatch.setattr(sys.stdin, "readline", lambda *a, **k: "b\n")
         monkeypatch.setattr(
             "password_manager.manager.select.select",
-            lambda *a, **k: (_ for _ in ()).throw(KeyboardInterrupt()),
+            lambda *a, **k: ([sys.stdin], [], []),
         )
 
         pm.handle_retrieve_entry()