From d0e9b945b5ca3e62fd1f9a9ecd2c443ce37bf064 Mon Sep 17 00:00:00 2001 From: Manan Tyagi Date: Fri, 17 Jul 2026 16:49:01 +0530 Subject: [PATCH] added nsf support on rotation script --- .../keepercli/commands/pam/pam_rotation.py | 232 ++++++++++++------ .../src/keepersdk/vault/nsf_management.py | 89 +++++++ .../src/keepersdk/vault/vault_extensions.py | 11 +- 3 files changed, 259 insertions(+), 73 deletions(-) diff --git a/keepercli-package/src/keepercli/commands/pam/pam_rotation.py b/keepercli-package/src/keepercli/commands/pam/pam_rotation.py index ab83c7cf..48579292 100644 --- a/keepercli-package/src/keepercli/commands/pam/pam_rotation.py +++ b/keepercli-package/src/keepercli/commands/pam/pam_rotation.py @@ -40,6 +40,8 @@ 'pamDatabase', 'pamDirectory', 'pamMachine', 'pamUser', 'pamRemoteBrowser', ) +_PAM_SCRIPT_RECORD_TYPES = ('pamUser', 'pamDirectory') + def _resolve_nsf_record_uid(vault: vault_online.VaultOnline, identifier: str) -> Optional[str]: """Resolve an NSF record UID from a UID or exact title.""" @@ -94,6 +96,130 @@ def _load_pam_typed_record( return None +def _is_nsf_pam_record(vault: vault_online.VaultOnline, record_uid: str) -> bool: + return bool(record_uid and vault.nsf_data and vault.nsf_data.get_record(record_uid)) + + +def _save_pam_typed_record( + vault: vault_online.VaultOnline, record: vault_record.TypedRecord) -> None: + """Persist a PAM typed record via NSF or classic update (including file/script links).""" + _attach_record_key(vault, record) + if _is_nsf_pam_record(vault, record.record_uid): + nsf_management.update_nsf_typed_record(vault, record) + else: + record_management.update_record(vault, record) + vault.sync_requested = True + + +def _iter_nsf_pam_script_records( + vault: vault_online.VaultOnline, pattern: Optional[str] = None): + """Yield NSF pamUser/pamDirectory records optionally filtered by UID/title pattern.""" + if not vault.nsf_data: + return + pattern_cf = pattern.casefold() if pattern else None + for entry in vault.nsf_data.records(): + typed = _load_nsf_typed_record(vault, entry.record_uid) + if not typed or typed.record_type not in _PAM_SCRIPT_RECORD_TYPES: + continue + if pattern_cf: + title = (typed.title or '').casefold() + if typed.record_uid != pattern and pattern_cf not in title and not fnmatch.fnmatch(title, pattern_cf): + continue + yield typed + + +def _find_pam_script_records( + vault: vault_online.VaultOnline, pattern: Optional[str] = None): + """Find pamUser/pamDirectory TypedRecords in classic vault and NSF.""" + found = [] + seen = set() + for rec in vault.vault_data.find_records( + criteria=pattern, record_version=3, record_type=_PAM_SCRIPT_RECORD_TYPES): + loaded = vault.vault_data.load_record(rec.record_uid) + if not isinstance(loaded, vault_record.TypedRecord): + continue + found.append(loaded) + seen.add(loaded.record_uid) + + if pattern: + nsf_uid = _resolve_nsf_record_uid(vault, pattern) + if nsf_uid and nsf_uid not in seen: + typed = _load_nsf_typed_record(vault, nsf_uid) + if typed and typed.record_type in _PAM_SCRIPT_RECORD_TYPES: + found.append(typed) + seen.add(typed.record_uid) + else: + for typed in _iter_nsf_pam_script_records(vault, pattern): + if typed.record_uid not in seen: + found.append(typed) + seen.add(typed.record_uid) + else: + for typed in _iter_nsf_pam_script_records(vault): + if typed.record_uid not in seen: + found.append(typed) + seen.add(typed.record_uid) + return found + + +def _get_unique_pam_script_record( + vault: vault_online.VaultOnline, record_name: str) -> vault_record.TypedRecord: + """Resolve a single pamUser/pamDirectory for script commands (classic or NSF).""" + records = _find_pam_script_records(vault, record_name) + if len(records) == 0: + raise base.CommandError(f'Record "{record_name}" not found') + if len(records) > 1: + raise base.CommandError(f'Record "{record_name}" is not unique. Use record UID.') + return records[0] + + +def _load_script_file_record(vault: vault_online.VaultOnline, file_ref: str): + """Load a script file attachment record from classic vault or NSF.""" + if not file_ref: + return None + file_record = vault.vault_data.load_record(file_ref) + if file_record: + return file_record + return _load_nsf_typed_record(vault, file_ref) + + +def _find_script_value(vault, script_field, script_name): + """Find a script value by fileRef UID or file title/name.""" + if not script_field or not script_name: + return None + script_value = next( + (x for x in script_field.value if isinstance(x, dict) and x.get('fileRef') == script_name), + None) + if script_value is not None: + return script_value + s_name = script_name.casefold() + for x in script_field.value: + if not isinstance(x, dict): + continue + file_uid = x.get('fileRef') + file_record = _load_script_file_record(vault, file_uid) + if not file_record: + continue + if getattr(file_record, 'record_uid', None) == script_name: + return x + title = (getattr(file_record, 'title', None) or '').casefold() + name = (getattr(file_record, 'name', None) or '').casefold() + if title == s_name or name == s_name: + return x + return None + + +def _resolve_script_credential_uid(vault: vault_online.VaultOnline, ref: str) -> Optional[str]: + """Resolve a credential UID for script recordRef (classic or NSF).""" + if not ref: + return None + loaded = _load_pam_typed_record(vault, ref) + if loaded: + return loaded.record_uid + if vault.vault_data.get_record_key(ref) or vault.vault_data.load_record(ref): + return ref + return None + + def _iter_nsf_pam_configurations(vault: vault_online.VaultOnline): """Yield NSF PAM configuration TypedRecords.""" if not vault.nsf_data: @@ -1275,25 +1401,22 @@ def execute(self, context: KeeperParams, **kwargs): table = [] header = ['record_uid', 'title', 'record_type', 'script_uid', 'script_name', 'records', 'command'] - for rec in vault.vault_data.find_records(criteria=pattern, record_version=3, - record_type=('pamUser', 'pamDirectory')): - record = vault.vault_data.load_record(rec.record_uid) - if not isinstance(record, vault_record.TypedRecord): - continue + for record in _find_pam_script_records(vault, pattern): for field in (x for x in record.fields if x.type == 'script'): - value = field.get_default_value(dict) - if not value: - continue - file_ref = value.get('fileRef') - if not file_ref: - continue - file_record = vault.vault_data.load_record(file_ref) - if not file_record: - continue - records = value.get('recordRef') - command = value.get('command') - table.append([record.record_uid, record.title, record.record_type, file_record.record_uid, - file_record.title, records, command]) + for value in (field.value or []): + if not isinstance(value, dict): + continue + file_ref = value.get('fileRef') + if not file_ref: + continue + file_record = _load_script_file_record(vault, file_ref) + if not file_record: + continue + records = value.get('recordRef') + command = value.get('command') + table.append([record.record_uid, record.title, record.record_type, + getattr(file_record, 'record_uid', file_ref), + getattr(file_record, 'title', file_ref), records, command]) fmt = kwargs.get('format') if fmt != 'json': header = [report_utils.field_to_title(x) for x in header] @@ -1323,12 +1446,7 @@ def execute(self, context: KeeperParams, **kwargs): record_name = kwargs.get('record') if not record_name: raise base.CommandError('"record" argument is required') - records = list(vault.vault_data.find_records(criteria=record_name, record_version=3, record_type=('pamUser', 'pamDirectory'))) - if len(records) == 0: - raise base.CommandError(f'Record "{record_name}" not found') - if len(records) > 1: - raise base.CommandError(f'Record "{record_name}" is not unique. Use record UID.') - record = vault.vault_data.load_record(records[0].record_uid) + record = _get_unique_pam_script_record(vault, record_name) if not isinstance(record, vault_record.TypedRecord): raise base.CommandError(f'Record "{record.title}" is not a rotation record.') @@ -1361,14 +1479,14 @@ def execute(self, context: KeeperParams, **kwargs): record_refs = kwargs.get('add_credential') if isinstance(record_refs, list): for ref in record_refs: - if ref in vault.vault_data._records: - script_value['recordRef'].append(ref) + resolved = _resolve_script_credential_uid(vault, ref) + if resolved: + script_value['recordRef'].append(resolved) cmd = kwargs.get('script_command') if cmd: script_value['command'] = cmd - record_management.update_record(vault, record) - vault.sync_data = True + _save_pam_typed_record(vault, record) class PAMScriptEditCommand(base.ArgparseCommand): @@ -1401,31 +1519,14 @@ def execute(self, context: KeeperParams, **kwargs): if not script_name: raise base.CommandError('"script" argument is required') - records = list(vault.vault_data.find_records(criteria=record_name, record_version=3, record_type=('pamUser', 'pamDirectory'))) - if len(records) == 0: - raise base.CommandError(f'Record "{record_name}" not found') - if len(records) > 1: - raise base.CommandError(f'Record "{record_name}" is not unique. Use record UID.') - record = vault.vault_data.load_record(records[0].record_uid) + record = _get_unique_pam_script_record(vault, record_name) if not isinstance(record, vault_record.TypedRecord): raise base.CommandError(f'Record "{record.title}" is not a rotation record.') script_field = next((x for x in record.fields if x.type == 'script'), None) if script_field is None: raise base.CommandError(f'Record "{record.title}" has no rotation scripts.') - script_value = next((x for x in script_field.value if x.get('fileRef') == script_name), None) - if script_value is None: - s_name = script_name.casefold() - for x in script_field.value: - file_uid = x.get('fileRef') - file_record = vault.vault_data.load_record(file_uid) - if isinstance(file_record, vault_record.FileRecord): - if file_record.record_uid == s_name: - script_value = x - break - elif file_record.title.casefold() == s_name: - script_value = x - break + script_value = _find_script_value(vault, script_field, script_name) if not isinstance(script_value, dict): raise base.CommandError(f'Record "{record.title}" does not have script "{script_name}"') @@ -1437,11 +1538,17 @@ def execute(self, context: KeeperParams, **kwargs): refs.update(record_refs) remove_credential = kwargs.get('remove_credential') if isinstance(remove_credential, list) and remove_credential: - refs.difference_update(remove_credential) + for ref in remove_credential: + resolved = _resolve_script_credential_uid(vault, ref) or ref + refs.discard(resolved) + refs.discard(ref) modified = True add_credential = kwargs.get('add_credential') if isinstance(add_credential, list) and add_credential: - refs.update(add_credential) + for ref in add_credential: + resolved = _resolve_script_credential_uid(vault, ref) + if resolved: + refs.add(resolved) modified = True if modified: script_value['recordRef'] = list(refs) @@ -1453,8 +1560,7 @@ def execute(self, context: KeeperParams, **kwargs): if not modified: raise base.CommandError('Nothing to do') - record_management.update_record(vault, record) - vault.sync_data = True + _save_pam_typed_record(vault, record) class PAMScriptDeleteCommand(base.ArgparseCommand): @@ -1480,35 +1586,17 @@ def execute(self, context: KeeperParams, **kwargs): if not script_name: raise base.CommandError('"script" argument is required') - records = list(vault.vault_data.find_records(criteria=record_name, record_version=3, record_type=('pamUser', 'pamDirectory'))) - if len(records) == 0: - raise base.CommandError(f'Record "{record_name}" not found') - if len(records) > 1: - raise base.CommandError(f'Record "{record_name}" is not unique. Use record UID.') - record = vault.vault_data.load_record(records[0].record_uid) + record = _get_unique_pam_script_record(vault, record_name) if not isinstance(record, vault_record.TypedRecord): raise base.CommandError(f'Record "{record.title}" is not a rotation record.') script_field = next((x for x in record.fields if x.type == 'script'), None) if script_field is None: raise base.CommandError(f'Record "{record.title}" has no rotation scripts.') - script_value = next((x for x in script_field.value if x.get('fileRef') == script_name), None) - if script_value is None: - s_name = script_name.casefold() - for x in script_field.value: - file_uid = x.get('fileRef') - file_record = vault.vault_data.load_record(file_uid) - if isinstance(file_record, vault_record.FileRecord): - if file_record.record_uid == s_name: - script_value = x - break - elif file_record.title.casefold() == s_name: - script_value = x - break + script_value = _find_script_value(vault, script_field, script_name) if not isinstance(script_value, dict): raise base.CommandError(f'Record "{record.title}" does not have script "{script_name}"') script_field.value.remove(script_value) - record_management.update_record(vault, record) - vault.sync_data = True + _save_pam_typed_record(vault, record) diff --git a/keepersdk-package/src/keepersdk/vault/nsf_management.py b/keepersdk-package/src/keepersdk/vault/nsf_management.py index 090d11a1..a3e230a3 100644 --- a/keepersdk-package/src/keepersdk/vault/nsf_management.py +++ b/keepersdk-package/src/keepersdk/vault/nsf_management.py @@ -694,6 +694,95 @@ def update_nsf_record( return result +def update_nsf_typed_record( + vault: VaultOnline, + record: 'vault_record.TypedRecord', + *, + request_sync: bool = True) -> NsfModifyResult: + """Update an NSF typed record, including record-link adds/removes for file/script refs. + + Mirrors classic ``record_management.update_record`` so PAM rotation scripts + (and other fileRef/script attachments) work on Keeper Drive records. + """ + from . import vault_record as vr + + if not isinstance(record, vr.TypedRecord) or not record.record_uid: + raise NsfError('TypedRecord with record_uid is required') + record_uid = resolve_nsf_record_uid(vault, record.record_uid) or record.record_uid + if not is_nsf_record(vault, record_uid): + raise NsfError(f'NSF record not found: {record.record_uid}') + record.record_uid = record_uid + + record_key = _get_record_key(vault, record_uid) + storage_row = _nsf_view(vault).storage.records.get_entity(record_uid) + revision = storage_row.revision if storage_row else 0 + + existing = vr.TypedRecord() + existing.record_uid = record_uid + try: + meta = load_nsf_record_metadata(vault, record_uid) + existing.load_record_data({ + 'type': meta.get('type') or '', + 'title': meta.get('title') or record_uid, + 'notes': meta.get('notes') or '', + 'fields': meta.get('fields') or [], + 'custom': meta.get('custom') or [], + }) + except NsfError: + pass + + data = vault_extensions.extract_typed_record_data(record, None) + ru = record_pb2.RecordUpdate() + ru.record_uid = utils.base64_url_decode(record_uid) + ru.client_modified_time = utils.current_milli_time() + ru.revision = revision + ru.data = crypto.encrypt_aes_v2(vault_extensions.get_padded_json_bytes(data), record_key) + + existing_refs = vault_extensions.extract_typed_record_refs(existing) + refs = vault_extensions.extract_typed_record_refs(record) + for ref_uid in refs.difference(existing_refs): + ref_key = None + if record.linked_keys and ref_uid in record.linked_keys: + ref_key = record.linked_keys[ref_uid] + if not ref_key: + try: + ref_key = vault.vault_data.get_record_key(ref_uid) + except Exception: + ref_key = None + if not ref_key and vault.nsf_data: + entry = vault.nsf_data.get_record(ref_uid) + if entry: + ref_key = entry.record_key + if not ref_key: + continue + link = record_pb2.RecordLink() + link.record_uid = utils.base64_url_decode(ref_uid) + link.record_key = crypto.encrypt_aes_v2(ref_key, record_key) + ru.record_links_add.append(link) + for ref_uid in existing_refs.difference(refs): + ru.record_links_remove.append(utils.base64_url_decode(ref_uid)) + + rq = record_pb2.RecordsUpdateRequest() + rq.client_time = utils.current_milli_time() + rq.records.append(ru) + + auth = vault.keeper_auth + response = auth.execute_auth_rest( + 'vault/records/v3/update', rq, response_type=record_pb2.RecordsModifyResponse) + if response is None: + response = auth.execute_auth_rest( + 'vault/records_update', rq, response_type=record_pb2.RecordsModifyResponse) + assert response is not None + + result = _parse_modify_response(response, record_uid) + if not result.success: + raise KeeperApiError(result.status, result.message) + if request_sync: + vault.sync_requested = True + vault.run_pending_jobs() + return result + + def get_nsf_record_details( vault: VaultOnline, record_uids: Iterable[str]) -> Dict[str, Any]: diff --git a/keepersdk-package/src/keepersdk/vault/vault_extensions.py b/keepersdk-package/src/keepersdk/vault/vault_extensions.py index 7fe695fc..356fc47e 100644 --- a/keepersdk-package/src/keepersdk/vault/vault_extensions.py +++ b/keepersdk-package/src/keepersdk/vault/vault_extensions.py @@ -236,11 +236,20 @@ def extract_audit_data(record: Union[vault_record.KeeperRecord, vault_record.Typ def extract_typed_record_refs(record: vault_record.TypedRecord) -> Set[str]: refs = set() for field in itertools.chain(record.fields, record.custom): - if field.type in {'fileRef', 'addressRef', 'cardRef'}: + if field.type in {'fileRef', 'addressRef', 'cardRef', 'recordRef'}: if isinstance(field.value, list): for ref in field.value: if isinstance(ref, str): refs.add(ref) + elif field.type == 'script': + if not isinstance(field.value, list): + continue + for script in field.value: + if not isinstance(script, dict): + continue + file_ref = script.get('fileRef') + if isinstance(file_ref, str) and file_ref: + refs.add(file_ref) return refs