From 47da67b37c544565f17388141e710a4e6530bfb6 Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Thu, 31 Jul 2025 15:47:24 -0400
Subject: [PATCH] Improve import handling

---
 src/seedpass/core/manager.py              | 54 +++++++++++++++++-----
 src/tests/test_manager_import_database.py | 55 +++++++++++++++++++++++
 2 files changed, 98 insertions(+), 11 deletions(-)
 create mode 100644 src/tests/test_manager_import_database.py

diff --git a/src/seedpass/core/manager.py b/src/seedpass/core/manager.py
index ebe832c..19e8767 100644
--- a/src/seedpass/core/manager.py
+++ b/src/seedpass/core/manager.py
@@ -32,6 +32,7 @@ from .password_generation import PasswordGenerator
 from .backup import BackupManager
 from .vault import Vault
 from .portable_backup import export_backup, import_backup
+from cryptography.fernet import InvalidToken
 from .totp import TotpManager
 from .entry_types import EntryType
 from .pubsub import bus
@@ -4013,26 +4014,57 @@ class PasswordManager:
 
     def handle_import_database(self, src: Path) -> None:
         """Import a portable database file, replacing the current index."""
-        try:
-            fp, parent_fp, child_fp = self.header_fingerprint_args
-            clear_header_with_notification(
-                self,
-                fp,
-                "Main Menu > Settings > Import database",
-                parent_fingerprint=parent_fp,
-                child_fingerprint=child_fp,
+
+        if not src.name.endswith(".json.enc"):
+            print(
+                colored(
+                    "Error: Selected file must be a SeedPass database backup (.json.enc).",
+                    "red",
+                )
             )
+            return
+
+        fp, parent_fp, child_fp = self.header_fingerprint_args
+        clear_header_with_notification(
+            self,
+            fp,
+            "Main Menu > Settings > Import database",
+            parent_fingerprint=parent_fp,
+            child_fingerprint=child_fp,
+        )
+
+        try:
             import_backup(
                 self.vault,
                 self.backup_manager,
                 src,
                 parent_seed=self.parent_seed,
             )
-            print(colored("Database imported successfully.", "green"))
-            self.sync_vault()
+        except InvalidToken:
+            logging.error("Invalid backup token during import", exc_info=True)
+            print(
+                colored(
+                    "Error: Invalid backup. Please import a file created by SeedPass.",
+                    "red",
+                )
+            )
+            return
+        except FileNotFoundError:
+            logging.error(f"Backup file not found: {src}", exc_info=True)
+            print(colored(f"Error: File '{src}' not found.", "red"))
+            return
         except Exception as e:
             logging.error(f"Failed to import database: {e}", exc_info=True)
-            print(colored(f"Error: Failed to import database: {e}", "red"))
+            print(
+                colored(
+                    f"Error: Failed to import database: {e}. Please verify the backup file.",
+                    "red",
+                )
+            )
+            return
+
+        print(colored("Database imported successfully.", "green"))
+        self.sync_vault()
 
     def handle_export_totp_codes(self) -> Path | None:
         """Export all 2FA codes to a JSON file for other authenticator apps."""
diff --git a/src/tests/test_manager_import_database.py b/src/tests/test_manager_import_database.py
new file mode 100644
index 0000000..9cfb18c
--- /dev/null
+++ b/src/tests/test_manager_import_database.py
@@ -0,0 +1,55 @@
+import queue
+from pathlib import Path
+from types import SimpleNamespace
+import sys
+
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+from seedpass.core.manager import PasswordManager, EncryptionMode
+
+
+def _make_pm() -> PasswordManager:
+    pm = PasswordManager.__new__(PasswordManager)
+    pm.encryption_mode = EncryptionMode.SEED_ONLY
+    pm.vault = SimpleNamespace()
+    pm.backup_manager = SimpleNamespace()
+    pm.parent_seed = "seed"
+    pm.profile_stack = []
+    pm.current_fingerprint = None
+    pm.sync_vault = lambda: None
+    pm.notifications = queue.Queue()
+    return pm
+
+
+def test_import_non_backup_file(monkeypatch, capsys):
+    pm = _make_pm()
+    called = {"called": False}
+
+    def fake_import(*_a, **_k):
+        called["called"] = True
+
+    monkeypatch.setattr("seedpass.core.manager.import_backup", fake_import)
+    monkeypatch.setattr(
+        "seedpass.core.manager.clear_header_with_notification", lambda *a, **k: None
+    )
+
+    pm.handle_import_database(Path("data.txt"))
+    out = capsys.readouterr().out
+    assert "json.enc" in out.lower()
+    assert called["called"] is False
+
+
+def test_import_missing_file(monkeypatch, capsys):
+    pm = _make_pm()
+
+    def raise_missing(*_a, **_k):
+        raise FileNotFoundError
+
+    monkeypatch.setattr("seedpass.core.manager.import_backup", raise_missing)
+    monkeypatch.setattr(
+        "seedpass.core.manager.clear_header_with_notification", lambda *a, **k: None
+    )
+
+    pm.handle_import_database(Path("missing.json.enc"))
+    out = capsys.readouterr().out
+    assert "not found" in out.lower()