diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index f51f32c..0115289 100644 --- a/src/password_manager/manager.py +++ b/src/password_manager/manager.py @@ -17,6 +17,7 @@ import os from typing import Optional import shutil import time +import select from termcolor import colored from password_manager.encryption import EncryptionManager @@ -937,7 +938,7 @@ class PasswordManager: self.last_update = time.time() print( colored( - f"\nImported \u2714 Codes for {label} are now stored in SeedPass.", + f"\nImported \u2714 Codes for {label} are now stored in SeedPass at ID {entry_id}.", "green", ) ) @@ -1184,32 +1185,50 @@ class PasswordManager: try: data = self.entry_manager.vault.load_index() entries = data.get("entries", {}) - totp_list: list[tuple[str, int, int]] = [] + totp_list: list[tuple[str, int, int, bool]] = [] for idx_str, entry in entries.items(): if entry.get("type") == EntryType.TOTP.value: label = entry.get("label", "") period = int(entry.get("period", 30)) - totp_list.append((label, int(idx_str), period)) + imported = "secret" in entry + totp_list.append((label, int(idx_str), period, imported)) if not totp_list: print(colored("No 2FA entries found.", "yellow")) return totp_list.sort(key=lambda t: t[0].lower()) - - print(colored("Press Ctrl+C to return to the menu.", "cyan")) + print(colored("Press 'b' then Enter to return to the menu.", "cyan")) while True: print("\033c", end="") - for label, idx, period in totp_list: - code = self.entry_manager.get_totp_code(idx, self.parent_seed) - remaining = self.entry_manager.get_totp_time_remaining(idx) - filled = int(20 * (period - remaining) / period) - bar = "[" + "#" * filled + "-" * (20 - filled) + "]" - print(f"{label}: {code} {bar} {remaining:2d}s") + print(colored("Press 'b' then Enter to return to the menu.", "cyan")) + generated = [t for t in totp_list if not t[3]] + imported_list = [t for t in totp_list if t[3]] + if generated: + print(colored("\nGenerated 2FA Codes:", "green")) + for label, idx, period, _ in generated: + code = self.entry_manager.get_totp_code(idx, self.parent_seed) + remaining = self.entry_manager.get_totp_time_remaining(idx) + filled = int(20 * (period - remaining) / period) + bar = "[" + "#" * filled + "-" * (20 - filled) + "]" + print(f"[{idx}] {label}: {code} {bar} {remaining:2d}s") + if imported_list: + print(colored("\nImported 2FA Codes:", "green")) + for label, idx, period, _ in imported_list: + code = self.entry_manager.get_totp_code(idx, self.parent_seed) + remaining = self.entry_manager.get_totp_time_remaining(idx) + filled = int(20 * (period - remaining) / period) + bar = "[" + "#" * filled + "-" * (20 - filled) + "]" + print(f"[{idx}] {label}: {code} {bar} {remaining:2d}s") sys.stdout.flush() - time.sleep(1) - except KeyboardInterrupt: - print() + try: + if sys.stdin in select.select([sys.stdin], [], [], 1)[0]: + user_input = sys.stdin.readline().strip().lower() + if user_input == "b": + break + except KeyboardInterrupt: + print() + break except Exception as e: logging.error(f"Error displaying TOTP codes: {e}", exc_info=True) print(colored(f"Error: Failed to display TOTP codes: {e}", "red")) diff --git a/src/tests/test_manager_add_totp.py b/src/tests/test_manager_add_totp.py index 6ee18b9..f7ccdf6 100644 --- a/src/tests/test_manager_add_totp.py +++ b/src/tests/test_manager_add_totp.py @@ -21,7 +21,7 @@ class FakeNostrClient: return None, "abcd" -def test_handle_add_totp(monkeypatch): +def test_handle_add_totp(monkeypatch, capsys): with TemporaryDirectory() as tmpdir: tmp_path = Path(tmpdir) vault, enc_mgr = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD) @@ -51,6 +51,7 @@ def test_handle_add_totp(monkeypatch): monkeypatch.setattr(pm, "sync_vault", lambda: None) pm.handle_add_totp() + out = capsys.readouterr().out entry = entry_mgr.retrieve_entry(0) assert entry == { @@ -60,3 +61,4 @@ def test_handle_add_totp(monkeypatch): "period": 30, "digits": 6, } + assert "ID 0" in out diff --git a/src/tests/test_manager_display_totp_codes.py b/src/tests/test_manager_display_totp_codes.py index 544b4e1..0bd562a 100644 --- a/src/tests/test_manager_display_totp_codes.py +++ b/src/tests/test_manager_display_totp_codes.py @@ -45,12 +45,14 @@ def test_handle_display_totp_codes(monkeypatch, capsys): pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30 ) - def interrupt(_): - raise KeyboardInterrupt() - - monkeypatch.setattr("password_manager.manager.time.sleep", interrupt) + # interrupt the loop after first iteration + monkeypatch.setattr( + "password_manager.manager.select.select", + lambda *a, **k: (_ for _ in ()).throw(KeyboardInterrupt()), + ) pm.handle_display_totp_codes() out = capsys.readouterr().out - assert "Example" in out + assert "Generated 2FA Codes" in out + assert "[0] Example" in out assert "123456" in out