Merge pull request #483 from PR0M3TH3AN/codex/avoid-storing-local-data-for-incomplete-profiles

Fix partial profile cleanup

src/password_manager/manager.py

@@ -719,39 +719,46 @@ class PasswordManager:
                 self.fingerprint_dir = fingerprint_dir
                 logging.info(f"Current seed profile set to {fingerprint}")
 
-                # Initialize EncryptionManager with key and fingerprint_dir
-                password = prompt_for_password()
-                index_key = derive_index_key(parent_seed)
-                iterations = self.config_manager.get_kdf_iterations()
-                seed_key = derive_key_from_password(password, iterations=iterations)
+                try:
+                    # Initialize EncryptionManager with key and fingerprint_dir
+                    password = prompt_for_password()
+                    index_key = derive_index_key(parent_seed)
+                    iterations = self.config_manager.get_kdf_iterations()
+                    seed_key = derive_key_from_password(password, iterations=iterations)
 
-                self.encryption_manager = EncryptionManager(index_key, fingerprint_dir)
-                seed_mgr = EncryptionManager(seed_key, fingerprint_dir)
-                self.vault = Vault(self.encryption_manager, fingerprint_dir)
+                    self.encryption_manager = EncryptionManager(
+                        index_key, fingerprint_dir
+                    )
+                    seed_mgr = EncryptionManager(seed_key, fingerprint_dir)
+                    self.vault = Vault(self.encryption_manager, fingerprint_dir)
 
-                # Ensure config manager is set for the new fingerprint
-                self.config_manager = ConfigManager(
-                    vault=self.vault,
-                    fingerprint_dir=fingerprint_dir,
-                )
+                    # Ensure config manager is set for the new fingerprint
+                    self.config_manager = ConfigManager(
+                        vault=self.vault,
+                        fingerprint_dir=fingerprint_dir,
+                    )
 
-                # Encrypt and save the parent seed
-                seed_mgr.encrypt_parent_seed(parent_seed)
-                logging.info("Parent seed encrypted and saved successfully.")
+                    # Encrypt and save the parent seed
+                    seed_mgr.encrypt_parent_seed(parent_seed)
+                    logging.info("Parent seed encrypted and saved successfully.")
 
-                # Store the hashed password
-                self.store_hashed_password(password)
-                logging.info("User password hashed and stored successfully.")
+                    # Store the hashed password
+                    self.store_hashed_password(password)
+                    logging.info("User password hashed and stored successfully.")
 
-                self.parent_seed = parent_seed  # Ensure this is a string
-                logger.debug(
-                    f"parent_seed set to: {self.parent_seed} (type: {type(self.parent_seed)})"
-                )
+                    self.parent_seed = parent_seed  # Ensure this is a string
+                    logger.debug(
+                        f"parent_seed set to: {self.parent_seed} (type: {type(self.parent_seed)})"
+                    )
 
-                self.initialize_bip85()
-                self.initialize_managers()
-                self.sync_index_from_nostr()
-                return fingerprint  # Return the generated or added fingerprint
+                    self.initialize_bip85()
+                    self.initialize_managers()
+                    self.sync_index_from_nostr()
+                    return fingerprint  # Return the generated or added fingerprint
+                except BaseException:
+                    # Clean up partial profile on failure or interruption
+                    self.fingerprint_manager.remove_fingerprint(fingerprint)
+                    raise
             else:
                 logging.error("Invalid BIP-85 seed phrase. Exiting.")
                 print(colored("Error: Invalid BIP-85 seed phrase.", "red"))
@@ -800,7 +807,12 @@ class PasswordManager:
             logging.info(f"Current seed profile set to {fingerprint}")
 
             # Now, save and encrypt the seed with the fingerprint_dir
-            self.save_and_encrypt_seed(new_seed, fingerprint_dir)
+            try:
+                self.save_and_encrypt_seed(new_seed, fingerprint_dir)
+            except BaseException:
+                # Clean up partial profile on failure or interruption
+                self.fingerprint_manager.remove_fingerprint(fingerprint)
+                raise
 
             return fingerprint  # Return the generated fingerprint
         else:

src/tests/test_profile_cleanup.py

@@ -0,0 +1,49 @@
+import sys
+import importlib
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import pytest
+from unittest.mock import patch
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+
+def setup_pm(tmp_path):
+    import constants
+    import password_manager.manager as manager_module
+
+    importlib.reload(constants)
+    importlib.reload(manager_module)
+
+    pm = manager_module.PasswordManager.__new__(manager_module.PasswordManager)
+    pm.encryption_mode = manager_module.EncryptionMode.SEED_ONLY
+    pm.fingerprint_manager = manager_module.FingerprintManager(constants.APP_DIR)
+    pm.current_fingerprint = None
+    return pm, constants, manager_module
+
+
+def test_generate_seed_cleanup_on_failure(monkeypatch):
+    with TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        monkeypatch.setattr(Path, "home", lambda: tmp_path)
+
+        pm, const, mgr = setup_pm(tmp_path)
+
+        with patch("password_manager.manager.confirm_action", return_value=True):
+            monkeypatch.setattr(
+                pm,
+                "save_and_encrypt_seed",
+                lambda seed, d: (_ for _ in ()).throw(RuntimeError("fail")),
+            )
+            with pytest.raises(RuntimeError):
+                pm.generate_new_seed()
+
+        # fingerprint list should be empty and only fingerprints.json should remain
+        assert pm.fingerprint_manager.list_fingerprints() == []
+        contents = list(const.APP_DIR.iterdir())
+        assert len(contents) == 1 and contents[0].name == "fingerprints.json"
+        fp_file = pm.fingerprint_manager.fingerprints_file
+        with open(fp_file) as f:
+            data = json.load(f)
+        assert data.get("fingerprints") == []