Merge pull request #84 from PR0M3TH3AN/codex/add-encryptionmode-class-and-update-callers

Introduce EncryptionMode and derive_index_key utility

src/password_manager/manager.py

@@ -24,7 +24,12 @@ from password_manager.entry_management import EntryManager
 from password_manager.password_generation import PasswordGenerator
 from password_manager.backup import BackupManager
 from password_manager.vault import Vault
-from utils.key_derivation import derive_key_from_parent_seed, derive_key_from_password
+from utils.key_derivation import (
+    derive_key_from_parent_seed,
+    derive_key_from_password,
+    derive_index_key,
+    DEFAULT_ENCRYPTION_MODE,
+)
 from utils.checksum import calculate_checksum, verify_checksum
 from utils.password_prompt import (
     prompt_for_password,
@@ -263,8 +268,15 @@ class PasswordManager:
             # Prompt for password if not provided
             if password is None:
                 password = prompt_existing_password("Enter your master password: ")
-            # Derive key from password
-            key = derive_key_from_password(password)
+            # Derive key using the configured encryption mode if seed is known
+            if self.parent_seed:
+                key = derive_index_key(
+                    self.parent_seed,
+                    password,
+                    DEFAULT_ENCRYPTION_MODE,
+                )
+            else:
+                key = derive_key_from_password(password)
             self.encryption_manager = EncryptionManager(key, fingerprint_dir)
             self.vault = Vault(self.encryption_manager, fingerprint_dir)
             logger.debug(
@@ -513,7 +525,11 @@ class PasswordManager:
 
                 # Initialize EncryptionManager with key and fingerprint_dir
                 password = prompt_for_password()
-                key = derive_key_from_password(password)
+                key = derive_index_key(
+                    parent_seed,
+                    password,
+                    DEFAULT_ENCRYPTION_MODE,
+                )
                 self.encryption_manager = EncryptionManager(key, fingerprint_dir)
                 self.vault = Vault(self.encryption_manager, fingerprint_dir)
 
@@ -644,8 +660,12 @@ class PasswordManager:
 
             # Prompt for password
             password = prompt_for_password()
-            # Derive key from password
-            key = derive_key_from_password(password)
+            # Derive key using the configured encryption mode
+            key = derive_index_key(
+                seed,
+                password,
+                DEFAULT_ENCRYPTION_MODE,
+            )
             # Re-initialize EncryptionManager with the new key and fingerprint_dir
             self.encryption_manager = EncryptionManager(key, fingerprint_dir)
 

src/tests/test_key_derivation.py

@@ -4,6 +4,8 @@ from utils.key_derivation import (
     derive_key_from_password,
     derive_index_key_seed_only,
     derive_index_key_seed_plus_pw,
+    derive_index_key,
+    EncryptionMode,
 )
 
 
@@ -36,3 +38,17 @@ def test_seed_plus_pw_differs_from_seed_only():
     k1 = derive_index_key_seed_only(seed)
     k2 = derive_index_key_seed_plus_pw(seed, pw)
     assert k1 != k2
+
+
+def test_derive_index_key_modes():
+    seed = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
+    pw = "hunter2"
+    assert derive_index_key(
+        seed, pw, EncryptionMode.SEED_ONLY
+    ) == derive_index_key_seed_only(seed)
+    assert derive_index_key(
+        seed, pw, EncryptionMode.SEED_PLUS_PW
+    ) == derive_index_key_seed_plus_pw(seed, pw)
+    assert derive_index_key(
+        seed, pw, EncryptionMode.PW_ONLY
+    ) == derive_key_from_password(pw)

src/utils/__init__.py

@@ -5,7 +5,13 @@ import traceback
 
 try:
     from .file_lock import exclusive_lock, shared_lock
-    from .key_derivation import derive_key_from_password, derive_key_from_parent_seed
+    from .key_derivation import (
+        derive_key_from_password,
+        derive_key_from_parent_seed,
+        derive_index_key,
+        EncryptionMode,
+        DEFAULT_ENCRYPTION_MODE,
+    )
     from .checksum import calculate_checksum, verify_checksum
     from .password_prompt import prompt_for_password
 
@@ -17,6 +23,9 @@ except Exception as e:
 __all__ = [
     "derive_key_from_password",
     "derive_key_from_parent_seed",
+    "derive_index_key",
+    "EncryptionMode",
+    "DEFAULT_ENCRYPTION_MODE",
     "calculate_checksum",
     "verify_checksum",
     "exclusive_lock",

src/utils/key_derivation.py

@@ -20,7 +20,8 @@ import base64
 import unicodedata
 import logging
 import traceback
-from typing import Union
+from enum import Enum
+from typing import Optional, Union
 from bip_utils import Bip39SeedGenerator
 from local_bip85.bip85 import BIP85
 
@@ -36,6 +37,17 @@ from cryptography.hazmat.backends import default_backend
 logger = logging.getLogger(__name__)
 
 
+class EncryptionMode(Enum):
+    """Supported key derivation modes for database encryption."""
+
+    SEED_ONLY = "seed-only"
+    SEED_PLUS_PW = "seed+pw"
+    PW_ONLY = "pw-only"
+
+
+DEFAULT_ENCRYPTION_MODE = EncryptionMode.SEED_ONLY
+
+
 def derive_key_from_password(password: str, iterations: int = 100_000) -> bytes:
     """
     Derives a Fernet-compatible encryption key from the provided password using PBKDF2-HMAC-SHA256.
@@ -196,3 +208,22 @@ def derive_index_key_seed_plus_pw(seed: str, password: str) -> bytes:
     )
     key = hkdf.derive(seed_bytes + b"|" + pw_bytes)
     return base64.urlsafe_b64encode(key)
+
+
+def derive_index_key(
+    seed: str,
+    password: Optional[str] = None,
+    mode: EncryptionMode = DEFAULT_ENCRYPTION_MODE,
+) -> bytes:
+    """Derive the index encryption key based on the selected mode."""
+    if mode == EncryptionMode.SEED_ONLY:
+        return derive_index_key_seed_only(seed)
+    if mode == EncryptionMode.SEED_PLUS_PW:
+        if password is None:
+            raise ValueError("Password required for seed+pw mode")
+        return derive_index_key_seed_plus_pw(seed, password)
+    if mode == EncryptionMode.PW_ONLY:
+        if password is None:
+            raise ValueError("Password required for pw-only mode")
+        return derive_key_from_password(password)
+    raise ValueError(f"Unsupported encryption mode: {mode}")