Fix TOTP seed retrieval and auto index

This commit is contained in:
thePR0M3TH3AN
2025-07-03 07:39:13 -04:00
parent b87f30a730
commit 82306a3a4b
5 changed files with 38 additions and 22 deletions

View File

@@ -179,7 +179,7 @@ python src/main.py
1. From the main menu choose **Add Entry** and select **2FA (TOTP)**. 1. From the main menu choose **Add Entry** and select **2FA (TOTP)**.
2. Provide a label for the account (for example, `GitHub`). 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. 4. Optionally specify the TOTP period and digit count.
5. SeedPass will display an `otpauth://` URI and secret that you can manually 5. SeedPass will display an `otpauth://` URI and secret that you can manually
enter into your authenticator app. enter into your authenticator app.

View File

@@ -153,11 +153,29 @@ class EntryManager:
print(colored(f"Error: Failed to add entry: {e}", "red")) print(colored(f"Error: Failed to add entry: {e}", "red"))
sys.exit(1) 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( 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: ) -> str:
"""Add a new TOTP entry and return the provisioning URI.""" """Add a new TOTP entry and return the provisioning URI."""
entry_id = self.get_next_index() entry_id = self.get_next_index()
if index is None:
index = self.get_next_totp_index()
data = self.vault.load_index() data = self.vault.load_index()
data.setdefault("entries", {}) data.setdefault("entries", {})
@@ -174,8 +192,7 @@ class EntryManager:
self.backup_index_file() self.backup_index_file()
try: try:
seed = self.vault.encryption_manager.decrypt_parent_seed() secret = TotpManager.derive_secret(parent_seed, index)
secret = TotpManager.derive_secret(seed, index)
return TotpManager.make_otpauth_uri(label, secret, period, digits) return TotpManager.make_otpauth_uri(label, secret, period, digits)
except Exception as e: except Exception as e:
logger.error(f"Failed to generate otpauth URI: {e}") logger.error(f"Failed to generate otpauth URI: {e}")
@@ -203,15 +220,16 @@ class EntryManager:
self.backup_index_file() self.backup_index_file()
raise NotImplementedError("Seed entry support not implemented yet") 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.""" """Return the current TOTP code for the specified entry."""
entry = self.retrieve_entry(index) entry = self.retrieve_entry(index)
if not entry or entry.get("type") != EntryType.TOTP.value: if not entry or entry.get("type") != EntryType.TOTP.value:
raise ValueError("Entry is not a TOTP entry") raise ValueError("Entry is not a TOTP entry")
seed = self.vault.encryption_manager.decrypt_parent_seed()
totp_index = int(entry.get("index", 0)) 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: def get_totp_time_remaining(self, index: int) -> int:
"""Return seconds remaining in the TOTP period for the given entry.""" """Return seconds remaining in the TOTP period for the given entry."""

View File

@@ -870,11 +870,7 @@ class PasswordManager:
print(colored("Error: Label cannot be empty.", "red")) print(colored("Error: Label cannot be empty.", "red"))
return return
index_input = input("Enter derivation index (number): ").strip() totp_index = self.entry_manager.get_next_totp_index()
if not index_input.isdigit():
print(colored("Error: Index must be a number.", "red"))
return
totp_index = int(index_input)
period_input = input("TOTP period in seconds (default 30): ").strip() period_input = input("TOTP period in seconds (default 30): ").strip()
period = 30 period = 30
@@ -893,7 +889,13 @@ class PasswordManager:
digits = int(digits_input) digits = int(digits_input)
entry_id = self.entry_manager.get_next_index() 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.is_dirty = True
self.last_update = time.time() self.last_update = time.time()

View File

@@ -52,8 +52,7 @@ def test_round_trip_entry_types(method, expected_type):
if method == "add_entry": if method == "add_entry":
index = entry_mgr.add_entry("example.com", 8) index = entry_mgr.add_entry("example.com", 8)
elif method == "add_totp": elif method == "add_totp":
with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): entry_mgr.add_totp("example", TEST_SEED)
entry_mgr.add_totp("example", 0)
index = 0 index = 0
else: else:
with pytest.raises(NotImplementedError): with pytest.raises(NotImplementedError):

View File

@@ -19,8 +19,7 @@ def test_add_totp_and_get_code():
vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD) vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD)
entry_mgr = EntryManager(vault, Path(tmpdir)) entry_mgr = EntryManager(vault, Path(tmpdir))
with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): uri = entry_mgr.add_totp("Example", TEST_SEED)
uri = entry_mgr.add_totp("Example", 0)
assert uri.startswith("otpauth://totp/") assert uri.startswith("otpauth://totp/")
entry = entry_mgr.retrieve_entry(0) entry = entry_mgr.retrieve_entry(0)
@@ -32,8 +31,7 @@ def test_add_totp_and_get_code():
"digits": 6, "digits": 6,
} }
with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): code = entry_mgr.get_totp_code(0, TEST_SEED, timestamp=0)
code = entry_mgr.get_totp_code(0, timestamp=0)
expected = TotpManager.current_code(TEST_SEED, 0, timestamp=0) expected = TotpManager.current_code(TEST_SEED, 0, timestamp=0)
assert code == expected assert code == expected
@@ -44,8 +42,7 @@ def test_totp_time_remaining(monkeypatch):
vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD) vault, enc_mgr = create_vault(Path(tmpdir), TEST_SEED, TEST_PASSWORD)
entry_mgr = EntryManager(vault, Path(tmpdir)) entry_mgr = EntryManager(vault, Path(tmpdir))
with patch.object(enc_mgr, "decrypt_parent_seed", return_value=TEST_SEED): entry_mgr.add_totp("Example", TEST_SEED)
entry_mgr.add_totp("Example", 0)
monkeypatch.setattr(TotpManager, "time_remaining", lambda period: 7) monkeypatch.setattr(TotpManager, "time_remaining", lambda period: 7)
remaining = entry_mgr.get_totp_time_remaining(0) remaining = entry_mgr.get_totp_time_remaining(0)