From 7fff164da72cfbfe79773a8e814a457bccdcb9c4 Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Wed, 2 Jul 2025 22:28:56 -0400 Subject: [PATCH] Update entry manager for generic entries --- src/password_manager/entry_management.py | 72 +++++++++++++++++++----- src/tests/test_entry_add.py | 2 + 2 files changed, 60 insertions(+), 14 deletions(-) diff --git a/src/password_manager/entry_management.py b/src/password_manager/entry_management.py index 00c07c3..911b854 100644 --- a/src/password_manager/entry_management.py +++ b/src/password_manager/entry_management.py @@ -28,6 +28,7 @@ from pathlib import Path from termcolor import colored from password_manager.migrations import LATEST_VERSION +from password_manager.entry_types import EntryType from password_manager.vault import Vault from utils.file_lock import exclusive_lock @@ -65,7 +66,7 @@ class EntryManager: return {"schema_version": LATEST_VERSION, "entries": {}} else: logger.info( - f"Index file '{self.index_file}' not found. Initializing new password database." + f"Index file '{self.index_file}' not found. Initializing new entries database." ) return {"schema_version": LATEST_VERSION, "entries": {}} @@ -79,7 +80,7 @@ class EntryManager: def get_next_index(self) -> int: """ - Retrieves the next available index for a new password entry. + Retrieves the next available index for a new entry. :return: The next index number as an integer. """ @@ -104,15 +105,17 @@ class EntryManager: username: Optional[str] = None, url: Optional[str] = None, blacklisted: bool = False, + notes: str = "", ) -> int: """ - Adds a new password entry to the encrypted JSON index file. + Adds a new entry to the encrypted JSON index file. :param website_name: The name of the website. :param length: The desired length of the password. :param username: (Optional) The username associated with the website. :param url: (Optional) The URL of the website. :param blacklisted: (Optional) Whether the password is blacklisted. Defaults to False. + :param notes: (Optional) Extra notes to attach to the entry. :return: The assigned index of the new entry. """ try: @@ -126,6 +129,8 @@ class EntryManager: "username": username if username else "", "url": url if url else "", "blacklisted": blacklisted, + "type": EntryType.PASSWORD.value, + "notes": notes, } logger.debug(f"Added entry at index {index}: {data['entries'][str(index)]}") @@ -144,6 +149,39 @@ class EntryManager: print(colored(f"Error: Failed to add entry: {e}", "red")) sys.exit(1) + def add_totp(self, notes: str = "") -> int: + """Placeholder for adding a TOTP entry.""" + index = self.get_next_index() + data = self.vault.load_index() + data.setdefault("entries", {}) + data["entries"][str(index)] = {"type": EntryType.TOTP.value, "notes": notes} + self._save_index(data) + self.update_checksum() + self.backup_index_file() + raise NotImplementedError("TOTP entry support not implemented yet") + + def add_ssh_key(self, notes: str = "") -> int: + """Placeholder for adding an SSH key entry.""" + index = self.get_next_index() + data = self.vault.load_index() + data.setdefault("entries", {}) + data["entries"][str(index)] = {"type": EntryType.SSH.value, "notes": notes} + self._save_index(data) + self.update_checksum() + self.backup_index_file() + raise NotImplementedError("SSH key entry support not implemented yet") + + def add_seed(self, notes: str = "") -> int: + """Placeholder for adding a seed entry.""" + index = self.get_next_index() + data = self.vault.load_index() + data.setdefault("entries", {}) + data["entries"][str(index)] = {"type": EntryType.SEED.value, "notes": notes} + self._save_index(data) + self.update_checksum() + self.backup_index_file() + raise NotImplementedError("Seed entry support not implemented yet") + def get_encrypted_index(self) -> Optional[bytes]: """ Retrieves the encrypted password index file's contents. @@ -161,9 +199,9 @@ class EntryManager: def retrieve_entry(self, index: int) -> Optional[Dict[str, Any]]: """ - Retrieves a password entry based on the provided index. + Retrieves an entry based on the provided index. - :param index: The index number of the password entry. + :param index: The index number of the entry. :return: A dictionary containing the entry details or None if not found. """ try: @@ -193,14 +231,16 @@ class EntryManager: username: Optional[str] = None, url: Optional[str] = None, blacklisted: Optional[bool] = None, + notes: Optional[str] = None, ) -> None: """ - Modifies an existing password entry based on the provided index and new values. + Modifies an existing entry based on the provided index and new values. - :param index: The index number of the password entry to modify. + :param index: The index number of the entry to modify. :param username: (Optional) The new username. :param url: (Optional) The new URL. :param blacklisted: (Optional) The new blacklist status. + :param notes: (Optional) New notes to attach to the entry. """ try: data = self.vault.load_index() @@ -232,6 +272,10 @@ class EntryManager: f"Updated blacklist status to '{blacklisted}' for index {index}." ) + if notes is not None: + entry["notes"] = notes + logger.debug(f"Updated notes for index {index}.") + data["entries"][str(index)] = entry logger.debug(f"Modified entry at index {index}: {entry}") @@ -252,7 +296,7 @@ class EntryManager: def list_entries(self) -> List[Tuple[int, str, Optional[str], Optional[str], bool]]: """ - Lists all password entries in the index. + Lists all entries in the index. :return: A list of tuples containing entry details: (index, website, username, url, blacklisted) """ @@ -261,8 +305,8 @@ class EntryManager: entries_data = data.get("entries", {}) if not entries_data: - logger.info("No password entries found.") - print(colored("No password entries found.", "yellow")) + logger.info("No entries found.") + print(colored("No entries found.", "yellow")) return [] entries = [] @@ -295,9 +339,9 @@ class EntryManager: def delete_entry(self, index: int) -> None: """ - Deletes a password entry based on the provided index. + Deletes an entry based on the provided index. - :param index: The index number of the password entry to delete. + :param index: The index number of the entry to delete. """ try: data = self.vault.load_index() @@ -424,7 +468,7 @@ class EntryManager: def list_all_entries(self) -> None: """ - Displays all password entries in a formatted manner. + Displays all entries in a formatted manner. """ try: entries = self.list_entries() @@ -432,7 +476,7 @@ class EntryManager: print(colored("No entries to display.", "yellow")) return - print(colored("\n[+] Listing All Password Entries:\n", "green")) + print(colored("\n[+] Listing All Entries:\n", "green")) for entry in entries: index, website, username, url, blacklisted = entry print(colored(f"Index: {index}", "cyan")) diff --git a/src/tests/test_entry_add.py b/src/tests/test_entry_add.py index 08d53e4..eac6db8 100644 --- a/src/tests/test_entry_add.py +++ b/src/tests/test_entry_add.py @@ -23,6 +23,8 @@ def test_add_and_retrieve_entry(): "username": "user", "url": "", "blacklisted": False, + "type": "password", + "notes": "", } data = enc_mgr.load_json_data(entry_mgr.index_file)