From 95213f71a1cc2840ee1eece258728eb94fb6a181 Mon Sep 17 00:00:00 2001 From: Ivan Dimov <78815270+idimov-keeper@users.noreply.github.com> Date: Wed, 17 Jun 2026 21:57:58 -0500 Subject: [PATCH 1/2] Add pam connection ai command for KeeperAI settings on PAM resources. Implements show, set/unset, and remove with sparse DAG merges, configure_resource meta bootstrap, GSE_DELETION removal, and CLI warnings for duplicate or mirrored options. Co-authored-by: Cursor --- .../commands/pam_import/keeper_ai_settings.py | 522 +++++++++++++++++- .../commands/tunnel_and_connections.py | 228 ++++++++ unit-tests/pam/test_dag_layer_b_migration.py | 29 +- unit-tests/pam/test_pam_connection_ai.py | 279 ++++++++++ 4 files changed, 1033 insertions(+), 25 deletions(-) create mode 100644 unit-tests/pam/test_pam_connection_ai.py diff --git a/keepercommander/commands/pam_import/keeper_ai_settings.py b/keepercommander/commands/pam_import/keeper_ai_settings.py index 7d4056c12..bda3693e9 100644 --- a/keepercommander/commands/pam_import/keeper_ai_settings.py +++ b/keepercommander/commands/pam_import/keeper_ai_settings.py @@ -21,11 +21,64 @@ from ... import vault from ...display import bcolors from ...proto import pam_pb2 +from ... import utils from ..pam._layer_b import should_fallback_on_layer_b_error, is_layer_b_feature_disabled from ..tunnel.port_forward.tunnel_helpers import get_config_uid, get_keeper_tokens from ...keeper_dag.crypto import encrypt_aes from keeper_secrets_manager_core.utils import url_safe_str_to_bytes +AI_RISK_LEVELS = ('critical', 'high', 'medium', 'low') +AI_SETTINGS_VERSION = 'v1.0.0' + + +def empty_keeper_ai_settings_dict() -> Dict[str, Any]: + """Vault-compatible default KeeperAI settings (``emptyKeeperAISettings`` in dag-pam-link.ts).""" + empty_allow_deny = {'allow': [], 'deny': []} + return { + 'version': AI_SETTINGS_VERSION, + 'riskLevels': { + 'critical': {'aiSessionTerminate': False, 'tags': dict(empty_allow_deny)}, + 'high': {'aiSessionTerminate': False, 'tags': dict(empty_allow_deny)}, + 'medium': {'aiSessionTerminate': False, 'tags': dict(empty_allow_deny)}, + 'low': {'aiSessionTerminate': False, 'tags': {'allow': []}}, + }, + } + + +def is_default_keeper_ai_settings(settings: Optional[Dict[str, Any]]) -> bool: + """True when settings are absent or match the vault empty/default template.""" + if not settings: + return True + risk_levels = settings.get('riskLevels') + if not isinstance(risk_levels, dict) or not risk_levels: + return True + for level, level_data in risk_levels.items(): + if not isinstance(level_data, dict): + return False + if level_data.get('aiSessionTerminate', False): + return False + tags = level_data.get('tags') + if not isinstance(tags, dict): + continue + if tags.get('allow'): + return False + if level != 'low' and tags.get('deny'): + return False + return True + + +def _find_highest_path_edge(vertex, head_uid: str, dag_path: str): + """Return the highest-version edge for a self-loop DATA path (any edge type).""" + best = None + best_version = -1 + for edge in vertex.edges or []: + if not edge or edge.head_uid != head_uid or edge.path != dag_path: + continue + if edge.version > best_version: + best_version = edge.version + best = edge + return best + def list_resource_data_edges( params: KeeperParams, @@ -129,7 +182,8 @@ def get_resource_settings( params: KeeperParams, resource_uid: str, dag_path: str, - config_uid: Optional[str] = None + config_uid: Optional[str] = None, + quiet_if_missing_vertex: bool = False, ) -> Optional[Dict[str, Any]]: """ Generic function to retrieve settings from a DAG DATA edge with the specified path for a resource. @@ -146,6 +200,8 @@ def get_resource_settings( resource_uid: UID of the PAM resource (pamMachine, pamDatabase, etc.) dag_path: Path of the DATA edge (e.g., 'ai_settings', 'jit_settings') config_uid: Optional PAM config UID. If not provided, will be looked up. + quiet_if_missing_vertex: Log at debug instead of warning when the resource + vertex is absent (expected before first link to a PAM Configuration). Returns: Dictionary containing settings if found, None otherwise. @@ -209,21 +265,17 @@ def get_resource_settings( # Get the resource vertex resource_vertex = linking_dag.get_vertex_by_uid(resource_uid) if not resource_vertex: - logging.warning(f"Resource vertex {resource_uid} not found in DAG") + log = logging.debug if quiet_if_missing_vertex else logging.warning + log(f"Resource vertex {resource_uid} not found in DAG") return None - # Find the DATA edge with the specified path (get highest version, regardless of active status) - settings_edge = None - highest_version = -1 - for edge in resource_vertex.edges: - if edge and (edge.edge_type == EdgeType.DATA and - edge.path == dag_path): - if edge.version > highest_version: - highest_version = edge.version - settings_edge = edge - - if not settings_edge: - logging.debug(f"No '{dag_path}' DATA edge found for resource {resource_uid}") + # Highest-version edge for this path; DELETION means "no settings" (vault GSE_DELETION). + settings_edge = _find_highest_path_edge(resource_vertex, resource_uid, dag_path) + if settings_edge is None or settings_edge.edge_type == EdgeType.DELETION: + logging.debug(f"No active '{dag_path}' DATA edge for resource {resource_uid}") + return None + if settings_edge.edge_type != EdgeType.DATA: + logging.debug(f"Latest '{dag_path}' edge is not DATA for resource {resource_uid}") return None # Get the content from the edge @@ -352,7 +404,8 @@ def get_resource_jit_settings( def get_resource_keeper_ai_settings( params: KeeperParams, resource_uid: str, - config_uid: Optional[str] = None + config_uid: Optional[str] = None, + quiet_if_missing_vertex: bool = False, ) -> Optional[Dict[str, Any]]: """ Retrieve KeeperAI settings from the DAG DATA edge with path 'ai_settings' for a resource. @@ -393,7 +446,10 @@ def get_resource_keeper_ai_settings( } Returns None if settings not found or error occurred. """ - return get_resource_settings(params, resource_uid, 'ai_settings', config_uid) + return get_resource_settings( + params, resource_uid, 'ai_settings', config_uid, + quiet_if_missing_vertex=quiet_if_missing_vertex, + ) def set_resource_keeper_ai_settings( @@ -431,6 +487,10 @@ def set_resource_keeper_ai_settings( return False record_key, resolved_config_uid = common + if not settings: + logging.debug(f"KeeperAI settings empty for {resource_uid}, skipping save") + return False + encrypted_content = encrypt_aes(json.dumps(settings).encode(), record_key) # krouter's configure_resource only writes a settings edge when it loads the @@ -438,11 +498,15 @@ def set_resource_keeper_ai_settings( # carry meta/jit/connection (UserRest.kt). A keeperAiSettings-only request # leaves loopEdges null and the ai_settings write is silently dropped. The Web # Vault avoids this by always sending meta alongside the AI settings, so mirror - # that: include the resource's current meta in the same request. - meta_bytes = None - current_meta = get_resource_settings(params, resource_uid, 'meta', resolved_config_uid) - if isinstance(current_meta, dict): - meta_bytes = json.dumps(current_meta).encode() + # that: include the resource's current meta in the same request. When the + # resource is not in the DAG yet (first link), bootstrap v1 meta instead. + from ..tunnel.port_forward.TunnelGraph import build_resource_meta_v1 + + current_meta = get_resource_settings( + params, resource_uid, 'meta', resolved_config_uid, quiet_if_missing_vertex=True) + if not isinstance(current_meta, dict): + current_meta = build_resource_meta_v1({}, False) + meta_bytes = json.dumps(current_meta).encode() # Primary: Layer-B configure_resource (permission-checked). from ..pam.router_helper import router_configure_resource, get_router_url @@ -453,9 +517,8 @@ def set_resource_keeper_ai_settings( recordUid=url_safe_str_to_bytes(resource_uid), networkUid=url_safe_str_to_bytes(resolved_config_uid), keeperAiSettings=encrypted_content, + meta=meta_bytes, ) - if meta_bytes is not None: - rq.meta = meta_bytes try: router_configure_resource(params, rq) logging.debug(f"Saved KeeperAI settings via configure_resource for {resource_uid}") @@ -579,6 +642,73 @@ def _set_resource_keeper_ai_settings_legacy( return False +def _delete_resource_data_edge_legacy( + params: KeeperParams, + resource_uid: str, + config_uid: str, + record_key: bytes, + dag_path: str, +) -> Optional[bool]: + """Delete a path-scoped self-loop DATA edge via GSE_DELETION (vault ``createDeletionEvent``). + + Returns: + True if a DELETION edge was written, + None if the path was already absent, + False on error. + """ + try: + encrypted_session_token, encrypted_transmission_key, transmission_key = get_keeper_tokens(params) + + dag_record = PasswordRecord() + dag_record.record_uid = config_uid + dag_record.record_key = record_key + + conn = Connection( + params=params, + encrypted_transmission_key=encrypted_transmission_key, + encrypted_session_token=encrypted_session_token, + transmission_key=transmission_key, + use_read_protobuf=False, + use_write_protobuf=False, + ) + + linking_dag = DAG( + conn=conn, + record=dag_record, + graph_id=PamGraphId.PAM.value, + decrypt=True, + ) + linking_dag.load() + + resource_vertex = linking_dag.get_vertex_by_uid(resource_uid) + if not resource_vertex: + logging.debug(f"No resource vertex {resource_uid} in DAG; '{dag_path}' already absent") + return None + + highest = _find_highest_path_edge(resource_vertex, resource_uid, dag_path) + if highest is None or highest.edge_type == EdgeType.DELETION: + logging.debug(f"'{dag_path}' already deleted for resource {resource_uid}") + return None + + for edge in resource_vertex.edges or []: + if (edge and edge.edge_type == EdgeType.DATA and edge.path == dag_path + and edge.head_uid == resource_uid and edge.active): + edge.active = False + + resource_vertex.belongs_to( + resource_vertex, + EdgeType.DELETION, + path=dag_path, + ) + linking_dag.save() + + logging.debug(f"Deleted '{dag_path}' via DELETION edge for resource {resource_uid}") + return True + except Exception as e: + logging.error(f"Error deleting '{dag_path}' for {resource_uid}: {e}", exc_info=True) + return False + + def set_resource_jit_settings( params: KeeperParams, resource_uid: str, @@ -824,6 +954,350 @@ def refresh_link_to_config_to_latest( return False +def _get_audit_user_id(params: KeeperParams) -> str: + if getattr(params, 'account_uid_bytes', None): + return utils.base64_url_encode(params.account_uid_bytes) + return getattr(params, 'user', '') or '' + + +def _make_tag_entry(tag: str, action: str, user_id: str) -> Dict[str, Any]: + return { + 'tag': tag, + 'auditLog': [{ + 'date': utils.current_milli_time(), + 'userId': user_id, + 'action': action, + }], + } + + +def _parse_tag_name(tag_item: Any) -> str: + if isinstance(tag_item, dict): + return str(tag_item.get('tag', '')).strip() + return str(tag_item).strip() + + +def _get_tag_list(level_data: Dict[str, Any], list_name: str) -> List[Dict[str, Any]]: + tags = level_data.setdefault('tags', {}) + entries = tags.get(list_name) + if not isinstance(entries, list): + entries = [] + tags[list_name] = entries + return entries + + +def parse_ai_setting_spec(spec: str) -> tuple: + """ + Parse CLI spec ``LEVEL``, ``LEVEL.SETTING``, or ``LEVEL.SETTING=VALUE``. + + Returns (level, setting_or_none, value_or_none). ``value_or_none`` is None when + ``=`` is absent (unset-without-value forms). + """ + if not spec or not str(spec).strip(): + raise ValueError('empty AI setting spec') + + text = str(spec).strip() + level_part, sep, value_part = text.partition('=') + if sep and value_part == '': + raise ValueError( + f'invalid AI setting spec (missing value): {spec}. ' + f'To remove a setting, use --unset|-u (e.g. -u high.terminate or -u high.allow=chmod).' + ) + + if '.' in level_part: + level_name, setting_name = level_part.split('.', 1) + else: + level_name, setting_name = level_part, None + + level_name = level_name.strip().lower() + if level_name not in AI_RISK_LEVELS: + raise ValueError(f'invalid risk level "{level_name}" (expected: {", ".join(AI_RISK_LEVELS)})') + + if setting_name is not None: + setting_name = setting_name.strip().lower() + if setting_name not in ('terminate', 'allow', 'deny'): + raise ValueError(f'invalid setting "{setting_name}" (expected: terminate, allow, deny)') + if level_name == 'low' and setting_name == 'deny': + raise ValueError('deny is not supported for the low risk level') + + value = value_part if sep else None + return level_name, setting_name, value + + +def dedupe_ai_cli_option_specs( + specs: Optional[List[str]], + option_label: str, +) -> tuple: + """Return unique specs in first-seen order and warnings for duplicate CLI options. + + ``option_label`` is shown in warnings (e.g. ``--set/-s``). + """ + if not specs: + return [], [] + + order: List[str] = [] + counts: Dict[str, int] = {} + for spec in specs: + counts[spec] = counts.get(spec, 0) + 1 + if spec not in order: + order.append(spec) + + warnings = [ + f'duplicate {option_label} ignored: {spec} ({counts[spec]}x)' + for spec in order + if counts[spec] > 1 + ] + return order, warnings + + +def _is_full_ai_setting_spec(spec: str) -> bool: + """True when spec is ``LEVEL.SETTING=VALUE`` (value present after ``=``).""" + _, setting, value = parse_ai_setting_spec(spec) + return setting is not None and value is not None + + +def _reconcile_set_unset_specs( + unset_specs: Optional[List[str]], + set_specs: Optional[List[str]], +) -> tuple: + """Drop ``--unset`` specs mirrored by a ``--set`` on the same ``LEVEL.SETTING=VALUE``. + + Only applies when both sides use the full ``level.setting=value`` form. Partial + unsets (e.g. ``-u high.allow`` or ``-u high.terminate``) are not reconciled. + """ + unset_list = list(unset_specs or []) + set_list = list(set_specs or []) + warnings: List[str] = [] + if not unset_list or not set_list: + return unset_list, set_list, warnings + + parsed_sets = [parse_ai_setting_spec(spec) for spec in set_list] + drop_unset_specs = set() + + for u_spec in unset_list: + if not _is_full_ai_setting_spec(u_spec): + continue + u_level, u_setting, _u_value = parse_ai_setting_spec(u_spec) + + for set_spec, (s_level, s_setting, _s_value) in zip(set_list, parsed_sets): + if u_level != s_level or u_setting != s_setting: + continue + if not _is_full_ai_setting_spec(set_spec): + continue + drop_unset_specs.add(u_spec) + if u_spec == set_spec: + warnings.append(f'--set/-s overrides --unset/-u: {u_spec}') + else: + warnings.append( + f'--set/-s overrides --unset/-u: {u_spec} (mirrors -s {set_spec})' + ) + break + + return [spec for spec in unset_list if spec not in drop_unset_specs], set_list, warnings + + +_LOW_TERMINATE_WARNING = 'risk level low.terminate always defaults to false.' + + +def _parse_terminate_value(value: str) -> bool: + normalized = str(value).strip().casefold() + if normalized == 'true': + return True + if normalized == 'false': + return False + raise ValueError(f'invalid terminate value "{value}" (expected true or false)') + + +def _existing_terminate_value(risk_levels: Dict[str, Any], level: str) -> Optional[bool]: + level_data = risk_levels.get(level) + if not isinstance(level_data, dict): + return None + value = level_data.get('aiSessionTerminate') + if value is None: + return None + return bool(value) + + +def _ensure_ai_settings_dict(existing: Optional[Dict[str, Any]]) -> Dict[str, Any]: + settings = dict(existing) if isinstance(existing, dict) else {} + settings['version'] = settings.get('version') or AI_SETTINGS_VERSION + risk_levels = settings.get('riskLevels') + if not isinstance(risk_levels, dict): + risk_levels = {} + settings['riskLevels'] = risk_levels + return settings + + +def _ensure_level_dict(risk_levels: Dict[str, Any], level: str) -> Dict[str, Any]: + level_data = risk_levels.get(level) + if not isinstance(level_data, dict): + level_data = {} + risk_levels[level] = level_data + return level_data + + +def _prune_level(level_data: Dict[str, Any]) -> bool: + tags = level_data.get('tags') + if isinstance(tags, dict): + for key in list(tags.keys()): + entries = tags.get(key) + if not entries: + tags.pop(key, None) + if not tags: + level_data.pop('tags', None) + return not level_data + + +def _level_is_only_default_low(level_data: Dict[str, Any]) -> bool: + if not isinstance(level_data, dict): + return True + tags = level_data.get('tags') + has_tags = isinstance(tags, dict) and any(tags.get(k) for k in tags) + if has_tags: + return False + terminate = level_data.get('aiSessionTerminate') + return terminate is None or terminate is False + + +def _normalize_ai_settings_for_save(settings: Dict[str, Any]) -> Dict[str, Any]: + risk_levels = settings.get('riskLevels') + if not isinstance(risk_levels, dict): + return {} + + for level in list(risk_levels.keys()): + if level not in AI_RISK_LEVELS: + risk_levels.pop(level, None) + continue + if _prune_level(risk_levels[level]): + risk_levels.pop(level, None) + + low_data = risk_levels.get('low') + if isinstance(low_data, dict): + low_data['aiSessionTerminate'] = False + tags = low_data.get('tags') + if isinstance(tags, dict): + tags.pop('deny', None) + if not tags: + low_data.pop('tags', None) + if _prune_level(low_data) or _level_is_only_default_low(low_data): + risk_levels.pop('low', None) + + if not risk_levels: + return {} + + return { + 'version': settings.get('version', AI_SETTINGS_VERSION), + 'riskLevels': risk_levels, + } + + +def apply_ai_setting_changes( + existing: Optional[Dict[str, Any]], + set_specs: Optional[List[str]], + unset_specs: Optional[List[str]], + params: KeeperParams, +) -> tuple: + """Merge CLI --set/--unset operations into KeeperAI DAG settings. + + Returns ``(settings_dict, warnings)`` where ``warnings`` is a list of user-facing + messages (e.g. silent conversions applied during save). + """ + settings = _ensure_ai_settings_dict(existing) + risk_levels = settings['riskLevels'] + user_id = _get_audit_user_id(params) + warnings: List[str] = [] + + unset_specs, set_specs, reconcile_warnings = _reconcile_set_unset_specs(unset_specs, set_specs) + warnings.extend(reconcile_warnings) + + for spec in unset_specs or []: + level, setting, value = parse_ai_setting_spec(spec) + if setting is None: + risk_levels.pop(level, None) + continue + + level_data = risk_levels.get(level) + if not isinstance(level_data, dict): + continue + + if setting == 'terminate': + level_data.pop('aiSessionTerminate', None) + else: + tags = level_data.get('tags') + if not isinstance(tags, dict): + continue + entries = tags.get(setting) + if not isinstance(entries, list): + continue + if value is None: + tags.pop(setting, None) + else: + tags[setting] = [e for e in entries if _parse_tag_name(e) != value] + if not tags.get(setting): + tags.pop(setting, None) + if not tags: + level_data.pop('tags', None) + + if _prune_level(level_data): + risk_levels.pop(level, None) + + for spec in set_specs or []: + level, setting, value = parse_ai_setting_spec(spec) + if setting is None: + raise ValueError(f'--set requires LEVEL.SETTING=VALUE (got "{spec}")') + if value is None: + raise ValueError(f'--set requires a value (got "{spec}")') + + level_data = _ensure_level_dict(risk_levels, level) + if setting == 'terminate': + terminate_value = _parse_terminate_value(value) + effective_value = False if level == 'low' else terminate_value + if level == 'low' and terminate_value: + if _LOW_TERMINATE_WARNING not in warnings: + warnings.append(_LOW_TERMINATE_WARNING) + if _existing_terminate_value(risk_levels, level) == effective_value: + continue + level_data['aiSessionTerminate'] = effective_value + continue + + tag_value = value.strip() + if not tag_value: + raise ValueError(f'--set tag value cannot be empty (got "{spec}")') + + entries = _get_tag_list(level_data, setting) + if not any(_parse_tag_name(e) == tag_value for e in entries): + action = 'added_to_allow' if setting == 'allow' else 'added_to_deny' + entries.append(_make_tag_entry(tag_value, action, user_id)) + + return _normalize_ai_settings_for_save(settings), warnings + + +def remove_resource_keeper_ai_settings( + params: KeeperParams, + resource_uid: str, + config_uid: Optional[str] = None +) -> Optional[bool]: + """Remove the ``ai_settings`` DATA edge (GSE_DELETION), restoring pre-AI DAG state. + + Web Vault treats a missing ``ai_settings`` edge as ``emptyKeeperAISettings`` in + memory. Writing ``{}`` or the empty template leaves a DATA edge and can break WV; + ``configure_resource`` does not clear ``keeperAiSettings`` when omitted — use graph-sync + DELETION like ``DagOperations.createDeletionEvent`` + ``dagPamLinkAddData``. + + Returns: + True if a DELETION edge was written, + None if ``ai_settings`` was already absent, + False on error. + """ + common = _resolve_resource_settings_inputs(params, resource_uid, {}, config_uid) + if common is None: + return False + record_key, resolved_config_uid = common + + return _delete_resource_data_edge_legacy( + params, resource_uid, resolved_config_uid, record_key, 'ai_settings') + + def print_keeper_ai_settings(params: KeeperParams, resource_uid: str, config_uid: Optional[str] = None): """ Print KeeperAI settings in a human-readable format. @@ -835,7 +1309,7 @@ def print_keeper_ai_settings(params: KeeperParams, resource_uid: str, config_uid """ settings = get_resource_keeper_ai_settings(params, resource_uid, config_uid) - if not settings: + if is_default_keeper_ai_settings(settings): print(f"{bcolors.WARNING}No KeeperAI settings found for resource {resource_uid}{bcolors.ENDC}") return diff --git a/keepercommander/commands/tunnel_and_connections.py b/keepercommander/commands/tunnel_and_connections.py index 0314248b6..743ec280d 100644 --- a/keepercommander/commands/tunnel_and_connections.py +++ b/keepercommander/commands/tunnel_and_connections.py @@ -108,6 +108,7 @@ def __init__(self): # self.register_command('stop', PAMConnectionStopCommand(), 'Stop Connection', 'x') self.register_command('edit', PAMConnectionEditCommand(), 'Edit Connection settings', 'e') self.register_command('jit', PAMConnectionJitCommand(), 'View/update JIT settings', 'j') + self.register_command('ai', PAMConnectionAiCommand(), 'View/update KeeperAI settings', 'a') self.default_verb = 'edit' @@ -3281,6 +3282,233 @@ def _do_set(self, params, kwargs, record_uid, config_uid): print(f'{bcolors.OKGREEN}JIT settings saved successfully.{bcolors.ENDC}') +class PAMConnectionAiCommand(Command): + parser = argparse.ArgumentParser(prog='pam connection ai') + parser.add_argument('record', type=str, action='store', + help='Record UID, path, or title (pamMachine, pamDatabase, or pamDirectory)') + parser.add_argument('--configuration', '-c', type=str, dest='configuration', action='store', default=None, + help='PAM Configuration UID or title (required when record is not linked yet or is linked to 2+ configs)') + parser.add_argument('--set', '-s', dest='set_values', action='append', default=None, + metavar='LEVEL.SETTING=VALUE', + help='Set a KeeperAI value (e.g. -s low.terminate=false -s low.allow=chmod -s "high.deny=kill -9")') + parser.add_argument('--unset', '-u', dest='unset_values', action='append', default=None, + metavar='LEVEL[.SETTING[=VALUE]]', + help='Unset KeeperAI values (e.g. -u low removes the level; -u low.allow removes all allow tags; ' + '-u low.allow=chmod removes one tag)') + parser.add_argument('--remove', dest='remove', action='store_true', default=False, + help='Remove all KeeperAI settings (ai_settings path only)') + parser.add_argument('--show', dest='show', action='store_true', default=False, + help='Show current KeeperAI settings') + + def get_parser(self): + return PAMConnectionAiCommand.parser + + def execute(self, params, **kwargs): + record_name = kwargs.get('record') + configuration = kwargs.get('configuration') + remove_flag = kwargs.get('remove', False) + show_flag = kwargs.get('show', False) + set_values = kwargs.get('set_values') or [] + unset_values = kwargs.get('unset_values') or [] + change_options_provided = bool(set_values or unset_values) + + if remove_flag and show_flag: + raise CommandError('', f'{bcolors.FAIL}--remove cannot be used with --show.{bcolors.ENDC}') + if remove_flag and change_options_provided: + raise CommandError('', f'{bcolors.FAIL}--remove cannot be used with any other option.{bcolors.ENDC}') + if show_flag and change_options_provided: + raise CommandError('', f'{bcolors.FAIL}--show cannot be used with --remove or any KeeperAI option.{bcolors.ENDC}') + if not remove_flag and not show_flag and not change_options_provided: + raise CommandError('', f'{bcolors.FAIL}Provide at least one --set/-s, --unset/-u, --remove, or --show.{bcolors.ENDC}') + + record = RecordMixin.resolve_single_record(params, record_name) + if record is None: + pam_resource_types = {'pamMachine', 'pamDatabase', 'pamDirectory'} + matches = [] + for uid in params.record_cache: + rec = vault.KeeperRecord.load(params, uid) + if rec and rec.record_type in pam_resource_types and rec.title.casefold() == record_name.casefold(): + matches.append(rec) + if len(matches) == 0: + raise CommandError('', f'{bcolors.FAIL}Record "{record_name}" not found.{bcolors.ENDC}') + if len(matches) > 1: + raise CommandError('', + f'{bcolors.FAIL}Multiple records match title "{record_name}"; use UID or path.{bcolors.ENDC}') + record = matches[0] + + if not isinstance(record, vault.TypedRecord) or \ + record.record_type not in ('pamMachine', 'pamDatabase', 'pamDirectory'): + raise CommandError('', + f'{bcolors.FAIL}KeeperAI settings are only supported on pamMachine, pamDatabase, ' + f'and pamDirectory records.{bcolors.ENDC}') + + record_uid = record.record_uid + config_uid = self._resolve_config_uid( + params, record_uid, configuration, allow_unlinked_without_config=show_flag) + + if show_flag: + self._do_show(params, record, record_uid, config_uid) + elif remove_flag: + self._do_remove(params, record_uid, config_uid) + else: + self._do_apply(params, set_values, unset_values, record_uid, config_uid) + + def _resolve_config_uid(self, params, record_uid, configuration, allow_unlinked_without_config=False): + encrypted_session_token, encrypted_transmission_key, _ = get_keeper_tokens(params) + config_leafs = get_dag_leafs(params, encrypted_session_token, encrypted_transmission_key, record_uid) + config_uids = [leaf.get('value') for leaf in (config_leafs or []) if leaf.get('value')] + + if len(config_uids) == 0: + if not configuration: + if allow_unlinked_without_config: + return None + raise CommandError('', + f'{bcolors.FAIL}Record is not linked to a PAM Configuration; ' + f'specify --configuration|-c.{bcolors.ENDC}') + config_uid = self._resolve_configuration_record(params, configuration, linked_config_uids=None) + return config_uid + + if len(config_uids) == 1: + config_uid = config_uids[0] + if configuration: + specified_uid = self._resolve_configuration_record(params, configuration, linked_config_uids=config_uids) + if specified_uid != config_uid: + raise CommandError('', + f'{bcolors.FAIL}PAM Configuration "{configuration}" is not linked to this record.{bcolors.ENDC}') + return config_uid + + if not configuration: + raise CommandError('', + f'{bcolors.FAIL}Record is linked to multiple PAM Configurations; ' + f'specify --configuration|-c.{bcolors.ENDC}') + return self._resolve_configuration_record(params, configuration, linked_config_uids=config_uids) + + @staticmethod + def _resolve_configuration_record(params, configuration, linked_config_uids): + config_rec = RecordMixin.resolve_single_record(params, configuration) + if config_rec is None: + for uid in params.record_cache: + if params.record_cache[uid].get('version', 0) == 6: + r = vault.KeeperRecord.load(params, uid) + if r and r.title.casefold() == configuration.casefold(): + config_rec = r + break + if config_rec is None: + raise CommandError('', + f'{bcolors.FAIL}PAM Configuration "{configuration}" not found.{bcolors.ENDC}') + config_uid = config_rec.record_uid + if linked_config_uids is not None and config_uid not in linked_config_uids: + raise CommandError('', + f'{bcolors.FAIL}PAM Configuration "{configuration}" is not linked to this record.{bcolors.ENDC}') + return config_uid + + def _do_show(self, params, record, record_uid, config_uid): + from .pam_import.keeper_ai_settings import ( + get_resource_keeper_ai_settings, + is_default_keeper_ai_settings, + ) + + print(f'\nKeeperAI Settings for {bcolors.OKBLUE}{record.title}{bcolors.ENDC} ({record_uid}):') + if config_uid is None: + print(f' {bcolors.WARNING}No KeeperAI settings configured ' + f'(record is not linked to a PAM Configuration).{bcolors.ENDC}\n') + return + + settings = get_resource_keeper_ai_settings(params, record_uid, config_uid) + if is_default_keeper_ai_settings(settings): + print(f' {bcolors.WARNING}No KeeperAI settings configured.{bcolors.ENDC}\n') + return + + print(f' Version: {settings.get("version", "unknown")}') + risk_levels = settings.get('riskLevels', {}) + for level in ('critical', 'high', 'medium', 'low'): + level_data = risk_levels.get(level) + if not level_data: + continue + terminate = level_data.get('aiSessionTerminate', False) + tags = level_data.get('tags', {}) if isinstance(level_data.get('tags'), dict) else {} + allow_tags = tags.get('allow', []) + deny_tags = tags.get('deny', []) if level != 'low' else [] + level_color = { + 'critical': bcolors.FAIL, + 'high': bcolors.WARNING, + 'medium': bcolors.OKBLUE, + 'low': bcolors.OKGREEN, + }.get(level, bcolors.ENDC) + print(f'\n {level_color}{level.upper()}{bcolors.ENDC}:') + print(f' Terminate Session: {terminate}') + if allow_tags: + print(' Allow:') + for tag_item in allow_tags: + tag_name = tag_item.get('tag', '') if isinstance(tag_item, dict) else str(tag_item) + print(f' - {tag_name}') + if deny_tags: + print(' Deny:') + for tag_item in deny_tags: + tag_name = tag_item.get('tag', '') if isinstance(tag_item, dict) else str(tag_item) + print(f' - {tag_name}') + print() + + def _do_remove(self, params, record_uid, config_uid): + from .pam_import.keeper_ai_settings import remove_resource_keeper_ai_settings + + result = remove_resource_keeper_ai_settings(params, record_uid, config_uid) + if result is True: + print(f'{bcolors.OKGREEN}KeeperAI settings removed successfully.{bcolors.ENDC}') + elif result is None: + print(f'{bcolors.WARNING}No KeeperAI settings configured.{bcolors.ENDC}') + else: + raise CommandError('', f'{bcolors.FAIL}Failed to remove KeeperAI settings.{bcolors.ENDC}') + + def _do_apply(self, params, set_values, unset_values, record_uid, config_uid): + from .pam_import.keeper_ai_settings import ( + apply_ai_setting_changes, + dedupe_ai_cli_option_specs, + get_resource_keeper_ai_settings, + is_default_keeper_ai_settings, + refresh_link_to_config_to_latest, + refresh_meta_to_latest, + remove_resource_keeper_ai_settings, + set_resource_keeper_ai_settings, + ) + + set_values, set_dup_warnings = dedupe_ai_cli_option_specs(set_values, '--set/-s') + unset_values, unset_dup_warnings = dedupe_ai_cli_option_specs(unset_values, '--unset/-u') + for warning in set_dup_warnings + unset_dup_warnings: + print(f'{bcolors.WARNING}Warning: {warning}{bcolors.ENDC}') + + existing = get_resource_keeper_ai_settings( + params, record_uid, config_uid, quiet_if_missing_vertex=True) + try: + ai_dict, warnings = apply_ai_setting_changes(existing, set_values, unset_values, params) + except ValueError as err: + raise CommandError('', f'{bcolors.FAIL}{err}{bcolors.ENDC}') from err + + for warning in warnings: + print(f'{bcolors.WARNING}Warning: {warning}{bcolors.ENDC}') + + if not ai_dict: + # e.g. -s low.terminate=true|false when nothing is configured: defaults to {}. + if is_default_keeper_ai_settings(existing): + return + result = remove_resource_keeper_ai_settings(params, record_uid, config_uid) + if result is False: + raise CommandError('', f'{bcolors.FAIL}Failed to remove KeeperAI settings.{bcolors.ENDC}') + if result is None: + print(f'{bcolors.WARNING}No KeeperAI settings configured.{bcolors.ENDC}') + else: + print(f'{bcolors.OKGREEN}KeeperAI settings removed successfully.{bcolors.ENDC}') + return + + ok = set_resource_keeper_ai_settings(params, record_uid, ai_dict, config_uid) + if not ok: + raise CommandError('', f'{bcolors.FAIL}Failed to save KeeperAI settings.{bcolors.ENDC}') + + refresh_meta_to_latest(params, record_uid, config_uid) + refresh_link_to_config_to_latest(params, record_uid, config_uid) + print(f'{bcolors.OKGREEN}KeeperAI settings saved successfully.{bcolors.ENDC}') + + class PAMRbiEditCommand(Command): choices = ['on', 'off', 'default'] parser = argparse.ArgumentParser(prog='pam rbi edit') diff --git a/unit-tests/pam/test_dag_layer_b_migration.py b/unit-tests/pam/test_dag_layer_b_migration.py index b7be169d1..3ba9090c7 100644 --- a/unit-tests/pam/test_dag_layer_b_migration.py +++ b/unit-tests/pam/test_dag_layer_b_migration.py @@ -96,6 +96,11 @@ def _capture(params, rq): assert rq.keeperAiSettings == b'CIPHER_BYTES' # Critical: must NOT be set on jitSettings field assert rq.jitSettings == b'' + assert rq.meta == json.dumps({ + 'version': 1, + 'allowedSettings': {}, + 'rotateOnTermination': False, + }).encode() def test_happy_path_bundles_current_meta_so_krouter_persists_ai_edge(self): """Regression: krouter's configure_resource only writes a settings edge @@ -125,8 +130,30 @@ def _capture(params, rq): # the ai_settings edge. Without it the write is a silent no-op. assert rq.meta == json.dumps(meta_dict).encode() # meta is read from the resource's current 'meta' DATA edge. + assert meta_mock.call_args.kwargs.get('quiet_if_missing_vertex') is True + + def test_bootstraps_default_meta_when_resource_not_in_dag_yet(self): + captured = {} + + def _capture(params, rq): + captured['rq'] = rq + return None + + with _patch_inputs(), \ + patch.object(ai_mod, 'encrypt_aes', return_value=b'CIPHER_BYTES'), \ + patch.object(ai_mod, 'get_resource_settings', return_value=None) as meta_mock, \ + patch('keepercommander.commands.pam.router_helper.router_configure_resource', side_effect=_capture): + ok = ai_mod.set_resource_keeper_ai_settings( + _mock_params(), RESOURCE_UID_STR, {'riskLevels': {'high': {}}}, config_uid=CONFIG_UID_STR + ) + assert ok is True + rq = captured['rq'] + assert rq.meta == json.dumps({ + 'version': 1, + 'allowedSettings': {}, + 'rotateOnTermination': False, + }).encode() meta_mock.assert_called_once() - assert meta_mock.call_args.args[2] == 'meta' def test_permission_denied_with_fallback_enabled_calls_legacy(self): legacy_called = {'count': 0} diff --git a/unit-tests/pam/test_pam_connection_ai.py b/unit-tests/pam/test_pam_connection_ai.py new file mode 100644 index 000000000..b97ed8529 --- /dev/null +++ b/unit-tests/pam/test_pam_connection_ai.py @@ -0,0 +1,279 @@ +import pytest +from unittest.mock import patch + +from keepercommander.commands.pam_import.keeper_ai_settings import ( + apply_ai_setting_changes, + dedupe_ai_cli_option_specs, + empty_keeper_ai_settings_dict, + is_default_keeper_ai_settings, + parse_ai_setting_spec, +) + + +class _Params: + account_uid_bytes = b'\x01\x02\x03' + user = 'user@example.com' + + +def test_dedupe_ai_cli_option_specs_reports_counts(): + specs, warnings = dedupe_ai_cli_option_specs( + ['critical.allow=chmod', 'high.allow=wget', 'critical.allow=chmod', 'critical.allow=chmod'], + '--set/-s', + ) + assert specs == ['critical.allow=chmod', 'high.allow=wget'] + assert warnings == ['duplicate --set/-s ignored: critical.allow=chmod (3x)'] + + +def test_dedupe_ai_cli_option_specs_no_warnings_when_unique(): + specs, warnings = dedupe_ai_cli_option_specs(['low'], '--unset/-u') + assert specs == ['low'] + assert warnings == [] + + +def test_is_default_keeper_ai_settings(): + assert is_default_keeper_ai_settings(None) + assert is_default_keeper_ai_settings({}) + assert is_default_keeper_ai_settings(empty_keeper_ai_settings_dict()) + assert not is_default_keeper_ai_settings({ + 'version': 'v1.0.0', + 'riskLevels': {'high': {'aiSessionTerminate': True}}, + }) + + +def test_remove_resource_keeper_ai_settings_emits_deletion(): + params = _Params() + captured = {} + + def _capture(_params, _resource_uid, _config_uid, _record_key, dag_path): + captured['dag_path'] = dag_path + return True + + from keepercommander.commands.pam_import import keeper_ai_settings as ai_mod + with patch.object(ai_mod, '_resolve_resource_settings_inputs', return_value=(b'key', 'config')), \ + patch.object(ai_mod, '_delete_resource_data_edge_legacy', side_effect=_capture): + result = ai_mod.remove_resource_keeper_ai_settings(params, 'resource', 'config') + assert result is True + assert captured['dag_path'] == 'ai_settings' + + +def test_remove_resource_keeper_ai_settings_already_absent(): + params = _Params() + + from keepercommander.commands.pam_import import keeper_ai_settings as ai_mod + with patch.object(ai_mod, '_resolve_resource_settings_inputs', return_value=(b'key', 'config')), \ + patch.object(ai_mod, '_delete_resource_data_edge_legacy', return_value=None): + result = ai_mod.remove_resource_keeper_ai_settings(params, 'resource', 'config') + assert result is None + + +def test_parse_ai_setting_spec_set(): + assert parse_ai_setting_spec('low.terminate=false') == ('low', 'terminate', 'false') + assert parse_ai_setting_spec('high.allow=chmod') == ('high', 'allow', 'chmod') + assert parse_ai_setting_spec('critical.deny=kill -9') == ('critical', 'deny', 'kill -9') + + +def test_parse_ai_setting_spec_unset(): + assert parse_ai_setting_spec('low') == ('low', None, None) + assert parse_ai_setting_spec('medium.allow') == ('medium', 'allow', None) + assert parse_ai_setting_spec('high.deny=wget') == ('high', 'deny', 'wget') + + +def test_parse_ai_setting_spec_rejects_low_deny(): + with pytest.raises(ValueError, match='deny is not supported'): + parse_ai_setting_spec('low.deny=bash') + + +def test_parse_ai_setting_spec_missing_value_suggests_unset(): + with pytest.raises(ValueError, match='--unset\\|-u'): + parse_ai_setting_spec('high.terminate=') + + +def test_apply_ai_setting_changes_merge_without_default_low_stub(): + params = _Params() + result, warnings = apply_ai_setting_changes(None, ['high.allow=chmod'], None, params) + + assert warnings == [] + assert result['version'] == 'v1.0.0' + assert result['riskLevels']['high']['tags']['allow'][0]['tag'] == 'chmod' + assert 'low' not in result['riskLevels'] + + +def test_apply_ai_setting_changes_tag_upsert_no_duplicate(): + params = _Params() + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'high': { + 'tags': {'allow': [{'tag': 'chmod', 'auditLog': [{'action': 'added_to_allow'}]}]}, + }, + }, + } + + result, warnings = apply_ai_setting_changes(existing, ['high.allow=chmod'], None, params) + + assert warnings == [] + assert len(result['riskLevels']['high']['tags']['allow']) == 1 + + +def test_apply_ai_setting_changes_unset_level_and_tag(): + params = _Params() + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'critical': { + 'aiSessionTerminate': True, + 'tags': {'allow': [{'tag': 'mount', 'auditLog': []}]}, + }, + 'low': { + 'aiSessionTerminate': False, + 'tags': {'allow': [{'tag': 'chmod', 'auditLog': []}, {'tag': 'wget', 'auditLog': []}]}, + }, + }, + } + + result, warnings = apply_ai_setting_changes(existing, None, ['critical', 'low.allow=chmod'], params) + + assert warnings == [] + assert 'critical' not in result['riskLevels'] + low_allow = [t['tag'] for t in result['riskLevels']['low']['tags']['allow']] + assert low_allow == ['wget'] + assert result['riskLevels']['low']['aiSessionTerminate'] is False + + +def test_apply_ai_setting_changes_set_terminate_rejects_invalid_value(): + params = _Params() + with pytest.raises(ValueError, match='invalid terminate value "truez"'): + apply_ai_setting_changes(None, ['high.terminate=truez'], None, params) + + +def test_apply_ai_setting_changes_set_terminate_accepts_true_false_only(): + params = _Params() + result_true, _ = apply_ai_setting_changes(None, ['high.terminate=true'], None, params) + assert result_true['riskLevels']['high']['aiSessionTerminate'] is True + + result_false, _ = apply_ai_setting_changes(result_true, ['medium.terminate=FALSE'], None, params) + assert result_false['riskLevels']['medium']['aiSessionTerminate'] is False + + +def test_apply_ai_setting_changes_set_terminate(): + params = _Params() + result, warnings = apply_ai_setting_changes(None, ['medium.terminate=true'], None, params) + assert warnings == [] + assert result['riskLevels']['medium']['aiSessionTerminate'] is True + + +def test_apply_ai_setting_changes_low_terminate_true_warns_and_stays_false(): + params = _Params() + result, warnings = apply_ai_setting_changes(None, ['low.terminate=true'], None, params) + + assert warnings == ['risk level low.terminate always defaults to false.'] + assert result == {} + + +def test_apply_ai_setting_changes_low_terminate_false_is_noop_when_empty(): + params = _Params() + result, warnings = apply_ai_setting_changes(None, ['low.terminate=false'], None, params) + + assert warnings == [] + assert result == {} + assert is_default_keeper_ai_settings(None) + + +def test_apply_ai_setting_changes_paired_unset_set_tag_preserves_audit_log(): + params = _Params() + original_audit = [{'date': 1, 'userId': 'u1', 'action': 'added_to_allow'}] + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'high': { + 'tags': {'allow': [{'tag': 'chmod', 'auditLog': original_audit}]}, + }, + }, + } + + result, warnings = apply_ai_setting_changes( + existing, + ['high.allow=chmod'], + ['high.allow=chmod'], + params, + ) + + assert warnings == ['--set/-s overrides --unset/-u: high.allow=chmod'] + entry = result['riskLevels']['high']['tags']['allow'][0] + assert entry['tag'] == 'chmod' + assert entry['auditLog'] == original_audit + + +def test_apply_ai_setting_changes_paired_unset_set_terminate_noop_when_already_true(): + params = _Params() + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'high': {'aiSessionTerminate': True}, + }, + } + + result, warnings = apply_ai_setting_changes( + existing, + ['high.terminate=true'], + ['high.terminate'], + params, + ) + + assert warnings == [] + assert result['riskLevels']['high']['aiSessionTerminate'] is True + + +def test_apply_ai_setting_changes_mirrored_unset_set_different_tag_values_warns(): + params = _Params() + result, warnings = apply_ai_setting_changes( + None, + ['critical.allow=chmod'], + ['critical.allow=chmo'], + params, + ) + + assert warnings == [ + '--set/-s overrides --unset/-u: critical.allow=chmo (mirrors -s critical.allow=chmod)', + ] + assert result['riskLevels']['critical']['tags']['allow'][0]['tag'] == 'chmod' + + +def test_apply_ai_setting_changes_partial_unset_not_mirrored_with_set(): + params = _Params() + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'critical': { + 'tags': {'allow': [{'tag': 'wget', 'auditLog': []}]}, + }, + }, + } + result, warnings = apply_ai_setting_changes( + existing, + ['critical.allow=chmod'], + ['critical.allow'], + params, + ) + + assert warnings == [] + tags = [t['tag'] for t in result['riskLevels']['critical']['tags']['allow']] + assert tags == ['chmod'] + + +def test_apply_ai_setting_changes_unset_all_leaves_empty(): + params = _Params() + existing = { + 'version': 'v1.0.0', + 'riskLevels': { + 'low': { + 'aiSessionTerminate': False, + 'tags': {'allow': [{'tag': 'chmod', 'auditLog': []}]}, + }, + }, + } + + result, warnings = apply_ai_setting_changes(existing, None, ['low'], params) + + assert warnings == [] + assert result == {} From b97200b70d21a5054a94ac2954e768ca71037c00 Mon Sep 17 00:00:00 2001 From: Ivan Dimov <78815270+idimov-keeper@users.noreply.github.com> Date: Wed, 17 Jun 2026 22:12:47 -0500 Subject: [PATCH 2/2] Extend pam connection ai to pamRemoteBrowser records. Co-authored-by: Cursor --- .../commands/pam_import/keeper_ai_settings.py | 8 +++---- .../commands/tunnel_and_connections.py | 23 ++++++++++++++----- unit-tests/pam/test_pam_connection_ai.py | 9 ++++++++ 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/keepercommander/commands/pam_import/keeper_ai_settings.py b/keepercommander/commands/pam_import/keeper_ai_settings.py index bda3693e9..9a1bc1df2 100644 --- a/keepercommander/commands/pam_import/keeper_ai_settings.py +++ b/keepercommander/commands/pam_import/keeper_ai_settings.py @@ -197,7 +197,7 @@ def get_resource_settings( Args: params: KeeperParams instance - resource_uid: UID of the PAM resource (pamMachine, pamDatabase, etc.) + resource_uid: UID of the PAM resource (pamMachine, pamDatabase, pamDirectory, pamRemoteBrowser) dag_path: Path of the DATA edge (e.g., 'ai_settings', 'jit_settings') config_uid: Optional PAM config UID. If not provided, will be looked up. quiet_if_missing_vertex: Log at debug instead of warning when the resource @@ -392,7 +392,7 @@ def get_resource_jit_settings( Args: params: KeeperParams instance - resource_uid: UID of the PAM resource (pamMachine, pamDatabase, etc.) + resource_uid: UID of the PAM resource (pamMachine, pamDatabase, pamDirectory, pamRemoteBrowser) config_uid: Optional PAM config UID. If not provided, will be looked up. Returns: @@ -419,7 +419,7 @@ def get_resource_keeper_ai_settings( Args: params: KeeperParams instance - resource_uid: UID of the PAM resource (pamMachine, pamDatabase, etc.) + resource_uid: UID of the PAM resource (pamMachine, pamDatabase, pamDirectory, pamRemoteBrowser) config_uid: Optional PAM config UID. If not provided, will be looked up. Returns: @@ -474,7 +474,7 @@ def set_resource_keeper_ai_settings( Args: params: KeeperParams instance - resource_uid: UID of the PAM resource (pamMachine, pamDatabase, etc.) + resource_uid: UID of the PAM resource (pamMachine, pamDatabase, pamDirectory, pamRemoteBrowser) settings: Dictionary containing KeeperAI settings to save config_uid: Optional PAM config UID. If not provided, will be looked up. diff --git a/keepercommander/commands/tunnel_and_connections.py b/keepercommander/commands/tunnel_and_connections.py index 743ec280d..aaa857b64 100644 --- a/keepercommander/commands/tunnel_and_connections.py +++ b/keepercommander/commands/tunnel_and_connections.py @@ -3282,10 +3282,21 @@ def _do_set(self, params, kwargs, record_uid, config_uid): print(f'{bcolors.OKGREEN}JIT settings saved successfully.{bcolors.ENDC}') +_PAM_CONNECTION_AI_RESOURCE_TYPES = frozenset({ + 'pamMachine', + 'pamDatabase', + 'pamDirectory', + 'pamRemoteBrowser', +}) +_PAM_CONNECTION_AI_RESOURCE_TYPES_LABEL = ( + 'pamMachine, pamDatabase, pamDirectory, or pamRemoteBrowser' +) + + class PAMConnectionAiCommand(Command): parser = argparse.ArgumentParser(prog='pam connection ai') parser.add_argument('record', type=str, action='store', - help='Record UID, path, or title (pamMachine, pamDatabase, or pamDirectory)') + help=f'Record UID, path, or title ({_PAM_CONNECTION_AI_RESOURCE_TYPES_LABEL})') parser.add_argument('--configuration', '-c', type=str, dest='configuration', action='store', default=None, help='PAM Configuration UID or title (required when record is not linked yet or is linked to 2+ configs)') parser.add_argument('--set', '-s', dest='set_values', action='append', default=None, @@ -3323,11 +3334,11 @@ def execute(self, params, **kwargs): record = RecordMixin.resolve_single_record(params, record_name) if record is None: - pam_resource_types = {'pamMachine', 'pamDatabase', 'pamDirectory'} matches = [] for uid in params.record_cache: rec = vault.KeeperRecord.load(params, uid) - if rec and rec.record_type in pam_resource_types and rec.title.casefold() == record_name.casefold(): + if rec and rec.record_type in _PAM_CONNECTION_AI_RESOURCE_TYPES \ + and rec.title.casefold() == record_name.casefold(): matches.append(rec) if len(matches) == 0: raise CommandError('', f'{bcolors.FAIL}Record "{record_name}" not found.{bcolors.ENDC}') @@ -3337,10 +3348,10 @@ def execute(self, params, **kwargs): record = matches[0] if not isinstance(record, vault.TypedRecord) or \ - record.record_type not in ('pamMachine', 'pamDatabase', 'pamDirectory'): + record.record_type not in _PAM_CONNECTION_AI_RESOURCE_TYPES: raise CommandError('', - f'{bcolors.FAIL}KeeperAI settings are only supported on pamMachine, pamDatabase, ' - f'and pamDirectory records.{bcolors.ENDC}') + f'{bcolors.FAIL}KeeperAI settings are only supported on ' + f'{_PAM_CONNECTION_AI_RESOURCE_TYPES_LABEL} records.{bcolors.ENDC}') record_uid = record.record_uid config_uid = self._resolve_config_uid( diff --git a/unit-tests/pam/test_pam_connection_ai.py b/unit-tests/pam/test_pam_connection_ai.py index b97ed8529..ea4ad03d1 100644 --- a/unit-tests/pam/test_pam_connection_ai.py +++ b/unit-tests/pam/test_pam_connection_ai.py @@ -277,3 +277,12 @@ def test_apply_ai_setting_changes_unset_all_leaves_empty(): assert warnings == [] assert result == {} + + +def test_pam_connection_ai_resource_types_include_rbi(): + from keepercommander.commands.tunnel_and_connections import _PAM_CONNECTION_AI_RESOURCE_TYPES + + assert 'pamRemoteBrowser' in _PAM_CONNECTION_AI_RESOURCE_TYPES + assert 'pamMachine' in _PAM_CONNECTION_AI_RESOURCE_TYPES + assert 'pamDatabase' in _PAM_CONNECTION_AI_RESOURCE_TYPES + assert 'pamDirectory' in _PAM_CONNECTION_AI_RESOURCE_TYPES