Merge pull request #193 from PR0M3TH3AN/codex/fix-2fa-blacklist-behavior-and-exit-issue

Fix TOTP blacklist and retrieval exit

src/password_manager/manager.py

@@ -1024,15 +1024,31 @@ class PasswordManager:
                         print(colored(f"Code: {code}", "yellow"))
                         if notes:
                             print(colored(f"Notes: {notes}", "cyan"))
-                        TotpManager.print_progress_bar(period)
+                        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], [], [], 0)[0]:
+                                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)
                     print(colored(f"Error: Failed to generate TOTP code: {e}", "red"))
@@ -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

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

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()