From 82306a3a4bedaa6e49fd51a75967646e995f6a21 Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Thu, 3 Jul 2025 07:39:13 -0400 Subject: [PATCH] Fix TOTP seed retrieval and auto index --- README.md | 2 +- src/password_manager/entry_management.py | 30 +++++++++++++++++++----- src/password_manager/manager.py | 14 ++++++----- src/tests/test_entry_add.py | 3 +-- src/tests/test_totp_entry.py | 11 ++++----- 5 files changed, 38 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 6256f76..bc5a3d8 100644 --- a/README.md +++ b/README.md @@ -179,7 +179,7 @@ python src/main.py 1. From the main menu choose **Add Entry** and select **2FA (TOTP)**. 2. Provide a label for the account (for example, `GitHub`). -3. Enter the derivation index you wish to use for this 2FA code. +3. SeedPass automatically chooses the next available derivation index. 4. Optionally specify the TOTP period and digit count. 5. SeedPass will display an `otpauth://` URI and secret that you can manually enter into your authenticator app. diff --git a/src/password_manager/entry_management.py b/src/password_manager/entry_management.py index 42043b2..70f1832 100644 --- a/src/password_manager/entry_management.py +++ b/src/password_manager/entry_management.py @@ -153,11 +153,29 @@ class EntryManager: print(colored(f"Error: Failed to add entry: {e}", "red")) sys.exit(1) + def get_next_totp_index(self) -> int: + """Return the next available derivation index for TOTP secrets.""" + data = self.vault.load_index() + entries = data.get("entries", {}) + indices = [ + int(v.get("index", 0)) + for v in entries.values() + if v.get("type") == EntryType.TOTP.value + ] + return (max(indices) + 1) if indices else 0 + def add_totp( - self, label: str, index: int, period: int = 30, digits: int = 6 + self, + label: str, + parent_seed: str, + index: int | None = None, + period: int = 30, + digits: int = 6, ) -> str: """Add a new TOTP entry and return the provisioning URI.""" entry_id = self.get_next_index() + if index is None: + index = self.get_next_totp_index() data = self.vault.load_index() data.setdefault("entries", {}) @@ -174,8 +192,7 @@ class EntryManager: self.backup_index_file() try: - seed = self.vault.encryption_manager.decrypt_parent_seed() - secret = TotpManager.derive_secret(seed, index) + secret = TotpManager.derive_secret(parent_seed, index) return TotpManager.make_otpauth_uri(label, secret, period, digits) except Exception as e: logger.error(f"Failed to generate otpauth URI: {e}") @@ -203,15 +220,16 @@ class EntryManager: self.backup_index_file() raise NotImplementedError("Seed entry support not implemented yet") - def get_totp_code(self, index: int, timestamp: int | None = None) -> str: + def get_totp_code( + self, index: int, parent_seed: str, timestamp: int | None = None + ) -> str: """Return the current TOTP code for the specified entry.""" entry = self.retrieve_entry(index) if not entry or entry.get("type") != EntryType.TOTP.value: raise ValueError("Entry is not a TOTP entry") - seed = self.vault.encryption_manager.decrypt_parent_seed() totp_index = int(entry.get("index", 0)) - return TotpManager.current_code(seed, totp_index, timestamp) + return TotpManager.current_code(parent_seed, totp_index, timestamp) def get_totp_time_remaining(self, index: int) -> int: """Return seconds remaining in the TOTP period for the given entry.""" diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index 154602c..f324547 100644 --- a/src/password_manager/manager.py +++ b/src/password_manager/manager.py @@ -870,11 +870,7 @@ class PasswordManager: print(colored("Error: Label cannot be empty.", "red")) return - index_input = input("Enter derivation index (number): ").strip() - if not index_input.isdigit(): - print(colored("Error: Index must be a number.", "red")) - return - totp_index = int(index_input) + totp_index = self.entry_manager.get_next_totp_index() period_input = input("TOTP period in seconds (default 30): ").strip() period = 30 @@ -893,7 +889,13 @@ class PasswordManager: digits = int(digits_input) entry_id = self.entry_manager.get_next_index() - uri = self.entry_manager.add_totp(label, totp_index, period, digits) + uri = self.entry_manager.add_totp( + label, + self.parent_seed, + index=totp_index, + period=period, + digits=digits, + ) self.is_dirty = True self.last_update = time.time() diff --git a/src/tests/test_entry_add.py b/src/tests/test_entry_add.py index a64b068..71731b8 100644 --- a/src/tests/test_entry_add.py +++ b/src/tests/test_entry_add.py @@ -52,8 +52,7 @@ def test_round_trip_entry_types(method, expected_type): if method == "add_entry": index = entry_mgr.add_entry("example.com", 8) elif method == "add_totp": - with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): - entry_mgr.add_totp("example", 0) + entry_mgr.add_totp("example", TEST_SEED) index = 0 else: with pytest.raises(NotImplementedError): diff --git a/src/tests/test_totp_entry.py b/src/tests/test_totp_entry.py index 371ceaa..a6d946a 100644 --- a/src/tests/test_totp_entry.py +++ b/src/tests/test_totp_entry.py @@ -19,9 +19,8 @@ def test_add_totp_and_get_code(): vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD) entry_mgr = EntryManager(vault, Path(tmpdir)) - with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): - uri = entry_mgr.add_totp("Example", 0) - assert uri.startswith("otpauth://totp/") + uri = entry_mgr.add_totp("Example", TEST_SEED) + assert uri.startswith("otpauth://totp/") entry = entry_mgr.retrieve_entry(0) assert entry == { @@ -32,8 +31,7 @@ def test_add_totp_and_get_code(): "digits": 6, } - with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): - code = entry_mgr.get_totp_code(0, timestamp=0) + code = entry_mgr.get_totp_code(0, TEST_SEED, timestamp=0) expected = TotpManager.current_code(TEST_SEED, 0, timestamp=0) assert code == expected @@ -44,8 +42,7 @@ def test_totp_time_remaining(monkeypatch): vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD) entry_mgr = EntryManager(vault, Path(tmpdir)) - with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): - entry_mgr.add_totp("Example", 0) + entry_mgr.add_totp("Example", TEST_SEED) monkeypatch.setattr(TotpManager, "time_remaining", lambda period: 7) remaining = entry_mgr.get_totp_time_remaining(0)