Merge pull request #830 from PR0M3TH3AN/codex/implement-hkdf-helper-and-update-sub-key-usage

Refactor key derivation with hierarchical HKDF

src/main.py

@@ -1415,9 +1415,10 @@ def main(argv: list[str] | None = None, *, fingerprint: str | None = None) -> in
         if entry.get("type") != EntryType.TOTP.value:
             print(colored("Entry is not a TOTP entry.", "red"))
             return 1
-        code = password_manager.entry_manager.get_totp_code(
-            idx, password_manager.parent_seed
+        key = getattr(password_manager, "KEY_TOTP_DET", None) or getattr(
+            password_manager, "parent_seed", None
         )
+        code = password_manager.entry_manager.get_totp_code(idx, key)
         print(code)
         try:
             if copy_to_clipboard(code, password_manager.clipboard_clear_delay):

src/seedpass/api.py

@@ -464,7 +464,8 @@ def export_totp(
     _check_token(request, authorization)
     _require_password(request, password)
     pm = _get_pm(request)
-    return pm.entry_manager.export_totp_entries(pm.parent_seed)
+    key = getattr(pm, "KEY_TOTP_DET", None) or getattr(pm, "parent_seed", None)
+    return pm.entry_manager.export_totp_entries(key)
 
 
 @app.get("/api/v1/totp")
@@ -482,7 +483,8 @@ def get_totp_codes(
     )
     codes = []
     for idx, label, _u, _url, _arch in entries:
-        code = pm.entry_manager.get_totp_code(idx, pm.parent_seed)
+        key = getattr(pm, "KEY_TOTP_DET", None) or getattr(pm, "parent_seed", None)
+        code = pm.entry_manager.get_totp_code(idx, key)
 
         rem = pm.entry_manager.get_totp_time_remaining(idx)
 

src/seedpass/core/api.py

@@ -305,9 +305,10 @@ class EntryService:
 
     def get_totp_code(self, entry_id: int) -> str:
         with self._lock:
-            return self._manager.entry_manager.get_totp_code(
-                entry_id, self._manager.parent_seed
+            key = getattr(self._manager, "KEY_TOTP_DET", None) or getattr(
+                self._manager, "parent_seed", None
             )
+            return self._manager.entry_manager.get_totp_code(entry_id, key)
 
     def add_entry(
         self,
@@ -515,9 +516,10 @@ class EntryService:
 
     def export_totp_entries(self) -> dict:
         with self._lock:
-            return self._manager.entry_manager.export_totp_entries(
-                self._manager.parent_seed
+            key = getattr(self._manager, "KEY_TOTP_DET", None) or getattr(
+                self._manager, "parent_seed", None
             )
+            return self._manager.entry_manager.export_totp_entries(key)
 
     def display_totp_codes(self) -> None:
         with self._lock:

src/seedpass/core/entry_management.py

@@ -257,7 +257,7 @@ class EntryManager:
     def add_totp(
         self,
         label: str,
-        parent_seed: str,
+        parent_seed: str | bytes,
         *,
         archived: bool = False,
         secret: str | None = None,
@@ -689,7 +689,10 @@ class EntryManager:
         return derive_seed_phrase(bip85, seed_index, words)
 
     def get_totp_code(
-        self, index: int, parent_seed: str | None = None, timestamp: int | None = None
+        self,
+        index: int,
+        parent_seed: str | bytes | None = None,
+        timestamp: int | None = None,
     ) -> str:
         """Return the current TOTP code for the specified entry."""
         entry = self.retrieve_entry(index)
@@ -719,7 +722,9 @@ class EntryManager:
         period = int(entry.get("period", 30))
         return TotpManager.time_remaining(period)
 
-    def export_totp_entries(self, parent_seed: str) -> dict[str, list[dict[str, Any]]]:
+    def export_totp_entries(
+        self, parent_seed: str | bytes
+    ) -> dict[str, list[dict[str, Any]]]:
         """Return all TOTP secrets and metadata for external use."""
         data = self._load_index()
         entries = data.get("entries", {})

src/seedpass/core/manager.py

@@ -49,6 +49,7 @@ from utils.key_derivation import (
     EncryptionMode,
     KdfConfig,
 )
+from utils.key_hierarchy import kd
 from utils.checksum import (
     calculate_checksum,
     verify_checksum,
@@ -232,6 +233,13 @@ class PasswordManager:
     verification, ensuring the integrity and confidentiality of the stored password database.
     """
 
+    # Class-level fallbacks so attributes exist even if ``__init__`` is bypassed
+    master_key: bytes | None = None
+    KEY_STORAGE: bytes | None = None
+    KEY_INDEX: bytes | None = None
+    KEY_PW_DERIVE: bytes | None = None
+    KEY_TOTP_DET: bytes | None = None
+
     def __init__(
         self, fingerprint: Optional[str] = None, *, password: Optional[str] = None
     ) -> None:
@@ -264,6 +272,13 @@ class PasswordManager:
         self._bip85_cache: dict[tuple[int, int], bytes] = {}
         self.audit_logger: Optional[AuditLogger] = None
 
+        # Derived key hierarchy
+        self.master_key: bytes | None = None
+        self.KEY_STORAGE: bytes | None = None
+        self.KEY_INDEX: bytes | None = None
+        self.KEY_PW_DERIVE: bytes | None = None
+        self.KEY_TOTP_DET: bytes | None = None
+
         # Track changes to trigger periodic Nostr sync
         self.is_dirty: bool = False
         self.last_update: float = time.time()
@@ -324,6 +339,30 @@ class PasswordManager:
 
         self._bip85_cache.clear()
 
+    def derive_key_hierarchy(self, seed_bytes: bytes) -> None:
+        """Populate sub-keys from ``seed_bytes`` using HKDF."""
+
+        master = kd(seed_bytes, b"seedpass:v1:master")
+        self.master_key = master
+        self.KEY_STORAGE = kd(master, b"seedpass:v1:storage")
+        self.KEY_INDEX = kd(master, b"seedpass:v1:index")
+        self.KEY_PW_DERIVE = kd(master, b"seedpass:v1:pw")
+        self.KEY_TOTP_DET = kd(master, b"seedpass:v1:totp")
+
+    def ensure_key_hierarchy(self) -> None:
+        """Ensure sub-keys are derived from the current parent seed."""
+        if (
+            self.KEY_STORAGE is None
+            or self.KEY_INDEX is None
+            or self.KEY_PW_DERIVE is None
+            or self.KEY_TOTP_DET is None
+        ) and getattr(self, "parent_seed", None):
+            try:
+                seed_bytes = Bip39SeedGenerator(self.parent_seed).Generate()
+            except Exception:
+                seed_bytes = hashlib.sha256(self.parent_seed.encode()).digest()
+            self.derive_key_hierarchy(seed_bytes)
+
     def ensure_script_checksum(self) -> None:
         """Initialize or verify the checksum of the manager script."""
         script_path = Path(__file__).resolve()
@@ -481,15 +520,12 @@ class PasswordManager:
         self.setup_encryption_manager(self.fingerprint_dir, password)
         self.initialize_bip85()
         self.initialize_managers()
+        self.ensure_key_hierarchy()
         self.is_locked = False
         self.locked = False
         self.update_activity()
-        if (
-            getattr(self, "audit_logger", None) is None
-            and getattr(self, "_parent_seed_secret", None) is not None
-        ):
-            key = hashlib.sha256(self.parent_seed.encode("utf-8")).digest()
-            self.audit_logger = AuditLogger(key)
+        if getattr(self, "audit_logger", None) is None and self.KEY_INDEX is not None:
+            self.audit_logger = AuditLogger(self.KEY_INDEX)
         if (
             getattr(self, "config_manager", None)
             and self.config_manager.get_quick_unlock()
@@ -720,9 +756,10 @@ class PasswordManager:
                     password = None
                     continue
 
-                key = derive_index_key(self.parent_seed)
-
-                self.encryption_manager = EncryptionManager(key, fingerprint_dir)
+                seed_bytes = Bip39SeedGenerator(self.parent_seed).Generate()
+                self.derive_key_hierarchy(seed_bytes)
+                key_b64 = base64.urlsafe_b64encode(self.KEY_STORAGE)
+                self.encryption_manager = EncryptionManager(key_b64, fingerprint_dir)
                 self.vault = Vault(self.encryption_manager, fingerprint_dir)
 
                 self.config_manager = ConfigManager(
@@ -783,6 +820,7 @@ class PasswordManager:
             seed_mgr = EncryptionManager(seed_key, fingerprint_dir)
             self.parent_seed = seed_mgr.decrypt_parent_seed()
             seed_bytes = Bip39SeedGenerator(self.parent_seed).Generate()
+            self.derive_key_hierarchy(seed_bytes)
             self.bip85 = BIP85(seed_bytes)
         except Exception as e:
             logger.error(f"Failed to load parent seed: {e}", exc_info=True)
@@ -812,8 +850,10 @@ class PasswordManager:
         self.fingerprint_dir = account_dir
         self.parent_seed = seed
 
-        key = derive_index_key(seed)
-        self.encryption_manager = EncryptionManager(key, account_dir)
+        seed_bytes = Bip39SeedGenerator(seed).Generate()
+        self.derive_key_hierarchy(seed_bytes)
+        key_b64 = base64.urlsafe_b64encode(self.KEY_STORAGE)
+        self.encryption_manager = EncryptionManager(key_b64, account_dir)
         self.vault = Vault(self.encryption_manager, account_dir)
 
         self.initialize_bip85()
@@ -832,9 +872,13 @@ class PasswordManager:
         self.current_fingerprint = fp
         self.fingerprint_dir = path
         self.parent_seed = seed
-
-        key = derive_index_key(seed)
-        self.encryption_manager = EncryptionManager(key, path)
+        try:
+            seed_bytes = Bip39SeedGenerator(seed).Generate()
+            self.derive_key_hierarchy(seed_bytes)
+            key_b64 = base64.urlsafe_b64encode(self.KEY_STORAGE)
+        except Exception:
+            key_b64 = derive_index_key(seed)
+        self.encryption_manager = EncryptionManager(key_b64, path)
         self.vault = Vault(self.encryption_manager, path)
 
         self.initialize_bip85()
@@ -900,10 +944,14 @@ class PasswordManager:
                 password, selected_fingerprint, iterations=iterations
             )
 
-            # Initialize EncryptionManager with key and fingerprint_dir
-            self.encryption_manager = EncryptionManager(key, fingerprint_dir)
+            seed_mgr = EncryptionManager(key, fingerprint_dir)
+            self.vault = Vault(seed_mgr, fingerprint_dir)
+            self.parent_seed = seed_mgr.decrypt_parent_seed()
+            seed_bytes = Bip39SeedGenerator(self.parent_seed).Generate()
+            self.derive_key_hierarchy(seed_bytes)
+            key_b64 = base64.urlsafe_b64encode(self.KEY_STORAGE)
+            self.encryption_manager = EncryptionManager(key_b64, fingerprint_dir)
             self.vault = Vault(self.encryption_manager, fingerprint_dir)
-            self.parent_seed = self.encryption_manager.decrypt_parent_seed()
 
             # Log the type and content of parent_seed
             logger.debug(
@@ -1045,7 +1093,9 @@ class PasswordManager:
             try:
                 if password is None:
                     password = prompt_for_password()
-                index_key = derive_index_key(parent_seed)
+                seed_bytes = Bip39SeedGenerator(parent_seed).Generate()
+                self.derive_key_hierarchy(seed_bytes)
+                index_key = base64.urlsafe_b64encode(self.KEY_STORAGE)
                 iterations = (
                     self.config_manager.get_kdf_iterations()
                     if getattr(self, "config_manager", None)
@@ -1224,7 +1274,9 @@ class PasswordManager:
             if password is None:
                 password = prompt_for_password()
 
-            index_key = derive_index_key(seed)
+            seed_bytes = Bip39SeedGenerator(seed).Generate()
+            self.derive_key_hierarchy(seed_bytes)
+            index_key = base64.urlsafe_b64encode(self.KEY_STORAGE)
             iterations = (
                 self.config_manager.get_kdf_iterations()
                 if getattr(self, "config_manager", None)
@@ -1270,6 +1322,7 @@ class PasswordManager:
         """
         try:
             seed_bytes = Bip39SeedGenerator(self.parent_seed).Generate()
+            self.derive_key_hierarchy(seed_bytes)
             self.bip85 = BIP85(seed_bytes)
             self._bip85_cache = {}
             orig_derive = self.bip85.derive_entropy
@@ -1302,6 +1355,9 @@ class PasswordManager:
             if not self.encryption_manager:
                 raise ValueError("EncryptionManager is not initialized.")
 
+            # Derive sub-keys if needed
+            self.ensure_key_hierarchy()
+
             # Reinitialize the managers with the updated EncryptionManager and current fingerprint context
             self.config_manager = ConfigManager(
                 vault=self.vault,
@@ -1334,10 +1390,11 @@ class PasswordManager:
                 backup_manager=self.backup_manager,
             )
 
+            pw_bip85 = BIP85(self.KEY_PW_DERIVE)
             self.password_generator = PasswordGenerator(
                 encryption_manager=self.encryption_manager,
-                parent_seed=self.parent_seed,
-                bip85=self.bip85,
+                parent_seed=self.KEY_PW_DERIVE,
+                bip85=pw_bip85,
                 policy=self.config_manager.get_password_policy(),
             )
 
@@ -1821,16 +1878,17 @@ class PasswordManager:
                     )
                     totp_index = self.entry_manager.get_next_totp_index()
                     entry_id = self.entry_manager.get_next_index()
+                    key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None)
                     uri = self.entry_manager.add_totp(
                         label,
-                        self.parent_seed,
+                        key,
                         index=totp_index,
                         period=int(period),
                         digits=int(digits),
                         notes=notes,
                         tags=tags,
                     )
-                    secret = TotpManager.derive_secret(self.parent_seed, totp_index)
+                    secret = TotpManager.derive_secret(key, totp_index)
                     self.is_dirty = True
                     self.last_update = time.time()
                     print(
@@ -1873,9 +1931,10 @@ class PasswordManager:
                             else []
                         )
                         entry_id = self.entry_manager.get_next_index()
+                        key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None)
                         uri = self.entry_manager.add_totp(
                             label,
-                            self.parent_seed,
+                            key,
                             secret=secret,
                             period=period,
                             digits=digits,
@@ -2637,7 +2696,8 @@ class PasswordManager:
             print(colored("Press Enter to return to the menu.", "cyan"))
             try:
                 while True:
-                    code = self.entry_manager.get_totp_code(index, self.parent_seed)
+                    key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None)
+                    code = self.entry_manager.get_totp_code(index, key)
                     if self.secret_mode_enabled:
                         if copy_to_clipboard(code, self.clipboard_clear_delay):
                             print(
@@ -4114,6 +4174,7 @@ class PasswordManager:
     def handle_export_totp_codes(self) -> Path | None:
         """Export all 2FA codes to a JSON file for other authenticator apps."""
         try:
+            self.ensure_key_hierarchy()
             fp, parent_fp, child_fp = self.header_fingerprint_args
             clear_header_with_notification(
                 self,
@@ -4135,7 +4196,8 @@ class PasswordManager:
                         secret = entry["secret"]
                     else:
                         idx = int(entry.get("index", 0))
-                        secret = TotpManager.derive_secret(self.parent_seed, idx)
+                        key = self.KEY_TOTP_DET or getattr(self, "parent_seed", None)
+                        secret = TotpManager.derive_secret(key, idx)
                     uri = TotpManager.make_otpauth_uri(label, secret, period, digits)
                     totp_entries.append(
                         {
@@ -4372,6 +4434,7 @@ class PasswordManager:
     def change_password(self, old_password: str, new_password: str) -> None:
         """Change the master password used for encryption."""
         try:
+            self.ensure_key_hierarchy()
             if not self.verify_password(old_password):
                 raise ValueError("Incorrect password")
 
@@ -4380,7 +4443,7 @@ class PasswordManager:
             config_data = self.config_manager.load_config(require_pin=False)
 
             # Create a new encryption manager with the new password
-            new_key = derive_index_key(self.parent_seed)
+            new_key = base64.urlsafe_b64encode(self.KEY_STORAGE)
 
             iterations = self.config_manager.get_kdf_iterations()
             seed_key = derive_key_from_password(

src/seedpass/core/menu_handler.py

@@ -131,7 +131,10 @@ class MenuHandler:
                 if generated:
                     print(colored("\nGenerated 2FA Codes:", "green"))
                     for label, idx, period, _ in generated:
-                        code = pm.entry_manager.get_totp_code(idx, pm.parent_seed)
+                        key = getattr(pm, "KEY_TOTP_DET", None) or getattr(
+                            pm, "parent_seed", None
+                        )
+                        code = pm.entry_manager.get_totp_code(idx, key)
                         remaining = pm.entry_manager.get_totp_time_remaining(idx)
                         filled = int(20 * (period - remaining) / period)
                         bar = "[" + "#" * filled + "-" * (20 - filled) + "]"
@@ -149,7 +152,10 @@ class MenuHandler:
                 if imported_list:
                     print(colored("\nImported 2FA Codes:", "green"))
                     for label, idx, period, _ in imported_list:
-                        code = pm.entry_manager.get_totp_code(idx, pm.parent_seed)
+                        key = getattr(pm, "KEY_TOTP_DET", None) or getattr(
+                            pm, "parent_seed", None
+                        )
+                        code = pm.entry_manager.get_totp_code(idx, key)
                         remaining = pm.entry_manager.get_totp_time_remaining(idx)
                         filled = int(20 * (period - remaining) / period)
                         bar = "[" + "#" * filled + "-" * (20 - filled) + "]"

src/seedpass/core/password_generation.py

@@ -113,10 +113,12 @@ class PasswordGenerator:
             self.bip85 = bip85
             self.policy = policy or PasswordPolicy()
 
-            # Derive seed bytes from parent_seed using BIP39 (handled by EncryptionManager)
-            self.seed_bytes = self.encryption_manager.derive_seed_from_mnemonic(
-                self.parent_seed
-            )
+            if isinstance(parent_seed, (bytes, bytearray)):
+                self.seed_bytes = bytes(parent_seed)
+            else:
+                self.seed_bytes = self.encryption_manager.derive_seed_from_mnemonic(
+                    self.parent_seed
+                )
 
             logger.debug("PasswordGenerator initialized successfully.")
         except Exception as e:

src/seedpass/core/totp.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import sys
 import time
+from typing import Union
 from urllib.parse import quote
 from urllib.parse import urlparse, parse_qs, unquote
 
@@ -18,13 +19,15 @@ class TotpManager:
     """Helper methods for TOTP secrets and codes."""
 
     @staticmethod
-    def derive_secret(seed: str, index: int) -> str:
-        """Derive a TOTP secret from a BIP39 seed and index."""
+    def derive_secret(seed: Union[str, bytes], index: int) -> str:
+        """Derive a TOTP secret from a seed or raw key and index."""
         return key_derivation.derive_totp_secret(seed, index)
 
     @classmethod
-    def current_code(cls, seed: str, index: int, timestamp: int | None = None) -> str:
-        """Return the TOTP code for the given seed and index."""
+    def current_code(
+        cls, seed: Union[str, bytes], index: int, timestamp: int | None = None
+    ) -> str:
+        """Return the TOTP code for the given seed/key and index."""
         secret = cls.derive_secret(seed, index)
         totp = pyotp.TOTP(secret)
         if timestamp is None:

src/tests/test_key_hierarchy.py

@@ -0,0 +1,19 @@
+import base64
+from bip_utils import Bip39SeedGenerator
+from utils.key_hierarchy import kd
+from utils.key_derivation import derive_index_key
+
+
+def test_kd_distinct_infos():
+    root = b"root" * 8
+    k1 = kd(root, b"info1")
+    k2 = kd(root, b"info2")
+    assert k1 != k2
+
+
+def test_derive_index_key_matches_hierarchy():
+    seed = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
+    seed_bytes = Bip39SeedGenerator(seed).Generate()
+    master = kd(seed_bytes, b"seedpass:v1:master")
+    expected = base64.urlsafe_b64encode(kd(master, b"seedpass:v1:storage"))
+    assert derive_index_key(seed) == expected

src/utils/key_derivation.py

@@ -25,6 +25,7 @@ from typing import Optional, Union, Dict, Any
 
 from bip_utils import Bip39SeedGenerator
 from local_bip85 import BIP85
+from .key_hierarchy import kd
 
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 from cryptography.hazmat.primitives import hashes
@@ -208,16 +209,10 @@ def derive_key_from_parent_seed(parent_seed: str, fingerprint: str = None) -> by
 
 
 def derive_index_key_seed_only(seed: str) -> bytes:
-    """Derive a deterministic Fernet key from only the BIP-39 seed."""
+    """Derive the index encryption key using the v1 hierarchy."""
     seed_bytes = Bip39SeedGenerator(seed).Generate()
-    hkdf = HKDF(
-        algorithm=hashes.SHA256(),
-        length=32,
-        salt=None,
-        info=b"password-db",
-        backend=default_backend(),
-    )
-    key = hkdf.derive(seed_bytes)
+    master = kd(seed_bytes, b"seedpass:v1:master")
+    key = kd(master, b"seedpass:v1:storage")
     return base64.urlsafe_b64encode(key)
 
 
@@ -226,23 +221,21 @@ def derive_index_key(seed: str) -> bytes:
     return derive_index_key_seed_only(seed)
 
 
-def derive_totp_secret(seed: str, index: int) -> str:
-    """Derive a base32-encoded TOTP secret from a BIP39 seed."""
+def derive_totp_secret(seed: Union[str, bytes], index: int) -> str:
+    """Derive a base32-encoded TOTP secret from a seed or raw key."""
     try:
-        # Initialize BIP85 from the BIP39 seed bytes
-        seed_bytes = Bip39SeedGenerator(seed).Generate()
+        if isinstance(seed, (bytes, bytearray)):
+            seed_bytes = bytes(seed)
+        else:
+            seed_bytes = Bip39SeedGenerator(seed).Generate()
         bip85 = BIP85(seed_bytes)
 
-        # Build the BIP32 path m/83696968'/39'/TOTP'/{index}'
         totp_int = int.from_bytes(b"TOTP", "big")
         path = f"m/83696968'/{TOTP_PURPOSE}'/{totp_int}'/{index}'"
-
-        # Derive entropy using the same scheme as BIP85
         child_key = bip85.bip32_ctx.DerivePath(path)
         key_bytes = child_key.PrivateKey().Raw().ToBytes()
         entropy = hmac.new(b"bip-entropy-from-k", key_bytes, hashlib.sha512).digest()
 
-        # Hash the first 32 bytes of entropy and encode the first 20 bytes
         hashed = hashlib.sha256(entropy[:32]).digest()
         secret = base64.b32encode(hashed[:20]).decode("utf-8")
         logger.debug(f"Derived TOTP secret for index {index}.")

src/utils/key_hierarchy.py

@@ -0,0 +1,28 @@
+"""Key hierarchy helper functions."""
+
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.backends import default_backend
+
+
+def kd(root: bytes, info: bytes, length: int = 32) -> bytes:
+    """Derive a sub-key from ``root`` using HKDF-SHA256.
+
+    Parameters
+    ----------
+    root:
+        Root key material.
+    info:
+        Domain separation string.
+    length:
+        Length of the derived key in bytes. Defaults to 32.
+    """
+
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=length,
+        salt=None,
+        info=info,
+        backend=default_backend(),
+    )
+    return hkdf.derive(root)