enterprise admin device and account lock/unlock actions

keepercli-package/src/keepercli/cli.py

@@ -1,12 +1,12 @@
 import logging
 import sys
-from typing import Optional, Any, Iterable, List
+from typing import Optional, Any, Iterable, List, Callable
 
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import History
 
 from . import prompt_utils, api, autocomplete
-from .commands import command_completer, base, command_history
+from .commands import command_completer, base, command_history, command_visibility
 from .helpers import report_utils
 from .params import KeeperParams, KeeperConfig
 from keepersdk import constants
@@ -30,6 +30,27 @@ def store_string(self, string: str) -> None:
             command_history.append(string)
 
 
+_SCOPE_DISPLAY_NAMES = {
+    base.CommandScope.Account: 'Account Commands',
+    base.CommandScope.Vault: 'Vault Commands',
+    base.CommandScope.DeviceManagement: 'Device Management Commands',
+    base.CommandScope.Enterprise: 'Enterprise Commands',
+    base.CommandScope.MSP: 'MSP Commands',
+    base.CommandScope.Distributor: 'Distributor Commands',
+    base.CommandScope.Common: 'Miscellaneous Commands',
+}
+
+_SCOPE_DISPLAY_ORDER = (
+    base.CommandScope.Account,
+    base.CommandScope.Vault,
+    base.CommandScope.DeviceManagement,
+    base.CommandScope.Enterprise,
+    base.CommandScope.MSP,
+    base.CommandScope.Distributor,
+    base.CommandScope.Common,
+)
+
+
 def do_command(command_line: str, context: KeeperParams, commands: base.CliCommands) -> Any:
     cmd, sep, args = command_line.partition(' ')
     orig_cmd = cmd
@@ -44,7 +65,7 @@ def do_command(command_line: str, context: KeeperParams, commands: base.CliComma
         command, _ = commands.commands[cmd]
         return command.execute_args(context, args.strip(), command=orig_cmd)
     else:
-        display_command_help(commands)
+        display_command_help(commands, context)
     return None
 
 
@@ -81,7 +102,8 @@ def get_prompt() -> str:
         if sys.stdin.isatty() and sys.stdout.isatty():
             from prompt_toolkit.enums import EditingMode
             from prompt_toolkit.shortcuts import CompleteStyle
-            completer = command_completer.CommandCompleter(commands, autocomplete.standard_completer(context))
+            completer = command_completer.CommandCompleter(
+                commands, autocomplete.standard_completer(context), context_getter=lambda: context)
             prompt_session = PromptSession(
                 multiline=False, editing_mode=EditingMode.EMACS, complete_style=CompleteStyle.MULTI_COLUMN,
                 complete_while_typing=False, completer=completer, auto_suggest=None, key_bindings=prompt_utils.kb,
@@ -174,18 +196,31 @@ def get_prompt() -> str:
     context.clear_session()
     return 0
 
-def display_command_help(commands: base.CliCommands):
+def display_command_help(commands: base.CliCommands, context: Optional[KeeperParams] = None):
     alias_lookup = {x[1]: x[0] for x in commands.aliases.items()}
-    all_scopes = {x[1]: x[1].name for x in commands.commands.values()}
-    scopes = sorted(all_scopes.keys())
+    available_scopes = {value[1] for value in commands.commands.values()}
     headers = ['', 'Command', 'Alias', '', 'Description']
     table = []
-    for scope in scopes:
-        scope_commands = [key for key, value in commands.commands.items() if value[1] == scope]
+    for scope in _SCOPE_DISPLAY_ORDER:
+        if scope not in available_scopes:
+            continue
+        scope_commands = [
+            key for key, value in commands.commands.items()
+            if value[1] == scope and command_visibility.is_command_visible(key, context)
+        ]
+        if not scope_commands:
+            continue
+        scope_name = _SCOPE_DISPLAY_NAMES.get(scope, scope.name)
         idx = 0
         for cmd in sorted(scope_commands):
             c = commands.commands[cmd][0]
-            table.append([all_scopes[scope] if idx == 0 else '', cmd, alias_lookup.get(cmd) or '', '...', c.description()])
+            table.append([
+                scope_name if idx == 0 else '',
+                cmd,
+                alias_lookup.get(cmd) or '',
+                '...',
+                c.description(),
+            ])
             idx += 1
 
     prompt_utils.output_text('\nCommands:')

keepercli-package/src/keepercli/commands/base.py

@@ -63,6 +63,7 @@ def description(self):
 class CommandScope(enum.IntFlag):
     Account = enum.auto()
     Vault = enum.auto()
+    DeviceManagement = enum.auto()
     Enterprise = enum.auto()
     MSP = enum.auto()
     Distributor = enum.auto()

keepercli-package/src/keepercli/commands/command_completer.py

@@ -4,14 +4,17 @@
 
 from . import base
 from .. import autocomplete
+from .command_visibility import is_command_visible
 
 
 class CommandCompleter(completion.Completer):
     def __init__(self,
                  command_collection: base.CommandCollection,
-                 on_complete: Optional[Callable[[str, str], Iterable[str]]] = None) -> None:
+                 on_complete: Optional[Callable[[str, str], Iterable[str]]] = None,
+                 context_getter: Optional[Callable[[], object]] = None) -> None:
         self.commands = command_collection
         self.on_complete = on_complete
+        self.context_getter = context_getter
 
     def get_completions(self, document, complete_event):
         if not document.is_cursor_at_the_end:
@@ -33,7 +36,11 @@ def get_completions(self, document, complete_event):
         command = self.commands.get_command_by_name(cmd)
         if command is None:
             if len(tokens) == 0 and document.char_before_cursor != ' ':
-                cmds = [x for x in self.commands.query_commands(cmd)]
+                context = self.context_getter() if self.context_getter else None
+                cmds = [
+                    x for x in self.commands.query_commands(cmd)
+                    if is_command_visible(x, context)
+                ]
                 cmds.sort()
                 for c in cmds:
                     yield completion.Completion(c, start_position=-len(document.text))

keepercli-package/src/keepercli/commands/command_visibility.py

@@ -0,0 +1,20 @@
+from typing import Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..params import KeeperParams
+
+DEVICE_ADMIN_COMMANDS = frozenset({'device-admin-list', 'device-admin-action'})
+
+
+def is_enterprise_admin(context: Optional['KeeperParams']) -> bool:
+    return bool(
+        context
+        and context.auth
+        and context.auth.auth_context.is_enterprise_admin
+    )
+
+
+def is_command_visible(command: str, context: Optional['KeeperParams']) -> bool:
+    if command in DEVICE_ADMIN_COMMANDS:
+        return is_enterprise_admin(context)
+    return True

keepercli-package/src/keepercli/commands/device_management.py

@@ -4,6 +4,7 @@
 from datetime import datetime
 from typing import Callable, Dict, List, Optional
 
+from keepersdk import errors
 from keepersdk.authentication import device_management
 
 from . import base
@@ -34,9 +35,35 @@ def _format_timestamp(dt: Optional[datetime]) -> str:
 
 
 def _sdk_error(exc: Exception) -> base.CommandError:
+    if isinstance(exc, errors.KeeperApiError):
+        if device_management.is_device_api_unavailable(exc):
+            return base.CommandError(device_management.DEVICE_FEATURE_UNAVAILABLE_MESSAGE)
     return base.CommandError(str(exc))
 
 
+def _validate_admin_enterprise_user_ids(
+    context: KeeperParams,
+    enterprise_user_ids: List[int],
+) -> List[int]:
+    """Return enterprise user IDs known to enterprise data; warn when IDs are not found."""
+    base.require_enterprise_admin(context)
+    resolved: List[int] = []
+    seen: set[int] = set()
+    for user_id in enterprise_user_ids:
+        if user_id in seen:
+            continue
+        seen.add(user_id)
+        if context.enterprise_data.users.get_entity(user_id) is None:
+            logger.warning(
+                "Warning: No enterprise_user_id found matching '%s'", user_id
+            )
+        else:
+            resolved.append(user_id)
+    if not resolved and enterprise_user_ids:
+        logger.info('No matching enterprise_user_id found')
+    return resolved
+
+
 def _run_device_action_command(
     context: KeeperParams,
     device_identifiers: List[str],
@@ -47,7 +74,7 @@ def _run_device_action_command(
     try:
         for name in action_fn(context.auth, device_identifiers):
             logger.info(success_message, name)
-    except ValueError as e:
+    except (ValueError, errors.KeeperApiError) as e:
         raise _sdk_error(e) from e
 
 
@@ -79,7 +106,7 @@ def _display_admin_devices(
     """Fetch and print the admin device list table for the given enterprise user IDs."""
     try:
         devices = device_management.list_admin_devices(context.auth, enterprise_user_ids)
-    except ValueError as e:
+    except (ValueError, errors.KeeperApiError) as e:
         raise _sdk_error(e) from e
 
     if not devices:
@@ -186,6 +213,31 @@ def _display_admin_devices(
         'handler': device_management.remove_admin_user_devices,
         'action_verb': 'removed',
     },
+    'lock': {
+        'description': (
+            'Lock the device for all users on the devices and the associated auto linked devices. '
+            'Logout all users from the device'
+        ),
+        'handler': device_management.lock_admin_user_devices,
+        'action_verb': 'locked',
+    },
+    'unlock': {
+        'description': (
+            'Unlock the devices and the associated auto linked devices for the calling user'
+        ),
+        'handler': device_management.unlock_admin_user_devices,
+        'action_verb': 'unlocked',
+    },
+    'account-lock': {
+        'description': 'Lock the device for the user only. If user is logged in, logout',
+        'handler': device_management.account_lock_admin_user_devices,
+        'action_verb': 'account locked',
+    },
+    'account-unlock': {
+        'description': 'Unlock the device for the user',
+        'handler': device_management.account_unlock_admin_user_devices,
+        'action_verb': 'account unlocked',
+    },
 }
 
 DEVICE_ADMIN_ACTION_CHOICES = list(DEVICE_ADMIN_ACTION_DEFINITIONS.keys())
@@ -231,7 +283,7 @@ def execute(self, context: KeeperParams, **kwargs):
         base.require_login(context)
         try:
             devices = device_management.list_user_devices(context.auth)
-        except ValueError as e:
+        except (ValueError, errors.KeeperApiError) as e:
             raise _sdk_error(e) from e
 
         if not devices:
@@ -277,7 +329,7 @@ class DeviceRenameCommand(base.ArgparseCommand):
     def __init__(self):
         parser = argparse.ArgumentParser(
             prog='device-rename',
-            description='Rename a device for the current user',
+            description='Rename user devices',
         )
         DeviceRenameCommand.add_arguments_to_parser(parser)
         super().__init__(parser)
@@ -302,7 +354,7 @@ def execute(self, context: KeeperParams, **kwargs):
             logger.info("Device name updated from '%s' to '%s'", old_name, updated_name)
             logger.info('')
             _display_user_devices(context, title_prefix='Updated ')
-        except ValueError as e:
+        except (ValueError, errors.KeeperApiError) as e:
             raise _sdk_error(e) from e
 
 
@@ -391,19 +443,23 @@ def add_arguments_to_parser(parser: argparse.ArgumentParser):
             'enterprise_user_ids',
             nargs='+',
             type=int,
-            help='List of Enterprise User IDs (required). You can get enterprise user IDs by running "ei --users" command',
+            help='List of Enterprise User IDs (required). You can get enterprise user IDs by running "enterprise-info user"',
         )
         parser.error = base.ArgparseCommand.raise_parse_exception
         parser.exit = base.ArgparseCommand.suppress_exit
 
     def execute(self, context: KeeperParams, **kwargs):
         """Display admin device list in table or JSON format for the given enterprise user IDs."""
         base.require_enterprise_admin(context)
-        enterprise_user_ids = kwargs.get('enterprise_user_ids') or []
+        enterprise_user_ids = _validate_admin_enterprise_user_ids(
+            context, kwargs.get('enterprise_user_ids') or []
+        )
+        if not enterprise_user_ids:
+            return
 
         try:
             devices = device_management.list_admin_devices(context.auth, enterprise_user_ids)
-        except ValueError as e:
+        except (ValueError, errors.KeeperApiError) as e:
             raise _sdk_error(e) from e
 
         if not devices:
@@ -437,7 +493,7 @@ class DeviceAdminActionCommand(base.ArgparseCommand):
     def __init__(self):
         parser = argparse.ArgumentParser(
             prog='device-admin-action',
-            description='Perform various action on one or more devices that the Admin has control of.',
+            description='Perform actions on devices across enterprise users',
         )
         DeviceAdminActionCommand.add_arguments_to_parser(parser)
         super().__init__(parser)
@@ -488,7 +544,12 @@ def execute(self, context: KeeperParams, **kwargs):
         """Run the requested admin device action and refresh the device list."""
         base.require_enterprise_admin(context)
         action = kwargs.get('action')
-        enterprise_user_id = kwargs.get('enterprise_user_id')
+        validated_user_ids = _validate_admin_enterprise_user_ids(
+            context, [kwargs.get('enterprise_user_id')]
+        )
+        if not validated_user_ids:
+            return
+        enterprise_user_id = validated_user_ids[0]
         devices = kwargs.get('devices') or []
         config = DEVICE_ADMIN_ACTION_DEFINITIONS.get(action or '')
         if not config:
@@ -506,7 +567,7 @@ def execute(self, context: KeeperParams, **kwargs):
                     "Device action successfully completed: '%s' %s for user %s",
                     name, action_verb, enterprise_user_id,
                 )
-        except ValueError as e:
+        except (ValueError, errors.KeeperApiError) as e:
             raise _sdk_error(e) from e
 
         logger.info('Updated device list for user %s:', enterprise_user_id)

keepercli-package/src/keepercli/register_commands.py

@@ -23,9 +23,9 @@ def register_commands(commands: base.CliCommands, scopes: Optional[base.CommandS
         commands.register_command('biometric', BiometricCommand(), base.CommandScope.Account)
         commands.register_command('logout', account_commands.LogoutCommand(), base.CommandScope.Account)
         commands.register_command('this-device', account_commands.ThisDeviceCommand(), base.CommandScope.Account)
-        commands.register_command('device-list', device_management.DeviceListCommand(), base.CommandScope.Account)
-        commands.register_command('device-action', device_management.DeviceActionCommand(), base.CommandScope.Account)
-        commands.register_command('device-rename', device_management.DeviceRenameCommand(), base.CommandScope.Account)
+        commands.register_command('device-list', device_management.DeviceListCommand(), base.CommandScope.DeviceManagement)
+        commands.register_command('device-action', device_management.DeviceActionCommand(), base.CommandScope.DeviceManagement)
+        commands.register_command('device-rename', device_management.DeviceRenameCommand(), base.CommandScope.DeviceManagement)
         commands.register_command('whoami', account_commands.WhoamiCommand(), base.CommandScope.Account)
         commands.register_command('reset-password', account_commands.ResetPasswordCommand(), base.CommandScope.Account)
         commands.register_command('2fa', two_fa.TwoFaCommand(), base.CommandScope.Account)
@@ -127,8 +127,8 @@ def register_commands(commands: base.CliCommands, scopes: Optional[base.CommandS
         commands.register_command('download-membership', importer_commands.DownloadMembershipCommand(), base.CommandScope.Enterprise)
         commands.register_command('apply-membership', importer_commands.ApplyMembershipCommand(), base.CommandScope.Enterprise)
         commands.register_command('device-approve', enterprise_user.EnterpriseDeviceApprovalCommand(), base.CommandScope.Enterprise)
-        commands.register_command('device-admin-list', device_management.DeviceAdminListCommand(), base.CommandScope.Enterprise)
-        commands.register_command('device-admin-action', device_management.DeviceAdminActionCommand(), base.CommandScope.Enterprise)
+        commands.register_command('device-admin-list', device_management.DeviceAdminListCommand(), base.CommandScope.DeviceManagement)
+        commands.register_command('device-admin-action', device_management.DeviceAdminActionCommand(), base.CommandScope.DeviceManagement)
         commands.register_command('pedm', pedm_admin.PedmCommand(), base.CommandScope.Enterprise)
         commands.register_command('msp-down', msp.MspDownCommand(), base.CommandScope.Enterprise, 'md')
         commands.register_command('msp-info', msp.MspInfoCommand(), base.CommandScope.Enterprise, 'mi')