diff --git a/src/seedpass/core/entry_management.py b/src/seedpass/core/entry_management.py index c3b003ad..1d19cfc4 100644 --- a/src/seedpass/core/entry_management.py +++ b/src/seedpass/core/entry_management.py @@ -694,6 +694,26 @@ class EntryManager: seed_index = int(entry.get("index", index)) return derive_seed_phrase(bip85, seed_index, words) + def get_totp_secret( + self, + index: int, + parent_seed: str | bytes | None = None, + ) -> str: + """Return the TOTP secret for the specified entry.""" + entry = self.retrieve_entry(index) + etype = entry.get("type") if entry else None + kind = entry.get("kind") if entry else None + if not entry or ( + etype != EntryType.TOTP.value and kind != EntryType.TOTP.value + ): + raise ValueError("Entry is not a TOTP entry") + if entry.get("deterministic", False) or "secret" not in entry: + if parent_seed is None: + raise ValueError("Seed required for derived TOTP") + totp_index = int(entry.get("index", 0)) + return TotpManager.derive_secret(parent_seed, totp_index) + return entry["secret"] + def get_totp_code( self, index: int, diff --git a/src/seedpass/core/manager.py b/src/seedpass/core/manager.py index 9c25101a..beecc915 100644 --- a/src/seedpass/core/manager.py +++ b/src/seedpass/core/manager.py @@ -2763,9 +2763,10 @@ class PasswordManager: print(colored(f"Retrieving 2FA code for '{label}'.", "cyan")) print(colored("Press Enter to return to the menu.", "cyan")) try: + key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None) + secret = self.entry_manager.get_totp_secret(index, key) while True: - key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None) - code = self.entry_manager.get_totp_code(index, key) + code = TotpManager.current_code_from_secret(secret) if self.secret_mode_enabled: if copy_to_clipboard(code, self.clipboard_clear_delay): print( diff --git a/src/seedpass/core/menu_handler.py b/src/seedpass/core/menu_handler.py index 6cbbfefa..ef60177c 100644 --- a/src/seedpass/core/menu_handler.py +++ b/src/seedpass/core/menu_handler.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING from termcolor import colored from .entry_types import EntryType, ALL_ENTRY_TYPES +from .totp import TotpManager import seedpass.core.manager as manager_module from utils.color_scheme import color_text from utils.terminal_utils import clear_header_with_notification @@ -118,6 +119,15 @@ class MenuHandler: return totp_list.sort(key=lambda t: t[0].lower()) + + totp_secrets: dict[int, str] = {} + key = getattr(pm, "KEY_TOTP_DET", None) or getattr(pm, "parent_seed", None) + for _, idx, _, _ in totp_list: + try: + totp_secrets[idx] = pm.entry_manager.get_totp_secret(idx, key) + except Exception: + pass + print(colored("Press Enter to return to the menu.", "cyan")) while True: fp, parent_fp, child_fp = pm.header_fingerprint_args @@ -134,10 +144,12 @@ class MenuHandler: if generated: print(colored("\nGenerated 2FA Codes:", "green")) for label, idx, period, _ in generated: - key = getattr(pm, "KEY_TOTP_DET", None) or getattr( - pm, "parent_seed", None + secret = totp_secrets.get(idx) + code = ( + TotpManager.current_code_from_secret(secret) + if secret + else "ERROR" ) - code = pm.entry_manager.get_totp_code(idx, key) remaining = pm.entry_manager.get_totp_time_remaining(idx) filled = int(20 * (period - remaining) / period) bar = "[" + "#" * filled + "-" * (20 - filled) + "]" @@ -155,10 +167,12 @@ class MenuHandler: if imported_list: print(colored("\nImported 2FA Codes:", "green")) for label, idx, period, _ in imported_list: - key = getattr(pm, "KEY_TOTP_DET", None) or getattr( - pm, "parent_seed", None + secret = totp_secrets.get(idx) + code = ( + TotpManager.current_code_from_secret(secret) + if secret + else "ERROR" ) - code = pm.entry_manager.get_totp_code(idx, key) remaining = pm.entry_manager.get_totp_time_remaining(idx) filled = int(20 * (period - remaining) / period) bar = "[" + "#" * filled + "-" * (20 - filled) + "]" diff --git a/src/tests/test_manager_display_totp_codes.py b/src/tests/test_manager_display_totp_codes.py index 9c567047..9e62e3d5 100644 --- a/src/tests/test_manager_display_totp_codes.py +++ b/src/tests/test_manager_display_totp_codes.py @@ -21,7 +21,10 @@ def test_handle_display_totp_codes(monkeypatch, capsys, password_manager): pm.entry_manager.add_totp("Example", TEST_SEED) - monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456") + monkeypatch.setattr( + "seedpass.core.totp.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", + ) monkeypatch.setattr(pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30) # interrupt the loop after first iteration @@ -45,7 +48,10 @@ def test_display_totp_codes_excludes_archived(monkeypatch, capsys, password_mana pm.entry_manager.add_totp("Hidden", TEST_SEED) pm.entry_manager.modify_entry(1, archived=True) - monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456") + monkeypatch.setattr( + "seedpass.core.totp.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", + ) monkeypatch.setattr(pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30) monkeypatch.setattr( diff --git a/src/tests/test_manager_list_entries.py b/src/tests/test_manager_list_entries.py index 49e4086b..816ec66b 100644 --- a/src/tests/test_manager_list_entries.py +++ b/src/tests/test_manager_list_entries.py @@ -346,7 +346,8 @@ def test_show_entry_details_sensitive(monkeypatch, capsys, entry_type): entry_mgr.add_totp("Example", TEST_SEED) idx = 0 monkeypatch.setattr( - pm.entry_manager, "get_totp_code", lambda *a, **k: "123456" + "seedpass.core.manager.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", ) monkeypatch.setattr( pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 1 @@ -383,7 +384,8 @@ def test_show_entry_details_with_enum_type(monkeypatch, capsys, entry_type): entry_mgr.add_totp("Example", TEST_SEED) idx = 0 monkeypatch.setattr( - pm.entry_manager, "get_totp_code", lambda *a, **k: "123456" + "seedpass.core.manager.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", ) monkeypatch.setattr( pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 1 diff --git a/src/tests/test_manager_retrieve_totp.py b/src/tests/test_manager_retrieve_totp.py index 8d01c28d..003084ec 100644 --- a/src/tests/test_manager_retrieve_totp.py +++ b/src/tests/test_manager_retrieve_totp.py @@ -45,7 +45,10 @@ def test_handle_retrieve_totp_entry(monkeypatch, capsys): inputs = iter(["0", "n", ""]) monkeypatch.setattr("builtins.input", lambda *a, **k: next(inputs)) - monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456") + monkeypatch.setattr( + "seedpass.core.manager.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", + ) monkeypatch.setattr( pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 1 ) diff --git a/src/tests/test_secret_mode.py b/src/tests/test_secret_mode.py index 4e3c0b10..be5222d3 100644 --- a/src/tests/test_secret_mode.py +++ b/src/tests/test_secret_mode.py @@ -62,7 +62,10 @@ def test_totp_display_secret_mode(monkeypatch, capsys): pm, entry_mgr = setup_pm(tmp) entry_mgr.add_totp("Example", TEST_SEED) - monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456") + monkeypatch.setattr( + "seedpass.core.totp.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", + ) monkeypatch.setattr( pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30 ) @@ -112,7 +115,10 @@ def test_totp_display_no_secret_mode(monkeypatch, capsys): pm.secret_mode_enabled = False entry_mgr.add_totp("Example", TEST_SEED) - monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456") + monkeypatch.setattr( + "seedpass.core.totp.TotpManager.current_code_from_secret", + lambda *a, **k: "123456", + ) monkeypatch.setattr( pm.entry_manager, "get_totp_time_remaining", lambda *a, **k: 30 )