Add 2FA codes menu

This commit is contained in:
thePR0M3TH3AN
2025-07-03 08:30:43 -04:00
parent 11bd97dcfe
commit c7ffdd6991
7 changed files with 59 additions and 18 deletions

View File

@@ -167,10 +167,11 @@ python src/main.py
1. Add Entry
2. Retrieve Entry
3. Modify an Existing Entry
4. Settings
5. Exit
4. 2FA Codes
5. Settings
6. Exit
Enter your choice (1-5):
Enter your choice (1-6):
```
When choosing **Add Entry**, you can now select **Password** or **2FA (TOTP)**.

View File

@@ -101,10 +101,11 @@ Select an option:
1. Add Entry
2. Retrieve Entry
3. Modify an Existing Entry
4. Settings
5. Exit
4. 2FA Codes
5. Settings
6. Exit
Enter your choice (1-5):
Enter your choice (1-6):
</pre>
</div>
</section>

View File

@@ -548,8 +548,9 @@ def display_menu(
1. Add Entry
2. Retrieve Entry
3. Modify an Existing Entry
4. Settings
5. Exit
4. 2FA Codes
5. Settings
6. Exit
"""
while True:
if time.time() - password_manager.last_activity > inactivity_timeout:
@@ -571,7 +572,7 @@ def display_menu(
print(colored(menu, "cyan"))
try:
choice = timed_input(
"Enter your choice (1-5): ", inactivity_timeout
"Enter your choice (1-6): ", inactivity_timeout
).strip()
except TimeoutError:
print(colored("Session timed out. Vault locked.", "yellow"))
@@ -582,7 +583,7 @@ def display_menu(
if not choice:
print(
colored(
"No input detected. Please enter a number between 1 and 5.",
"No input detected. Please enter a number between 1 and 6.",
"yellow",
)
)
@@ -613,8 +614,11 @@ def display_menu(
password_manager.handle_modify_entry()
elif choice == "4":
password_manager.update_activity()
handle_settings(password_manager)
password_manager.handle_display_totp_codes()
elif choice == "5":
password_manager.update_activity()
handle_settings(password_manager)
elif choice == "6":
logging.info("Exiting the program.")
print(colored("Exiting the program.", "green"))
password_manager.nostr_client.close_client_pool()

View File

@@ -1178,6 +1178,41 @@ class PasswordManager:
logging.error(f"Error during entry deletion: {e}", exc_info=True)
print(colored(f"Error: Failed to delete entry: {e}", "red"))
def handle_display_totp_codes(self) -> None:
"""Display all stored TOTP codes with a countdown progress bar."""
try:
data = self.entry_manager.vault.load_index()
entries = data.get("entries", {})
totp_list: list[tuple[str, int, int]] = []
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))
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"))
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")
sys.stdout.flush()
time.sleep(1)
except KeyboardInterrupt:
print()
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"))
def handle_verify_checksum(self) -> None:
"""
Handles verifying the script's checksum against the stored checksum to ensure integrity.

View File

@@ -31,7 +31,7 @@ def test_auto_sync_triggers_post(monkeypatch):
called = True
monkeypatch.setattr(main, "handle_post_to_nostr", fake_post)
monkeypatch.setattr(main, "timed_input", lambda *_: "5")
monkeypatch.setattr(main, "timed_input", lambda *_: "6")
with pytest.raises(SystemExit):
main.display_menu(pm, sync_interval=0.1)

View File

@@ -52,7 +52,7 @@ def _make_pm(called, locked=None):
def test_empty_and_non_numeric_choice(monkeypatch, capsys):
called = {"add": False, "retrieve": False, "modify": False}
pm, _ = _make_pm(called)
inputs = iter(["", "abc", "5"])
inputs = iter(["", "abc", "6"])
monkeypatch.setattr(main, "timed_input", lambda *_: next(inputs))
with pytest.raises(SystemExit):
main.display_menu(pm, sync_interval=1000, inactivity_timeout=1000)
@@ -65,7 +65,7 @@ def test_empty_and_non_numeric_choice(monkeypatch, capsys):
def test_out_of_range_menu(monkeypatch, capsys):
called = {"add": False, "retrieve": False, "modify": False}
pm, _ = _make_pm(called)
inputs = iter(["9", "5"])
inputs = iter(["9", "6"])
monkeypatch.setattr(main, "timed_input", lambda *_: next(inputs))
with pytest.raises(SystemExit):
main.display_menu(pm, sync_interval=1000, inactivity_timeout=1000)
@@ -77,7 +77,7 @@ def test_out_of_range_menu(monkeypatch, capsys):
def test_invalid_add_entry_submenu(monkeypatch, capsys):
called = {"add": False, "retrieve": False, "modify": False}
pm, _ = _make_pm(called)
inputs = iter(["1", "4", "3", "5"])
inputs = iter(["1", "4", "3", "6"])
monkeypatch.setattr(main, "timed_input", lambda *_: next(inputs))
monkeypatch.setattr("builtins.input", lambda *_: next(inputs))
with pytest.raises(SystemExit):
@@ -92,7 +92,7 @@ def test_inactivity_timeout_loop(monkeypatch, capsys):
pm, locked = _make_pm(called)
pm.last_activity = 0
monkeypatch.setattr(time, "time", lambda: 100.0)
monkeypatch.setattr(main, "timed_input", lambda *_: "5")
monkeypatch.setattr(main, "timed_input", lambda *_: "6")
with pytest.raises(SystemExit):
main.display_menu(pm, sync_interval=1000, inactivity_timeout=0.1)
out = capsys.readouterr().out

View File

@@ -36,7 +36,7 @@ def test_inactivity_triggers_lock(monkeypatch):
unlock_vault=unlock_vault,
)
monkeypatch.setattr(main, "timed_input", lambda *_: "5")
monkeypatch.setattr(main, "timed_input", lambda *_: "6")
with pytest.raises(SystemExit):
main.display_menu(pm, sync_interval=1000, inactivity_timeout=0.1)
@@ -72,7 +72,7 @@ def test_input_timeout_triggers_lock(monkeypatch):
unlock_vault=unlock_vault,
)
responses = iter([TimeoutError(), "5"])
responses = iter([TimeoutError(), "6"])
def fake_input(*_args, **_kwargs):
val = next(responses)