From fe41ae4a60dbb49917b0e11e7f5899248287a9cf Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Thu, 3 Jul 2025 10:08:35 -0400
Subject: [PATCH] Handle retrieval of 2FA entries

---
 src/password_manager/manager.py         | 26 +++++++++----
 src/tests/test_manager_retrieve_totp.py | 50 +++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 7 deletions(-)
 create mode 100644 src/tests/test_manager_retrieve_totp.py

diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py
index c71be79..de3b3cc 100644
--- a/src/password_manager/manager.py
+++ b/src/password_manager/manager.py
@@ -1003,21 +1003,35 @@ class PasswordManager:
                 return
             index = int(index_input)
 
-            # Retrieve entry details
             entry = self.entry_manager.retrieve_entry(index)
             if not entry:
                 return
 
-            # Display entry details
+            entry_type = entry.get("type", EntryType.PASSWORD.value)
+
+            if entry_type == EntryType.TOTP.value:
+                label = entry.get("label", "")
+                period = int(entry.get("period", 30))
+                notes = entry.get("notes", "")
+                print(colored(f"Retrieving 2FA code for '{label}'.", "cyan"))
+                try:
+                    code = self.entry_manager.get_totp_code(index, self.parent_seed)
+                    print(colored("\n[+] Retrieved 2FA Code:\n", "green"))
+                    print(colored(f"Code: {code}", "yellow"))
+                    if notes:
+                        print(colored(f"Notes: {notes}", "cyan"))
+                    TotpManager.print_progress_bar(period)
+                except Exception as e:
+                    logging.error(f"Error generating TOTP code: {e}", exc_info=True)
+                    print(colored(f"Error: Failed to generate TOTP code: {e}", "red"))
+                return
+
             website_name = entry.get("website")
             length = entry.get("length")
             username = entry.get("username")
             url = entry.get("url")
             blacklisted = entry.get("blacklisted")
             notes = entry.get("notes", "")
-            notes = entry.get("notes", "")
-            notes = entry.get("notes", "")
-            notes = entry.get("notes", "")
 
             print(
                 colored(
@@ -1037,10 +1051,8 @@ class PasswordManager:
                     )
                 )
 
-            # Generate the password
             password = self.password_generator.generate_password(length, index)
 
-            # Display the password and associated details
             if password:
                 print(
                     colored(f"\n[+] Retrieved Password for {website_name}:\n", "green")
diff --git a/src/tests/test_manager_retrieve_totp.py b/src/tests/test_manager_retrieve_totp.py
new file mode 100644
index 0000000..ea10084
--- /dev/null
+++ b/src/tests/test_manager_retrieve_totp.py
@@ -0,0 +1,50 @@
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from helpers import create_vault, TEST_SEED, TEST_PASSWORD
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+from password_manager.entry_management import EntryManager
+from password_manager.backup import BackupManager
+from password_manager.manager import PasswordManager, EncryptionMode, TotpManager
+
+
+class FakeNostrClient:
+    def __init__(self, *args, **kwargs):
+        self.published = []
+
+    def publish_snapshot(self, data: bytes):
+        self.published.append(data)
+        return None, "abcd"
+
+
+def test_handle_retrieve_totp_entry(monkeypatch, capsys):
+    with TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        vault, enc_mgr = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD)
+        entry_mgr = EntryManager(vault, tmp_path)
+        backup_mgr = BackupManager(tmp_path)
+
+        pm = PasswordManager.__new__(PasswordManager)
+        pm.encryption_mode = EncryptionMode.SEED_ONLY
+        pm.encryption_manager = enc_mgr
+        pm.vault = vault
+        pm.entry_manager = entry_mgr
+        pm.backup_manager = backup_mgr
+        pm.parent_seed = TEST_SEED
+        pm.nostr_client = FakeNostrClient()
+        pm.fingerprint_dir = tmp_path
+        pm.is_dirty = False
+
+        entry_mgr.add_totp("Example", TEST_SEED)
+
+        monkeypatch.setattr("builtins.input", lambda *a, **k: "0")
+        monkeypatch.setattr(pm.entry_manager, "get_totp_code", lambda *a, **k: "123456")
+        monkeypatch.setattr(TotpManager, "print_progress_bar", lambda period: None)
+
+        pm.handle_retrieve_entry()
+        out = capsys.readouterr().out
+        assert "Retrieved 2FA Code" in out
+        assert "123456" in out