diff --git a/examples/sdk_examples/nested_shared_folders/nsf_record_add.py b/examples/sdk_examples/nested_shared_folders/nsf_record_add.py index 2eaac9e1..9eea775b 100644 --- a/examples/sdk_examples/nested_shared_folders/nsf_record_add.py +++ b/examples/sdk_examples/nested_shared_folders/nsf_record_add.py @@ -511,22 +511,20 @@ def close_vault(vault: vault_online.VaultOnline, keeper_auth_context: keeper_aut keeper_auth_context.close() def nsf_record_add(vault: vault_online.VaultOnline) -> None: - """Add a record to NSF (nsf-record-add).""" - TITLE = "My NSF Login" - RECORD_TYPE = "login" # e.g. login, password, general - FOLDER_UID_OR_NAME = "Projects" # NSF folder name or UID; None for root - NOTES = "Created via SDK example" + """Add a single NSF record (nsf-record-add / create_nsf_record).""" + TITLE = 'My NSF Login' + RECORD_TYPE = 'login' + FOLDER = 'Projects' # NSF folder name or UID; set None for root + NOTES = 'Created via SDK example' FIELDS = { - "login": "user@example.com", - "password": "changeme", - "url": "https://example.com", + 'login': 'user@example.com', + 'password': 'changeme', + 'url': 'https://example.com', } folder_uid = None - if FOLDER_UID_OR_NAME: - folder_uid = nsf_management.resolve_nsf_folder_uid(vault, FOLDER_UID_OR_NAME) - if not folder_uid: - raise ValueError(f"NSF folder not found: {FOLDER_UID_OR_NAME}") + if FOLDER: + folder_uid = nsf_management.resolve_nsf_folder_uid(vault, FOLDER) or FOLDER result = nsf_management.create_nsf_record( vault, diff --git a/examples/sdk_examples/nested_shared_folders/nsf_record_add_batch.py b/examples/sdk_examples/nested_shared_folders/nsf_record_add_batch.py new file mode 100644 index 00000000..d7e9f6d7 --- /dev/null +++ b/examples/sdk_examples/nested_shared_folders/nsf_record_add_batch.py @@ -0,0 +1,565 @@ +import getpass +import json +import logging +import sqlite3 +from typing import Dict, Optional + +import fido2 +import webbrowser + +from keepersdk import errors, utils +from keepersdk.authentication import ( + configuration, + endpoint, + keeper_auth, + login_auth, +) +from keepersdk.authentication.yubikey import ( + IKeeperUserInteraction, + yubikey_authenticate, +) +from keepersdk.constants import KEEPER_PUBLIC_HOSTS +from keepersdk.vault import nsf_management, sqlite_storage, vault_online + +try: + import pyperclip +except ImportError: + pyperclip = None + +logger = utils.get_logger() +logger.setLevel(logging.INFO) +if not logger.handlers: + _handler = logging.StreamHandler() + _handler.setLevel(logging.INFO) + _handler.setFormatter( + logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s") + ) + logger.addHandler(_handler) + + +class FidoCliInteraction(fido2.client.UserInteraction, IKeeperUserInteraction): + def output_text(self, text: str) -> None: + print(text) + + def prompt_up(self) -> None: + print( + "\nTouch the flashing Security key to authenticate or " + "press Ctrl-C to resume with the primary two factor authentication..." + ) + + def request_pin(self, permissions, rd_id): + return getpass.getpass("Enter Security Key PIN: ") + + def request_uv(self, permissions, rd_id): + print("User Verification required.") + return True + + +# Two-factor duration codes (used by LoginFlow) +_TWO_FACTOR_DURATION_CODES: Dict[login_auth.TwoFactorDuration, str] = { + login_auth.TwoFactorDuration.EveryLogin: "login", + login_auth.TwoFactorDuration.Every12Hours: "12_hours", + login_auth.TwoFactorDuration.EveryDay: "24_hours", + login_auth.TwoFactorDuration.Every30Days: "30_days", + login_auth.TwoFactorDuration.Forever: "forever", +} + + +class LoginFlow: + """ + Handles the full login process: server selection, username, password, + device approval, 2FA, SSO data key, and SSO token. + """ + + def __init__(self) -> None: + self._config = configuration.JsonConfigurationStorage() + self._logged_in_with_persistent = True + self._endpoint: Optional[endpoint.KeeperEndpoint] = None + + @property + def endpoint(self) -> Optional[endpoint.KeeperEndpoint]: + return self._endpoint + + @property + def logged_in_with_persistent(self) -> bool: + """True if login succeeded by resuming an existing persistent session (no step loop).""" + return self._logged_in_with_persistent + + def run(self) -> Optional[keeper_auth.KeeperAuth]: + """ + Run the login flow. + + Returns: + Authenticated Keeper context, or None if login fails. + """ + server = self._ensure_server() + keeper_endpoint = endpoint.KeeperEndpoint(self._config, server) + self._endpoint = keeper_endpoint + login_auth_context = login_auth.LoginAuth(keeper_endpoint) + + username = self._config.get().last_login or input("Enter username: ") + login_auth_context.resume_session = True + login_auth_context.login(username) + + while not login_auth_context.login_step.is_final(): + step = login_auth_context.login_step + if isinstance(step, login_auth.LoginStepDeviceApproval): + self._handle_device_approval(step) + elif isinstance(step, login_auth.LoginStepTwoFactor): + self._handle_two_factor(step) + elif isinstance(step, login_auth.LoginStepPassword): + self._handle_password(step) + elif isinstance(step, login_auth.LoginStepSsoToken): + self._handle_sso_token(step) + elif isinstance(step, login_auth.LoginStepSsoDataKey): + self._handle_sso_data_key(step) + elif isinstance(step, login_auth.LoginStepError): + print(f"Login error: ({step.code}) {step.message}") + return None + else: + raise NotImplementedError( + f"Unsupported login step type: {type(step).__name__}" + ) + self._logged_in_with_persistent = False + + if self._logged_in_with_persistent: + print("Successfully logged in with persistent login") + + if isinstance(login_auth_context.login_step, login_auth.LoginStepConnected): + return login_auth_context.login_step.take_keeper_auth() + + return None + + def _ensure_server(self) -> str: + if not self._config.get().last_server: + print("Available server options:") + for region, host in KEEPER_PUBLIC_HOSTS.items(): + print(f" {region}: {host}") + server = ( + input("Enter server (default: keepersecurity.com): ").strip() + or "keepersecurity.com" + ) + self._config.get().last_server = server + else: + server = self._config.get().last_server + return server + + def _handle_device_approval( + self, step: login_auth.LoginStepDeviceApproval + ) -> None: + """Device approval: same options as keepercli verify_device (email, keeper push, 2FA, resume).""" + menu = [ + ("email_send", "to send email"), + ("email_code=", "to validate verification code sent via email"), + ("keeper_push", "to send Keeper Push notification"), + ("2fa_send", "to send 2FA code"), + ("2fa_code=", "to validate a code provided by 2FA application"), + ("", "to resume"), + ] + lines = ["Approve by selecting a method below"] + lines.extend(f" {cmd} {desc}" for cmd, desc in menu) + print("\n".join(lines)) + + selection = input("Type your selection or to resume: ").strip() + if selection is None: + return + if selection in ("email_send", "es"): + step.send_push(channel=login_auth.DeviceApprovalChannel.Email) + print("An email with instructions has been sent. Press when approved.") + elif selection.startswith("email_code="): + code = selection[len("email_code=") :] + step.send_code(channel=login_auth.DeviceApprovalChannel.Email, code=code) + print("Successfully verified email code.") + elif selection in ("keeper_push", "kp"): + step.send_push(channel=login_auth.DeviceApprovalChannel.KeeperPush) + print( + "Successfully made a push notification to the approved device. " + "Press when approved." + ) + elif selection in ("2fa_send", "2fs"): + step.send_push(channel=login_auth.DeviceApprovalChannel.TwoFactor) + print("2FA code was sent.") + elif selection.startswith("2fa_code="): + code = selection[len("2fa_code=") :] + step.send_code(channel=login_auth.DeviceApprovalChannel.TwoFactor, code=code) + print("Successfully verified 2FA code.") + else: + step.resume() + + def _handle_password(self, step: login_auth.LoginStepPassword) -> None: + """Password step: prompt for password and retry on auth_failed (aligned with keepercli handle_verify_password).""" + print(f"\nEnter password for {step.username}") + while True: + password = getpass.getpass("Password: ") + if not password: + raise KeyboardInterrupt() + try: + step.verify_password(password) + break + except errors.KeeperApiError as kae: + print( + "Invalid email or password combination, please re-enter." + if kae.result_code == "auth_failed" + else kae.message + ) + + def _handle_two_factor(self, step: login_auth.LoginStepTwoFactor) -> None: + channels = [ + x + for x in step.get_channels() + if x.channel_type != login_auth.TwoFactorChannel.Other + ] + menu = [] + for i, channel in enumerate(channels): + desc = self._two_factor_channel_desc(channel.channel_type) + menu.append( + ( + str(i + 1), + f"{desc} {channel.channel_name} {channel.phone}", + ) + ) + menu.append(("q", "Quit authentication attempt and return to Commander prompt.")) + + lines = ["", "This account requires 2FA Authentication"] + lines.extend(f" {a}. {t}" for a, t in menu) + print("\n".join(lines)) + + while True: + selection = input("Selection: ") + if selection is None: + return + if selection in ("q", "Q"): + raise KeyboardInterrupt() + try: + assert selection.isnumeric() + idx = 1 if not selection else int(selection) + assert 1 <= idx <= len(channels) + channel = channels[idx - 1] + desc = self._two_factor_channel_desc(channel.channel_type) + print(f"Selected {idx}. {desc}") + except AssertionError: + print( + "Invalid entry, additional factors of authentication shown " + "may be configured if not currently enabled." + ) + continue + + if channel.channel_type in ( + login_auth.TwoFactorChannel.TextMessage, + login_auth.TwoFactorChannel.KeeperDNA, + login_auth.TwoFactorChannel.DuoSecurity, + ): + action = next( + ( + x + for x in step.get_channel_push_actions(channel.channel_uid) + if x + in ( + login_auth.TwoFactorPushAction.TextMessage, + login_auth.TwoFactorPushAction.KeeperDna, + ) + ), + None, + ) + if action: + step.send_push(channel.channel_uid, action) + + if channel.channel_type == login_auth.TwoFactorChannel.SecurityKey: + try: + challenge = json.loads(channel.challenge) + signature = yubikey_authenticate(challenge, FidoCliInteraction()) + if signature: + print("Verified Security Key.") + step.send_code(channel.channel_uid, signature) + return + except Exception as e: + logger.error(e) + continue + + # 2FA code path + step.duration = min(step.duration, channel.max_expiration) + available_dura = sorted( + x for x in _TWO_FACTOR_DURATION_CODES if x <= channel.max_expiration + ) + available_codes = [ + _TWO_FACTOR_DURATION_CODES.get(x) or "login" for x in available_dura + ] + + while True: + mfa_desc = self._two_factor_duration_desc(step.duration) + prompt_exp = ( + f"\n2FA Code Duration: {mfa_desc}.\n" + f"To change duration: 2fa_duration={'|'.join(available_codes)}" + ) + print(prompt_exp) + + selection = input("\nEnter 2FA Code or Duration: ") + if not selection: + return + if selection in available_codes: + step.duration = self._two_factor_code_to_duration(selection) + elif selection.startswith("2fa_duration="): + code = selection[len("2fa_duration=") :] + if code in available_codes: + step.duration = self._two_factor_code_to_duration(code) + else: + print(f"Invalid 2FA duration: {code}") + else: + try: + step.send_code(channel.channel_uid, selection) + print("Successfully verified 2FA Code.") + return + except errors.KeeperApiError as kae: + print(f"Invalid 2FA code: ({kae.result_code}) {kae.message}") + + def _handle_sso_data_key( + self, step: login_auth.LoginStepSsoDataKey + ) -> None: + menu = [ + ("1", "Keeper Push. Send a push notification to your device."), + ("2", "Admin Approval. Request your admin to approve this device."), + ("r", "Resume SSO authentication after device is approved."), + ("q", "Quit SSO authentication attempt and return to Commander prompt."), + ] + lines = ["Approve this device by selecting a method below:"] + lines.extend(f" {cmd:>3}. {text}" for cmd, text in menu) + print("\n".join(lines)) + + while True: + answer = input("Selection: ") + if answer is None: + return + if answer == "q": + raise KeyboardInterrupt() + if answer == "r": + step.resume() + break + if answer in ("1", "2"): + step.request_data_key( + login_auth.DataKeyShareChannel.KeeperPush + if answer == "1" + else login_auth.DataKeyShareChannel.AdminApproval + ) + else: + print(f'Action "{answer}" is not supported.') + + def _handle_sso_token(self, step: login_auth.LoginStepSsoToken) -> None: + menu = [ + ("a", "SSO User with a Master Password."), + ] + if pyperclip: + menu.append(("c", "Copy SSO Login URL to clipboard.")) + else: + menu.append(("u", "Show SSO Login URL.")) + try: + wb = webbrowser.get() + menu.append(("o", "Navigate to SSO Login URL with the default web browser.")) + except Exception: + wb = None + if pyperclip: + menu.append(("p", "Paste SSO Token from clipboard.")) + menu.append(("t", "Enter SSO Token manually.")) + menu.append(("q", "Quit SSO authentication attempt and return to Commander prompt.")) + + lines = [ + "", + "SSO Login URL:", + step.sso_login_url, + "Navigate to SSO Login URL with your browser and complete authentication.", + "Copy a returned SSO Token into clipboard." + + (" Paste that token into Commander." if pyperclip else " Then use option 't' to enter the token manually."), + 'NOTE: To copy SSO Token please click "Copy authentication token" ' + 'button on "SSO Connect" page.', + "", + ] + lines.extend(f" {a:>3}. {t}" for a, t in menu) + print("\n".join(lines)) + + while True: + token = input("Selection: ") + if token == "q": + raise KeyboardInterrupt() + if token == "a": + step.login_with_password() + return + if token == "c": + token = None + if pyperclip: + try: + pyperclip.copy(step.sso_login_url) + print("SSO Login URL is copied to clipboard.") + except Exception: + print("Failed to copy SSO Login URL to clipboard.") + else: + print("Clipboard not available (install pyperclip).") + elif token == "u": + token = None + if not pyperclip: + print("\nSSO Login URL:", step.sso_login_url, "\n") + else: + print("Unsupported menu option (use 'c' to copy URL).") + elif token == "o": + token = None + if wb: + try: + wb.open_new_tab(step.sso_login_url) + except Exception: + print("Failed to open web browser.") + elif token == "p": + if pyperclip: + try: + token = pyperclip.paste() + except Exception: + token = "" + print("Failed to paste from clipboard") + else: + token = None + print("Clipboard not available (use 't' to enter token manually).") + elif token == "t": + token = getpass.getpass("Enter SSO Token: ").strip() + else: + if len(token) < 10: + print(f"Unsupported menu option: {token}") + continue + + if token: + try: + step.set_sso_token(token) + break + except errors.KeeperApiError as kae: + print(f"SSO Login error: ({kae.result_code}) {kae.message}") + + @staticmethod + def _two_factor_channel_desc( + channel_type: login_auth.TwoFactorChannel, + ) -> str: + return { + login_auth.TwoFactorChannel.Authenticator: "TOTP (Google and Microsoft Authenticator)", + login_auth.TwoFactorChannel.TextMessage: "Send SMS Code", + login_auth.TwoFactorChannel.DuoSecurity: "DUO", + login_auth.TwoFactorChannel.RSASecurID: "RSA SecurID", + login_auth.TwoFactorChannel.SecurityKey: "WebAuthN (FIDO2 Security Key)", + login_auth.TwoFactorChannel.KeeperDNA: "Keeper DNA (Watch)", + login_auth.TwoFactorChannel.Backup: "Backup Code", + }.get(channel_type, "Not Supported") + + @staticmethod + def _two_factor_duration_desc( + duration: login_auth.TwoFactorDuration, + ) -> str: + return { + login_auth.TwoFactorDuration.EveryLogin: "Require Every Login", + login_auth.TwoFactorDuration.Forever: "Save on this Device Forever", + login_auth.TwoFactorDuration.Every12Hours: "Ask Every 12 hours", + login_auth.TwoFactorDuration.EveryDay: "Ask Every 24 hours", + login_auth.TwoFactorDuration.Every30Days: "Ask Every 30 days", + }.get(duration, "Require Every Login") + + @staticmethod + def _two_factor_code_to_duration( + text: str, + ) -> login_auth.TwoFactorDuration: + for dura, code in _TWO_FACTOR_DURATION_CODES.items(): + if code == text: + return dura + return login_auth.TwoFactorDuration.EveryLogin + + +def enable_persistent_login(keeper_auth_context: keeper_auth.KeeperAuth) -> None: + """ + Enable persistent login and register data key for device. + Sets persistent_login to on and logout_timer to 30 days. + """ + keeper_auth.set_user_setting(keeper_auth_context, 'persistent_login', '1') + keeper_auth.register_data_key_for_device(keeper_auth_context) + mins_per_day = 60 * 24 + timeout_in_minutes = mins_per_day * 30 # 30 days + keeper_auth.set_user_setting(keeper_auth_context, 'logout_timer', str(timeout_in_minutes)) + print("Persistent login turned on successfully and device registered") + + +def login(): + """ + Handle the login process including server selection, authentication, + and multi-factor authentication steps (device approval, password, 2FA + with channel selection and Security Key, SSO data key, SSO token). + + Returns: + tuple: (keeper_auth_context, keeper_endpoint) on success, or (None, None) if login fails. + """ + flow = LoginFlow() + keeper_auth_context = flow.run() + if keeper_auth_context and not flow.logged_in_with_persistent: + enable_persistent_login(keeper_auth_context) + keeper_endpoint = flow.endpoint if keeper_auth_context else None + return keeper_auth_context, keeper_endpoint + + +def open_vault(keeper_auth_context: keeper_auth.KeeperAuth) -> vault_online.VaultOnline: + conn = sqlite3.Connection("file::memory:", uri=True) + vault_storage = sqlite_storage.SqliteVaultStorage( + lambda: conn, + vault_owner=bytes(keeper_auth_context.auth_context.username, "utf-8"), + ) + vault = vault_online.VaultOnline(keeper_auth_context, vault_storage) + vault.sync_down() + return vault + + +def close_vault(vault: vault_online.VaultOnline, keeper_auth_context: keeper_auth.KeeperAuth) -> None: + vault.close() + keeper_auth_context.close() + +def nsf_record_add_batch(vault: vault_online.VaultOnline) -> None: + """Batch record add execution (nsf-record-add --batch / create_nsf_records).""" + RECORDS = [ + { + 'title': 'My NSF Login', + 'record_type': 'login', + 'folder': 'Projects', # NSF folder name or UID; omit for root + 'notes': 'Created via SDK batch example', + 'fields': { + 'login': 'user@example.com', + 'password': 'changeme', + 'url': 'https://example.com', + }, + }, + { + 'title': 'My NSF Login 2', + 'record_type': 'login', + 'folder': 'Projects', + 'fields': { + 'login': 'user2@example.com', + 'password': 'changeme2', + }, + }, + ] + + results = nsf_management.create_nsf_records(vault, RECORDS) + for result in results: + if result.success: + print(f"NSF record created: {result.record_uid} (status: {result.status})") + else: + print(f"NSF record failed: {result.record_uid} ({result.message or result.status})") + + +def nsf_record_add_batch_run(keeper_auth_context: keeper_auth.KeeperAuth) -> None: + vault = open_vault(keeper_auth_context) + try: + nsf_record_add_batch(vault) + except Exception as e: + print(f"Error: {e}") + finally: + close_vault(vault, keeper_auth_context) + + +def main() -> None: + keeper_auth_context, _ = login() + if keeper_auth_context: + nsf_record_add_batch_run(keeper_auth_context) + else: + print("Login failed.") + + +if __name__ == "__main__": + main() diff --git a/keepercli-package/src/keepercli/commands/nsf_commands.py b/keepercli-package/src/keepercli/commands/nsf_commands.py index efd3f0a9..4444c0cd 100644 --- a/keepercli-package/src/keepercli/commands/nsf_commands.py +++ b/keepercli-package/src/keepercli/commands/nsf_commands.py @@ -1,12 +1,13 @@ import argparse import json -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Mapping, Optional, Union from keepersdk.vault import nsf_folder_records, nsf_management, nsf_sharing, vault_record, nsf_common from keepersdk.vault.share_management_utils import parse_nsf_share_expiration from keepersdk.vault.nsf_management import ( NsfError, NsfListRow, + NsfRecordAddSpec, NsfRemovePreviewItem, NsfRemoveResult, ) @@ -156,6 +157,77 @@ def build_nsf_record_data( return _typed_record_to_data(record, title, notes) +def _build_batch_record_data( + mixin: _NsfRecordDataMixin, + context: KeeperParams, + record_type: str, + title: str, + notes: Optional[str], + raw_fields: Union[Mapping[str, Any], List[str], None], +) -> Optional[Dict[str, Any]]: + if raw_fields is None: + return None + if isinstance(raw_fields, Mapping): + return None + record_fields: List[ParsedFieldValue] = [] + for field in raw_fields: + if not isinstance(field, str): + raise base.CommandError('Batch record fields must be strings or a field object') + parsed = RecordEditMixin.parse_field(field) + if parsed.type == 'file': + raise base.CommandError( + 'File attachments are not supported in nsf-record-add batch mode') + record_fields.append(parsed) + return mixin.build_nsf_record_data(context, record_type, title, notes, record_fields) + + +def _load_nsf_record_add_batch_specs( + mixin: _NsfRecordDataMixin, + context: KeeperParams, + batch_path: str, +) -> List[NsfRecordAddSpec]: + try: + with open(batch_path, 'r', encoding='utf-8') as handle: + payload = json.load(handle) + except OSError as exc: + raise base.CommandError(f'Unable to read batch file: {exc}') from exc + except json.JSONDecodeError as exc: + raise base.CommandError(f'Invalid JSON in batch file: {exc}') from exc + + if not isinstance(payload, list): + raise base.CommandError('Batch file must contain a JSON array of record definitions') + if not payload: + raise base.CommandError('Batch file must contain at least one record definition') + + specs: List[NsfRecordAddSpec] = [] + for index, entry in enumerate(payload, start=1): + if not isinstance(entry, dict): + raise base.CommandError(f'Batch record #{index} must be a JSON object') + title = entry.get('title') + record_type = entry.get('record_type') or entry.get('type') + if not title: + raise base.CommandError(f'Batch record #{index} is missing title') + if not record_type: + raise base.CommandError(f'Batch record #{index} is missing record_type') + + notes = entry.get('notes') + raw_fields = entry.get('fields') + record_data = entry.get('record_data') + field_map = raw_fields if isinstance(raw_fields, Mapping) else None + if record_data is None and isinstance(raw_fields, list): + record_data = _build_batch_record_data( + mixin, context, record_type, title, notes, raw_fields) + specs.append(NsfRecordAddSpec( + title=title, + record_type=record_type, + folder_uid=entry.get('folder_uid') or entry.get('folder'), + fields=field_map, + notes=notes, + record_data=record_data, + )) + return specs + + class NsfListCommand(base.ArgparseCommand): def __init__(self): @@ -485,6 +557,10 @@ def add_arguments_to_parser(parser: argparse.ArgumentParser) -> None: parser.add_argument('-n', '--notes', dest='notes', type=str, help='record notes') parser.add_argument('--folder', dest='folder_uid', metavar='FOLDER', type=str, help='folder name or UID to store record') + parser.add_argument( + '--batch-file', dest='batch_file', metavar='FILE', type=str, + help='JSON file containing up to 1000 NSF records per API batch', + ) parser.add_argument('fields', nargs='*', type=str, help='load record type data from strings with dot notation') @@ -494,6 +570,10 @@ def execute(self, context: KeeperParams, **kwargs): prompt_utils.output_text(record_fields_description) return + batch_file = kwargs.get('batch_file') + if batch_file: + return self._execute_batch(context, vault, batch_file, kwargs.get('force') is True) + title = kwargs.get('title') if not title: raise base.CommandError('Title parameter is required.') @@ -547,6 +627,46 @@ def _run(): logger.info('NSF record created: %s', result.record_uid) return result.record_uid + def _execute_batch( + self, + context: KeeperParams, + vault, + batch_file: str, + force: bool) -> List[str]: + specs = _load_nsf_record_add_batch_specs(self, context, batch_file) + if self.warnings: + for w in self.warnings: + logger.warning(w) + if not force: + return [] + + def _run(): + return nsf_management.create_nsf_records(vault, specs) + + results = _wrap_nsf('nsf-record-add', _run) + created: List[str] = [] + failed = 0 + for result in results: + if result.success: + created.append(result.record_uid) + logger.info('NSF record created: %s (%s)', result.record_uid, result.status) + else: + failed += 1 + logger.warning( + 'NSF record add failed: %s (%s)', + result.record_uid, + result.message or result.status, + ) + logger.info( + 'NSF batch add complete: %d created, %d failed, %d total', + len(created), + failed, + len(results), + ) + if failed and not created: + raise base.CommandError('All NSF records in the batch failed to create') + return created + class NsfRecordUpdateCommand(base.ArgparseCommand, _NsfRecordDataMixin): diff --git a/keepersdk-package/src/keepersdk/vault/nsf_management.py b/keepersdk-package/src/keepersdk/vault/nsf_management.py index 090d11a1..e00efb88 100644 --- a/keepersdk-package/src/keepersdk/vault/nsf_management.py +++ b/keepersdk-package/src/keepersdk/vault/nsf_management.py @@ -3,7 +3,7 @@ import json import os from dataclasses import dataclass -from typing import Any, Dict, Iterable, List, Mapping, Optional +from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union from .. import crypto, utils from ..errors import KeeperApiError @@ -14,6 +14,9 @@ ROOT_FOLDER_UID = 'AAAAAAAAAAAAAAAAAPmtNA' """Sentinel UID the server uses for the NSF root folder.""" +NSF_RECORD_ADD_BATCH_LIMIT = 1000 +"""Maximum number of NSF records per ``vault/records/v3/add`` request.""" + class NsfError(ValueError): """Raised when NSF operations cannot proceed (missing cache, bad identifier, etc.).""" @@ -37,6 +40,17 @@ class NsfModifyResult: revision: int = 0 +@dataclass(frozen=True) +class NsfRecordAddSpec: + """Specification for creating one NSF record in a batch add request.""" + title: str + record_type: str + folder_uid: Optional[str] = None + fields: Optional[Mapping[str, Any]] = None + notes: Optional[str] = None + record_data: Optional[Mapping[str, Any]] = None + + @dataclass class NsfFolderModifyResult: folder_uid: str @@ -534,42 +548,157 @@ def _build_record_add_message( return ra -def _build_legacy_record_add_message( - record_uid: str, - record_key: bytes, - data: Dict[str, Any], - auth_data_key: bytes, - folder_uid: Optional[str], - folder_key: Optional[bytes]) -> record_pb2.RecordAdd: - ra = record_pb2.RecordAdd() - ra.record_uid = utils.base64_url_decode(record_uid) - ra.client_modified_time = utils.current_milli_time() - json_bytes = vault_extensions.get_padded_json_bytes(data) - if folder_uid and folder_key: - ra.folder_uid = utils.base64_url_decode(folder_uid) - ra.record_key = crypto.encrypt_aes_v2(record_key, folder_key) - else: - ra.record_key = crypto.encrypt_aes_v2(record_key, auth_data_key) - ra.data = crypto.encrypt_aes_v2(json_bytes, record_key) - return ra +def _normalize_record_add_spec( + spec: Union[NsfRecordAddSpec, Mapping[str, Any]]) -> NsfRecordAddSpec: + if isinstance(spec, NsfRecordAddSpec): + return spec + record_type = spec.get('record_type') or spec.get('type') + if not record_type: + raise NsfError('record_type is required for each batch record spec') + title = spec.get('title') + if not title: + raise NsfError('title is required for each batch record spec') + return NsfRecordAddSpec( + title=title, + record_type=record_type, + folder_uid=spec.get('folder_uid') or spec.get('folder'), + fields=spec.get('fields'), + notes=spec.get('notes'), + record_data=spec.get('record_data'), + ) + + +def _resolve_nsf_folder_for_add(vault: VaultOnline, folder_uid: str) -> str: + resolved = resolve_nsf_folder_uid(vault, folder_uid) or folder_uid + if not is_nsf_folder(vault, resolved): + raise NsfError(f'NSF folder not found: {folder_uid}') + return resolved def _parse_modify_response( response: record_pb2.RecordsModifyResponse, record_uid: str) -> NsfModifyResult: + results = _parse_batch_modify_response(response, [record_uid]) + return results[0] + + +def _parse_batch_modify_response( + response: record_pb2.RecordsModifyResponse, + record_uids: List[str]) -> List[NsfModifyResult]: if not response.records: raise KeeperApiError('no_results', 'No results from record modify response') - for row in response.records: - if utils.base64_url_encode(row.record_uid) == record_uid: - status_name = record_pb2.RecordModifyResult.Name(row.status) - return NsfModifyResult( - record_uid=record_uid, - success=row.status == record_pb2.RS_SUCCESS, - status=status_name, - message=row.message, - revision=getattr(response, 'revision', 0), - ) - raise KeeperApiError('no_results', f'Record {record_uid} not present in modify response') + if len(response.records) != len(record_uids): + raise KeeperApiError( + 'no_results', + f'Expected {len(record_uids)} record results, received {len(response.records)}') + revision = getattr(response, 'revision', 0) + results: List[NsfModifyResult] = [] + for idx, row in enumerate(response.records): + status_name = record_pb2.RecordModifyResult.Name(row.status) + results.append(NsfModifyResult( + record_uid=record_uids[idx], + success=row.status == record_pb2.RS_SUCCESS, + status=status_name, + message=row.message, + revision=revision, + )) + return results + + +def _prepare_nsf_record_add_messages( + vault: VaultOnline, + specs: List[NsfRecordAddSpec], + folder_key_cache: Optional[Dict[str, bytes]] = None, +) -> Tuple[List[str], List[record_endpoints_pb2.RecordAdd]]: + auth = vault.keeper_auth + data_key = auth.auth_context.data_key + cache = folder_key_cache if folder_key_cache is not None else {} + record_uids: List[str] = [] + messages: List[record_endpoints_pb2.RecordAdd] = [] + + for spec in specs: + folder_uid = None + if spec.folder_uid: + folder_uid = _resolve_nsf_folder_for_add(vault, spec.folder_uid) + data = _build_record_data( + spec.record_type, spec.title, spec.fields, spec.notes, spec.record_data) + record_uid = utils.generate_uid() + record_key = os.urandom(32) + folder_key = None + if folder_uid: + if folder_uid not in cache: + cache[folder_uid] = _get_folder_key(vault, folder_uid) + folder_key = cache[folder_uid] + messages.append(_build_record_add_message( + record_uid, record_key, data, data_key, folder_uid, folder_key)) + record_uids.append(record_uid) + return record_uids, messages + + +def _execute_nsf_records_add( + vault: VaultOnline, + record_adds: List[record_endpoints_pb2.RecordAdd], +) -> record_pb2.RecordsModifyResponse: + """Create NSF records via ``vault/records/v3/add`` endpoint.""" + if not record_adds or len(record_adds) > NSF_RECORD_ADD_BATCH_LIMIT: + raise ValueError(f'Provide 1..{NSF_RECORD_ADD_BATCH_LIMIT} records') + + auth = vault.keeper_auth + rq = record_endpoints_pb2.RecordsAddRequest() + rq.clientTime = utils.current_milli_time() + rq.records.extend(record_adds) + response = auth.execute_auth_rest( + 'vault/records/v3/add', rq, response_type=record_pb2.RecordsModifyResponse) + if response is None: + raise KeeperApiError('no_results', 'No results from NSF record add') + return response + + +def create_nsf_records_batch( + vault: VaultOnline, + record_specs: Iterable[Union[NsfRecordAddSpec, Mapping[str, Any]]], + *, + request_sync: bool = True, + folder_key_cache: Optional[Dict[str, bytes]] = None) -> List[NsfModifyResult]: + """Create up to 1000 NSF records in a single ``vault/records/v3/add`` request.""" + specs = [_normalize_record_add_spec(spec) for spec in record_specs] + if not specs: + raise ValueError('At least one record spec is required') + if len(specs) > NSF_RECORD_ADD_BATCH_LIMIT: + raise ValueError(f'Maximum {NSF_RECORD_ADD_BATCH_LIMIT} records at a time') + + record_uids, record_adds = _prepare_nsf_record_add_messages( + vault, specs, folder_key_cache=folder_key_cache) + response = _execute_nsf_records_add(vault, record_adds) + results = _parse_batch_modify_response(response, record_uids) + if request_sync: + vault.sync_requested = True + vault.run_pending_jobs() + return results + + +def create_nsf_records( + vault: VaultOnline, + record_specs: Iterable[Union[NsfRecordAddSpec, Mapping[str, Any]]], + *, + request_sync: bool = True) -> List[NsfModifyResult]: + """Create NSF records, chunking into batches of up to 1000 per API request.""" + specs = [_normalize_record_add_spec(spec) for spec in record_specs] + if not specs: + raise ValueError('At least one record spec is required') + + all_results: List[NsfModifyResult] = [] + folder_key_cache: Dict[str, bytes] = {} + for batch_start in range(0, len(specs), NSF_RECORD_ADD_BATCH_LIMIT): + batch = specs[batch_start:batch_start + NSF_RECORD_ADD_BATCH_LIMIT] + is_last = batch_start + len(batch) >= len(specs) + all_results.extend(create_nsf_records_batch( + vault, + batch, + request_sync=request_sync and is_last, + folder_key_cache=folder_key_cache, + )) + return all_results def create_nsf_record( @@ -583,42 +712,18 @@ def create_nsf_record( record_data: Optional[Mapping[str, Any]] = None, request_sync: bool = True) -> NsfModifyResult: """Create an NSF record.""" - if folder_uid: - resolved = resolve_nsf_folder_uid(vault, folder_uid) or folder_uid - if not is_nsf_folder(vault, resolved): - raise NsfError(f'NSF folder not found: {folder_uid}') - folder_uid = resolved - - data = _build_record_data(record_type, title, fields, notes, record_data) - record_uid = utils.generate_uid() - record_key = os.urandom(32) - auth = vault.keeper_auth - folder_key = _get_folder_key(vault, folder_uid) if folder_uid else None - - ra = _build_record_add_message( - record_uid, record_key, data, auth.auth_context.data_key, folder_uid, folder_key) - rq = record_endpoints_pb2.RecordsAddRequest() - rq.clientTime = utils.current_milli_time() - rq.records.append(ra) - - response = auth.execute_auth_rest( - 'vault/records/v3/add', rq, response_type=record_pb2.RecordsModifyResponse) - if response is None: - legacy_ra = _build_legacy_record_add_message( - record_uid, record_key, data, auth.auth_context.data_key, folder_uid, folder_key) - legacy_rq = record_pb2.RecordsAddRequest() - legacy_rq.client_time = utils.current_milli_time() - legacy_rq.records.append(legacy_ra) - response = auth.execute_auth_rest( - 'vault/records_add', legacy_rq, response_type=record_pb2.RecordsModifyResponse) - assert response is not None - - result = _parse_modify_response(response, record_uid) + spec = NsfRecordAddSpec( + title=title, + record_type=record_type, + folder_uid=folder_uid, + fields=fields, + notes=notes, + record_data=record_data, + ) + results = create_nsf_records_batch(vault, [spec], request_sync=request_sync) + result = results[0] if not result.success: raise KeeperApiError(result.status, result.message) - if request_sync: - vault.sync_requested = True - vault.run_pending_jobs() return result diff --git a/keepersdk-package/unit_tests/test_nsf_record_add_batch.py b/keepersdk-package/unit_tests/test_nsf_record_add_batch.py new file mode 100644 index 00000000..5a5b794a --- /dev/null +++ b/keepersdk-package/unit_tests/test_nsf_record_add_batch.py @@ -0,0 +1,148 @@ +import unittest +from typing import List +from unittest.mock import MagicMock, patch + +from keepersdk import utils +from keepersdk.errors import KeeperApiError +from keepersdk.proto import record_pb2 +from keepersdk.vault import nsf_management + + +class TestNsfRecordAddBatch(unittest.TestCase): + def _vault(self): + vault = MagicMock() + vault.keeper_auth.auth_context.data_key = b'0' * 32 + vault.keeper_auth.execute_auth_rest.return_value = None + return vault + + def _success_response(self, record_uids): + response = record_pb2.RecordsModifyResponse() + for uid in record_uids: + row = response.records.add() + row.record_uid = utils.base64_url_decode(uid) + row.status = record_pb2.RS_SUCCESS + row.message = '' + return response + + @patch.object(nsf_management, 'is_nsf_folder', return_value=True) + @patch.object(nsf_management, 'resolve_nsf_folder_uid', side_effect=lambda _v, uid: uid) + @patch.object(nsf_management, '_get_folder_key', return_value=b'1' * 32) + @patch.object( + nsf_management.utils, + 'generate_uid', + side_effect=[utils.generate_uid() for _ in range(5)], + ) + def test_create_nsf_records_batch_single_request(self, *_mocks): + vault = self._vault() + expected_uids: List[str] = [] + + def _execute(rest_endpoint, request, response_type=None): + expected_uids.extend( + utils.base64_url_encode(record.recordUid) for record in request.records) + return self._success_response(expected_uids) + + vault.keeper_auth.execute_auth_rest.side_effect = _execute + + specs = [ + {'title': f'Record {i}', 'record_type': 'login', 'fields': {'login': f'user{i}'}} + for i in range(3) + ] + results = nsf_management.create_nsf_records_batch( + vault, specs, request_sync=False) + + self.assertEqual(len(results), 3) + self.assertTrue(all(result.success for result in results)) + vault.keeper_auth.execute_auth_rest.assert_called_once() + request = vault.keeper_auth.execute_auth_rest.call_args.args[1] + self.assertEqual(len(request.records), 3) + + def test_create_nsf_records_batch_rejects_over_limit(self): + vault = self._vault() + specs = [ + {'title': f'Record {i}', 'record_type': 'login'} + for i in range(nsf_management.NSF_RECORD_ADD_BATCH_LIMIT + 1) + ] + with self.assertRaisesRegex(ValueError, '1000'): + nsf_management.create_nsf_records_batch(vault, specs, request_sync=False) + + @patch.object(nsf_management, 'is_nsf_folder', return_value=True) + @patch.object(nsf_management, 'resolve_nsf_folder_uid', side_effect=lambda _v, uid: uid) + @patch.object(nsf_management, '_get_folder_key', return_value=b'1' * 32) + @patch.object( + nsf_management.utils, + 'generate_uid', + side_effect=[utils.generate_uid() for _ in range(1001)], + ) + def test_create_nsf_records_chunks_large_input(self, generate_uid_mock, *_mocks): + vault = self._vault() + + call_count = {'value': 0} + + def _execute(rest_endpoint, request, response_type=None): + call_count['value'] += 1 + uids = [utils.base64_url_encode(record.recordUid) for record in request.records] + return self._success_response(uids) + + vault.keeper_auth.execute_auth_rest.side_effect = _execute + + specs = [ + {'title': f'Record {i}', 'record_type': 'login'} + for i in range(1001) + ] + results = nsf_management.create_nsf_records(vault, specs, request_sync=False) + + self.assertEqual(len(results), 1001) + self.assertEqual(call_count['value'], 2) + first_batch_size = vault.keeper_auth.execute_auth_rest.call_args_list[0].args[1].records + second_batch_size = vault.keeper_auth.execute_auth_rest.call_args_list[1].args[1].records + self.assertEqual(len(first_batch_size), 1000) + self.assertEqual(len(second_batch_size), 1) + + @patch.object(nsf_management, 'is_nsf_folder', return_value=True) + @patch.object(nsf_management, 'resolve_nsf_folder_uid', side_effect=lambda _v, uid: uid) + @patch.object(nsf_management, '_get_folder_key', return_value=b'1' * 32) + @patch.object(nsf_management.utils, 'generate_uid', return_value='AAAAAAAAAAAAAAAAAAAAAA') + def test_create_nsf_record_raises_on_failure(self, *_mocks): + vault = self._vault() + response = record_pb2.RecordsModifyResponse() + row = response.records.add() + row.record_uid = utils.base64_url_decode('AAAAAAAAAAAAAAAAAAAAAA') + row.status = record_pb2.RS_ACCESS_DENIED + row.message = 'denied' + vault.keeper_auth.execute_auth_rest.return_value = response + + with self.assertRaises(KeeperApiError): + nsf_management.create_nsf_record( + vault, + title='Test', + record_type='login', + request_sync=False, + ) + + @patch.object(nsf_management.utils, 'generate_uid', return_value='AAAAAAAAAAAAAAAAAAAAAA') + def test_v3_none_response_raises_no_results(self, *_mocks): + vault = self._vault() + vault.keeper_auth.execute_auth_rest.return_value = None + + with self.assertRaises(KeeperApiError) as ctx: + nsf_management.create_nsf_records_batch( + vault, + [{'title': 'Test', 'record_type': 'login'}], + request_sync=False, + ) + self.assertEqual(ctx.exception.result_code, 'no_results') + vault.keeper_auth.execute_auth_rest.assert_called_once() + self.assertEqual( + vault.keeper_auth.execute_auth_rest.call_args.args[0], + 'vault/records/v3/add', + ) + + def test_normalize_record_add_spec_requires_title_and_type(self): + with self.assertRaisesRegex(nsf_management.NsfError, 'title'): + nsf_management._normalize_record_add_spec({'record_type': 'login'}) + with self.assertRaisesRegex(nsf_management.NsfError, 'record_type'): + nsf_management._normalize_record_add_spec({'title': 'Test'}) + + +if __name__ == '__main__': + unittest.main()