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

@@ -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."""

View File

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

View File

@@ -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):

View File

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