From f49daca4df273ca1ee631298a4eff8222382e91f Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Tue, 1 Jul 2025 11:01:22 -0400 Subject: [PATCH] Add EncryptionMode enum and integrate index key derivation --- src/password_manager/manager.py | 32 +++++++++++++++++++++++++------ src/tests/test_key_derivation.py | 16 ++++++++++++++++ src/utils/__init__.py | 11 ++++++++++- src/utils/key_derivation.py | 33 +++++++++++++++++++++++++++++++- 4 files changed, 84 insertions(+), 8 deletions(-) diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index 8fea00a..d53e4b2 100644 --- a/src/password_manager/manager.py +++ b/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) diff --git a/src/tests/test_key_derivation.py b/src/tests/test_key_derivation.py index f734eb6..06b5f6a 100644 --- a/src/tests/test_key_derivation.py +++ b/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) diff --git a/src/utils/__init__.py b/src/utils/__init__.py index e12cbff..dbde68e 100644 --- a/src/utils/__init__.py +++ b/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", diff --git a/src/utils/key_derivation.py b/src/utils/key_derivation.py index 3a9065f..eb022b2 100644 --- a/src/utils/key_derivation.py +++ b/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}")