Fix NSF folder titles for team-shared folders

keepersdk-package/src/keepersdk/vault/nsf_crypto.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import json
-from typing import Dict, List, Optional
+from dataclasses import dataclass
+from typing import Any, Dict, List, Mapping, Optional
 
 from .. import crypto, utils
 from ..authentication import keeper_auth
@@ -11,6 +12,15 @@
 
 _FOLDER_KEY_ENCRYPTION = folder_pb2.FolderKeyEncryptionType
 _ENCRYPTED_KEY_TYPE = folder_pb2.EncryptedKeyType
+_ACCESS_TYPE = folder_pb2.AccessType
+
+
+@dataclass(frozen=True)
+class TeamKeyMaterial:
+    """Decrypted team keys used to unwrap team-shared NSF folder keys."""
+    team_key: bytes
+    rsa_private_key: Optional[Any] = None
+    ec_private_key: Optional[Any] = None
 
 
 def try_decrypt_symmetric(encrypted_key: bytes, symmetric_key: bytes) -> Optional[bytes]:
@@ -42,30 +52,85 @@ def try_decrypt_with_user_keys(encrypted_key: bytes, auth_context: keeper_auth.A
     return None
 
 
+def try_decrypt_with_typed_key(
+        encrypted_key: bytes,
+        key_type: int,
+        *,
+        aes_key: Optional[bytes] = None,
+        rsa_key: Optional[Any] = None,
+        ecc_key: Optional[Any] = None) -> Optional[bytes]:
+    """Decrypt using the algorithm indicated by *key_type* (Vault decryptFolderKeyByType)."""
+    try:
+        if key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_data_key_gcm):
+            if aes_key is not None:
+                return crypto.decrypt_aes_v2(encrypted_key, aes_key)
+        elif key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_data_key):
+            if aes_key is not None:
+                return crypto.decrypt_aes_v1(encrypted_key, aes_key)
+        elif key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_public_key):
+            if rsa_key is not None:
+                return crypto.decrypt_rsa(encrypted_key, rsa_key)
+        elif key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_public_key_ecc):
+            if ecc_key is not None:
+                return crypto.decrypt_ec(encrypted_key, ecc_key)
+    except Exception:
+        return None
+    return None
+
+
 def try_decrypt_from_folder_access(
         folder_uid: str,
         storage: INSFStorage,
-        auth_context: keeper_auth.AuthContext) -> Optional[bytes]:
+        auth_context: keeper_auth.AuthContext,
+        teams: Optional[Mapping[str, TeamKeyMaterial]] = None) -> Optional[bytes]:
+    """Unwrap folder key from folderAccesses (user or team), mirroring Web Vault."""
+    teams = teams or {}
     for fa in storage.folder_accesses.get_links_by_subject(folder_uid):
         if not fa.folder_key_encrypted:
             continue
         try:
             enc_key = utils.base64_url_decode(fa.folder_key_encrypted)
             key_type = fa.folder_key_type
-            if key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_data_key_gcm):
-                return crypto.decrypt_aes_v2(enc_key, auth_context.data_key)
-            if key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_data_key):
-                return crypto.decrypt_aes_v1(enc_key, auth_context.data_key)
-            if key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_public_key):
-                if auth_context.rsa_private_key is not None:
-                    return crypto.decrypt_rsa(enc_key, auth_context.rsa_private_key)
-            elif key_type == int(_ENCRYPTED_KEY_TYPE.encrypted_by_public_key_ecc):
-                if auth_context.ec_private_key is not None:
-                    return crypto.decrypt_ec(enc_key, auth_context.ec_private_key)
-            else:
-                result = try_decrypt_with_user_keys(enc_key, auth_context)
-                if result is not None:
-                    return result
+            access_uid = fa.access_type_uid
+            use_team = (
+                fa.access_type == int(_ACCESS_TYPE.AT_TEAM)
+                or (access_uid in teams)
+            )
+
+            folder_key: Optional[bytes] = None
+            if use_team and access_uid in teams:
+                team = teams[access_uid]
+                folder_key = try_decrypt_with_typed_key(
+                    enc_key, key_type,
+                    aes_key=team.team_key,
+                    rsa_key=team.rsa_private_key,
+                    ecc_key=team.ec_private_key,
+                )
+                if folder_key is None:
+                    folder_key = try_decrypt_symmetric(enc_key, team.team_key)
+                if folder_key is None and team.rsa_private_key is not None:
+                    try:
+                        folder_key = crypto.decrypt_rsa(enc_key, team.rsa_private_key)
+                    except Exception:
+                        pass
+                if folder_key is None and team.ec_private_key is not None:
+                    try:
+                        folder_key = crypto.decrypt_ec(enc_key, team.ec_private_key)
+                    except Exception:
+                        pass
+
+            if folder_key is None:
+                folder_key = try_decrypt_with_typed_key(
+                    enc_key, key_type,
+                    aes_key=auth_context.data_key,
+                    rsa_key=auth_context.rsa_private_key,
+                    ecc_key=auth_context.ec_private_key,
+                )
+            if folder_key is None:
+                folder_key = try_decrypt_with_user_keys(enc_key, auth_context)
+
+            if folder_key is not None and len(folder_key) == 32:
+                return folder_key
         except Exception:
             continue
     return None
@@ -74,18 +139,28 @@ def try_decrypt_from_folder_access(
 def try_decrypt_folder_key(
         fk: nsf.NSFFolderKey,
         auth_context: keeper_auth.AuthContext,
-        decrypted_folder_keys: Dict[str, bytes]) -> Optional[bytes]:
+        decrypted_folder_keys: Dict[str, bytes],
+        teams: Optional[Mapping[str, TeamKeyMaterial]] = None) -> Optional[bytes]:
+    """
+    Attempt decrypt from a FolderKey link.
+
+    Returns None for ENCRYPTED_BY_TEAM_KEY (caller must use folderAccesses).
+    For PARENT_KEY without a decrypted parent, returns None so caller can fall back.
+    """
+    enc_key_type = fk.encrypted_by
+    if enc_key_type == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_TEAM_KEY):
+        return None  # key lives in folderAccesses
     if not fk.folder_key:
         return None
     try:
-        enc_key_type = fk.encrypted_by
         encrypted_key = utils.base64_url_decode(fk.folder_key)
         if enc_key_type == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_USER_KEY):
             return try_decrypt_with_user_keys(encrypted_key, auth_context)
         if enc_key_type == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_PARENT_KEY):
-            if not fk.parent_uid:
+            parent_uid = fk.parent_uid
+            if not parent_uid:
                 return None
-            parent_key = decrypted_folder_keys.get(fk.parent_uid)
+            parent_key = decrypted_folder_keys.get(parent_uid)
             if parent_key is None:
                 return None
             return try_decrypt_symmetric(encrypted_key, parent_key)
@@ -114,9 +189,31 @@ def try_decrypt_folder_entity_key(
     return None
 
 
+def _folder_needs_access_fallback(
+        folder_uid: str,
+        keys_by_folder: Mapping[str, List[nsf.NSFFolderKey]],
+        decrypted_keys: Mapping[str, bytes]) -> bool:
+    """True when FolderKey links require folderAccesses (TEAM_KEY or failed PARENT/USER)."""
+    if folder_uid in decrypted_keys:
+        return False
+    for fk in keys_by_folder.get(folder_uid, []):
+        if fk.encrypted_by == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_TEAM_KEY):
+            return True
+        if fk.encrypted_by == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_PARENT_KEY):
+            parent_uid = fk.parent_uid
+            if not parent_uid or parent_uid not in decrypted_keys:
+                return True
+        if fk.encrypted_by == int(_FOLDER_KEY_ENCRYPTION.ENCRYPTED_BY_USER_KEY):
+            return True  # USER_KEY already tried; fall back to accesses
+    return True
+
+
 def decrypt_folder_keys(
         storage: INSFStorage,
-        auth_context: keeper_auth.AuthContext) -> Dict[str, bytes]:
+        auth_context: keeper_auth.AuthContext,
+        teams: Optional[Mapping[str, TeamKeyMaterial]] = None) -> Dict[str, bytes]:
+    """Decrypt NSF folder keys. Pass *teams* for team-shared folder unwrap."""
+    teams = teams or {}
     decrypted_keys: Dict[str, bytes] = {}
     keys_by_folder: Dict[str, List[nsf.NSFFolderKey]] = {}
     for fk in storage.folder_keys.get_all_links():
@@ -130,7 +227,7 @@ def decrypt_folder_keys(
             if folder_uid in decrypted_keys:
                 continue
             for fk in folder_keys:
-                key = try_decrypt_folder_key(fk, auth_context, decrypted_keys)
+                key = try_decrypt_folder_key(fk, auth_context, decrypted_keys, teams)
                 if key is not None:
                     decrypted_keys[folder_uid] = key
                     progress = True
@@ -143,17 +240,16 @@ def decrypt_folder_keys(
                 decrypted_keys[row.folder_uid] = key
                 progress = True
 
-    for folder_uid in keys_by_folder:
-        if folder_uid not in decrypted_keys:
-            key = try_decrypt_from_folder_access(folder_uid, storage, auth_context)
-            if key is not None:
-                decrypted_keys[folder_uid] = key
-
-    for row in storage.folders.get_all_entities():
-        if row.folder_uid not in decrypted_keys:
-            key = try_decrypt_from_folder_access(row.folder_uid, storage, auth_context)
-            if key is not None:
-                decrypted_keys[row.folder_uid] = key
+    # folderAccesses fallback (TEAM_KEY, PARENT without parent, USER_KEY miss, bare accesses)
+    candidates = set(keys_by_folder.keys()) | {row.folder_uid for row in folder_rows}
+    for folder_uid in candidates:
+        if folder_uid in decrypted_keys:
+            continue
+        if not _folder_needs_access_fallback(folder_uid, keys_by_folder, decrypted_keys):
+            continue
+        key = try_decrypt_from_folder_access(folder_uid, storage, auth_context, teams)
+        if key is not None:
+            decrypted_keys[folder_uid] = key
 
     return decrypted_keys
 

keepersdk-package/src/keepersdk/vault/nsf_data.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, Iterable, List, Optional, Set
+from typing import Dict, Iterable, List, Mapping, Optional, Set
 
 from ..authentication import keeper_auth
 from . import nsf_crypto, nsf_storage_types as nsf
@@ -103,13 +103,17 @@ def rebuild_data(self, changes: Optional[NSFRebuildTask] = None) -> None:
         del changes
         self.rebuild_nsf(self._auth_context)
 
-    def rebuild_nsf(self, auth_context: Optional[keeper_auth.AuthContext]) -> None:
+    def rebuild_nsf(
+            self,
+            auth_context: Optional[keeper_auth.AuthContext],
+            teams: Optional[Mapping[str, nsf_crypto.TeamKeyMaterial]] = None) -> None:
         self._folders.clear()
         self._records.clear()
         if auth_context is None:
             return
 
-        decrypted_folder_keys = nsf_crypto.decrypt_folder_keys(self._storage, auth_context)
+        decrypted_folder_keys = nsf_crypto.decrypt_folder_keys(
+            self._storage, auth_context, teams=teams)
         decrypted_record_keys = nsf_crypto.decrypt_record_keys(
             self._storage, decrypted_folder_keys, auth_context)
 

keepersdk-package/src/keepersdk/vault/nsf_management.py

@@ -465,7 +465,8 @@ def _get_folder_key(vault: VaultOnline, folder_uid: str) -> bytes:
         return folder.folder_key
 
     auth_context = vault.keeper_auth.auth_context
-    decrypted = nsf_crypto.decrypt_folder_keys(view.storage, auth_context)
+    teams = vault.vault_data.get_nsf_team_key_materials()
+    decrypted = nsf_crypto.decrypt_folder_keys(view.storage, auth_context, teams=teams)
     key = decrypted.get(folder_uid)
     if key is None:
         label = folder.name if folder and folder.name and folder.name != '(NSF Folder)' else folder_uid

keepersdk-package/src/keepersdk/vault/vault_data.py

@@ -260,6 +260,18 @@ def get_team_key(self, team_uid: str) -> Optional[bytes]:
         if t:
             return t.team_key
 
+    def get_nsf_team_key_materials(self) -> Dict[str, 'nsf_crypto.TeamKeyMaterial']:
+        """Decrypted team keys for NSF folderAccesses unwrap (team-shared folders)."""
+        from . import nsf_crypto
+        return {
+            uid: nsf_crypto.TeamKeyMaterial(
+                team_key=t.team_key,
+                rsa_private_key=t.rsa_private_key,
+                ec_private_key=t.ec_private_key,
+            )
+            for uid, t in self._teams.items()
+        }
+
     def teams(self) -> Iterable[vault_types.TeamInfo]:
         return (x.info for x in self._teams.values())
 

keepersdk-package/src/keepersdk/vault/vault_online.py

@@ -116,7 +116,12 @@ def sync_down(self, force=False):
         self._sync_record_types = False
         self._vault_data.rebuild_data(result.vault)
         if self._nsf_data is not None:
-            self._nsf_data.rebuild_nsf(self._keeper_auth.auth_context)
+            # Teams are decrypted during vault rebuild; NSF team-shared folders
+            # need those keys to unwrap folderAccesses (ENCRYPTED_BY_TEAM_KEY).
+            self._nsf_data.rebuild_nsf(
+                self._keeper_auth.auth_context,
+                teams=self._vault_data.get_nsf_team_key_materials(),
+            )
 
     def _background_task(self):
         if self._keeper_auth.auth_context.enterprise_ec_public_key: