From 467bdd3d04e2b34ed73cc75d6492bd9e2498e7e3 Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Mon, 7 Jul 2025 15:19:42 -0400
Subject: [PATCH] Add notes support for TOTP entries

---
 src/password_manager/entry_management.py |  3 +++
 src/password_manager/manager.py          |  4 ++++
 src/tests/test_manager_add_totp.py       |  2 ++
 src/tests/test_totp_entry.py             | 13 +++++++++++++
 4 files changed, 22 insertions(+)

diff --git a/src/password_manager/entry_management.py b/src/password_manager/entry_management.py
index 7470c10..c63731f 100644
--- a/src/password_manager/entry_management.py
+++ b/src/password_manager/entry_management.py
@@ -195,6 +195,7 @@ class EntryManager:
         index: int | None = None,
         period: int = 30,
         digits: int = 6,
+        notes: str = "",
     ) -> str:
         """Add a new TOTP entry and return the provisioning URI."""
         entry_id = self.get_next_index()
@@ -213,6 +214,7 @@ class EntryManager:
                 "period": period,
                 "digits": digits,
                 "archived": archived,
+                "notes": notes,
             }
         else:
             entry = {
@@ -223,6 +225,7 @@ class EntryManager:
                 "period": period,
                 "digits": digits,
                 "archived": archived,
+                "notes": notes,
             }
 
         data["entries"][str(entry_id)] = entry
diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py
index 7ac1cd8..f694ab0 100644
--- a/src/password_manager/manager.py
+++ b/src/password_manager/manager.py
@@ -1001,6 +1001,7 @@ class PasswordManager:
                             colored("Error: Period and digits must be numbers.", "red")
                         )
                         continue
+                    notes = input("Notes (optional): ").strip()
                     totp_index = self.entry_manager.get_next_totp_index()
                     entry_id = self.entry_manager.get_next_index()
                     uri = self.entry_manager.add_totp(
@@ -1009,6 +1010,7 @@ class PasswordManager:
                         index=totp_index,
                         period=int(period),
                         digits=int(digits),
+                        notes=notes,
                     )
                     secret = TotpManager.derive_secret(self.parent_seed, totp_index)
                     self.is_dirty = True
@@ -1043,6 +1045,7 @@ class PasswordManager:
                             secret = raw.upper()
                             period = int(input("Period (default 30): ").strip() or 30)
                             digits = int(input("Digits (default 6): ").strip() or 6)
+                        notes = input("Notes (optional): ").strip()
                         entry_id = self.entry_manager.get_next_index()
                         uri = self.entry_manager.add_totp(
                             label,
@@ -1050,6 +1053,7 @@ class PasswordManager:
                             secret=secret,
                             period=period,
                             digits=digits,
+                            notes=notes,
                         )
                         self.is_dirty = True
                         self.last_update = time.time()
diff --git a/src/tests/test_manager_add_totp.py b/src/tests/test_manager_add_totp.py
index 8ad618b..0038511 100644
--- a/src/tests/test_manager_add_totp.py
+++ b/src/tests/test_manager_add_totp.py
@@ -47,6 +47,7 @@ def test_handle_add_totp(monkeypatch, capsys):
                 "Example",  # label
                 "",  # period
                 "",  # digits
+                "",  # notes
             ]
         )
         monkeypatch.setattr("builtins.input", lambda *args, **kwargs: next(inputs))
@@ -64,5 +65,6 @@ def test_handle_add_totp(monkeypatch, capsys):
             "period": 30,
             "digits": 6,
             "archived": False,
+            "notes": "",
         }
         assert "ID 0" in out
diff --git a/src/tests/test_totp_entry.py b/src/tests/test_totp_entry.py
index 87f24da..b25d610 100644
--- a/src/tests/test_totp_entry.py
+++ b/src/tests/test_totp_entry.py
@@ -36,6 +36,7 @@ def test_add_totp_and_get_code():
             "period": 30,
             "digits": 6,
             "archived": False,
+            "notes": "",
         }
 
         code = entry_mgr.get_totp_code(0, TEST_SEED, timestamp=0)
@@ -74,6 +75,18 @@ def test_add_totp_imported(tmp_path):
         "period": 30,
         "digits": 6,
         "archived": False,
+        "notes": "",
     }
     code = em.get_totp_code(0, timestamp=0)
     assert code == pyotp.TOTP(secret).at(0)
+
+
+def test_add_totp_with_notes(tmp_path):
+    vault, _ = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD)
+    cfg_mgr = ConfigManager(vault, tmp_path)
+    backup_mgr = BackupManager(tmp_path, cfg_mgr)
+    em = EntryManager(vault, backup_mgr)
+
+    em.add_totp("NoteLabel", TEST_SEED, notes="some note")
+    entry = em.retrieve_entry(0)
+    assert entry["notes"] == "some note"