KC-1469: Enterprise root-node isolation is not handled correctly

keepercommander/commands/enterprise.py

@@ -1244,19 +1244,44 @@ def execute(self, params, **kwargs):
             if not matched_nodes:
                 raise CommandError('enterprise-node', 'No nodes to toggle.')
 
+            toggled_nodes = []
             for mn in matched_nodes:
                 node_id = mn['node_id']
                 data = mn['data']
-                displayname = data['displayname']
+                displayname = data.get('displayname') or str(node_id)
+                was_isolated = bool(mn.get('restrict_visibility'))
+                is_root = not mn.get('parent_id')
                 request = enterprise_pb2.SetRestrictVisibilityRequest()
-                request.nodeId = node_id
+                # Root isolation is an enterprise-level flag returned in
+                # GeneralDataEntity rather than on the root Node entity.
+                request.nodeId = 0 if is_root else node_id
                 try:
                     api.communicate_rest(params, request, 'enterprise/set_restrict_visibility')
-                    mn['restrict_visibility'] = not (mn.get('restrict_visibility') or False)
-                    logging.warning('good result: {}'.format(displayname))
+                    toggled_nodes.append((node_id, displayname, was_isolated))
                 except Exception as e:
                     logging.warning('node \"%s\": toggle isolation failed: %s', displayname, e)
-            api.query_enterprise(params)
+            if toggled_nodes:
+                api.query_enterprise(params, force=True)
+                refreshed_nodes = {
+                    x['node_id']: x for x in (params.enterprise or {}).get('nodes', [])
+                }
+                for node_id, displayname, was_isolated in toggled_nodes:
+                    refreshed_node = refreshed_nodes.get(node_id)
+                    if not refreshed_node:
+                        logging.warning(
+                            'node \"%s\": isolation toggle could not be verified after refresh',
+                            displayname)
+                        continue
+                    is_isolated = bool(refreshed_node.get('restrict_visibility'))
+                    if is_isolated == was_isolated:
+                        logging.warning(
+                            'node \"%s\": server accepted the isolation toggle, '
+                            'but the state did not change',
+                            displayname)
+                    else:
+                        logging.info(
+                            'node \"%s\": isolation is now %s',
+                            displayname, 'enabled' if is_isolated else 'disabled')
         else:
             for node_name in unmatched_nodes:
                 logging.warning('Node \'%s\' is not found: Skipping', node_name)

keepercommander/enterprise.py

@@ -185,6 +185,7 @@ def load(self, params):  # type: (KeeperParams) -> None
 
             params.enterprise['keys'] = keys
         entities = set()
+        root_restrict_visibility = None
         while True:
             rq = proto.EnterpriseDataRequest()
             if self._continuationToken:
@@ -202,6 +203,8 @@ def load(self, params):  # type: (KeeperParams) -> None
                 params.enterprise['enterprise_name'] = self._enterprise.enterprise_name
                 if rs.generalData.distributor:
                     params.enterprise['distributor'] = True
+            if rs.HasField('generalData'):
+                root_restrict_visibility = rs.generalData.restrictVisibility
 
             for ed in rs.data:
                 entities.add(ed.entity)
@@ -212,6 +215,11 @@ def load(self, params):  # type: (KeeperParams) -> None
             self._continuationToken = rs.continuationToken
             if not rs.hasMore:
                 break
+        if root_restrict_visibility is not None:
+            root_node = next((x for x in params.enterprise.get('nodes', []) if not x.get('parent_id')), None)
+            if root_node:
+                _set_or_remove(root_node, 'restrict_visibility',
+                               True if root_restrict_visibility else None)
         if proto.MANAGED_NODES in entities:
             try:
                 self.load_missing_role_keys(params)
@@ -456,8 +464,9 @@ def to_keeper_entity(self, proto_entity, keeper_entity):  # type: (proto.Node, d
         _set_or_remove(keeper_entity, 'rsa_enabled', True if proto_entity.rsaEnabled else None)
         _set_or_remove(keeper_entity, 'sso_service_provider_id',
                        proto_entity.ssoServiceProviderId if proto_entity.ssoServiceProviderId > 0 else None)
-        _set_or_remove(keeper_entity, 'restrict_visibility',
-                       proto_entity.restrictVisibility if proto_entity.restrictVisibility else None)
+        if keeper_entity.get('parent_id'):
+            _set_or_remove(keeper_entity, 'restrict_visibility',
+                           proto_entity.restrictVisibility if proto_entity.restrictVisibility else None)
 
         data = {}
         if 'encrypted_data' in keeper_entity:

unit-tests/test_command_enterprise.py

@@ -5,9 +5,10 @@
 from unittest import TestCase, mock
 
 from data_enterprise import EnterpriseEnvironment, get_enterprise_data, enterprise_allocate_ids
-from keepercommander import api, crypto, utils, vault
+from keepercommander import api, crypto, enterprise as enterprise_data, utils, vault
 from keepercommander.params import KeeperParams, PublicKeys
 from keepercommander.error import CommandError
+from keepercommander.proto import enterprise_pb2
 from data_vault import VaultEnvironment, get_connected_params
 from keepercommander.commands import enterprise, aram
 
@@ -47,6 +48,41 @@ def test_get_enterprise_public_key(self):
         self.assertEqual(params.enterprise['unencrypted_tree_key'], ent_env.tree_key)
         self.assertEqual(len(params.enterprise['nodes']), 2)
 
+    def test_general_data_restrict_visibility_controls_root_node(self):
+        params = get_connected_params()
+        api.query_enterprise(params)
+        params.enterprise['keys'] = {}
+        root = next(x for x in params.enterprise['nodes'] if x['node_id'] == ent_env.node1_id)
+        child = next(x for x in params.enterprise['nodes'] if x['node_id'] == ent_env.node2_id)
+        root['restrict_visibility'] = True
+        child['restrict_visibility'] = True
+
+        response = enterprise_pb2.EnterpriseDataResponse()
+        response.generalData.enterpriseName = params.enterprise['enterprise_name']
+        response.generalData.restrictVisibility = False
+        response.hasMore = False
+
+        loader = enterprise_data._EnterpriseLoader(params.enterprise['unencrypted_tree_key'])
+        with mock.patch('keepercommander.enterprise.api.communicate_rest', return_value=response):
+            loader.load(params)
+
+        self.assertNotIn('restrict_visibility', root)
+        self.assertTrue(child['restrict_visibility'])
+
+        response.generalData.restrictVisibility = True
+        with mock.patch('keepercommander.enterprise.api.communicate_rest', return_value=response):
+            loader.load(params)
+
+        self.assertTrue(root['restrict_visibility'])
+        self.assertTrue(child['restrict_visibility'])
+
+        response = enterprise_pb2.EnterpriseDataResponse()
+        response.hasMore = False
+        with mock.patch('keepercommander.enterprise.api.communicate_rest', return_value=response):
+            loader.load(params)
+
+        self.assertTrue(root['restrict_visibility'])
+
     def test_enterprise_info_command(self):
         params = get_connected_params()
         api.query_enterprise(params)
@@ -164,6 +200,93 @@ def test_enterprise_node_move_sets_selected_parent_id(self):
         request = execute_batch.call_args.args[1][0]
         self.assertEqual(request['parent_id'], ent_env.node1_id)
 
+    def test_enterprise_node_toggle_root_isolation(self):
+        for was_isolated in (False, True):
+            with self.subTest(was_isolated=was_isolated):
+                params = get_connected_params()
+                api.query_enterprise(params)
+                root = next(x for x in params.enterprise['nodes']
+                            if x['node_id'] == ent_env.node1_id)
+                root['data']['displayname'] = 'Enterprise 1'
+                if was_isolated:
+                    root['restrict_visibility'] = True
+
+                def refresh_enterprise(p, force=False, tree_key=None):
+                    self.assertTrue(force)
+                    refreshed_root = next(x for x in p.enterprise['nodes']
+                                          if x['node_id'] == ent_env.node1_id)
+                    if was_isolated:
+                        refreshed_root.pop('restrict_visibility', None)
+                    else:
+                        refreshed_root['restrict_visibility'] = True
+
+                cmd = enterprise.EnterpriseNodeCommand()
+                with mock.patch(
+                        'keepercommander.commands.enterprise.api.communicate_rest'
+                ) as communicate_rest, mock.patch(
+                        'keepercommander.commands.enterprise.api.query_enterprise',
+                        side_effect=refresh_enterprise
+                ) as query_enterprise:
+                    cmd.execute(
+                        params,
+                        node=[str(ent_env.node1_id)],
+                        toggle_isolated=True,
+                    )
+
+                request = communicate_rest.call_args.args[1]
+                self.assertEqual(request.nodeId, 0)
+                query_enterprise.assert_called_once_with(params, force=True)
+
+    def test_enterprise_node_toggle_child_isolation(self):
+        params = get_connected_params()
+        api.query_enterprise(params)
+
+        def refresh_enterprise(p, force=False, tree_key=None):
+            self.assertTrue(force)
+            child = next(x for x in p.enterprise['nodes']
+                         if x['node_id'] == ent_env.node2_id)
+            child['restrict_visibility'] = True
+
+        cmd = enterprise.EnterpriseNodeCommand()
+        with mock.patch(
+                'keepercommander.commands.enterprise.api.communicate_rest'
+        ) as communicate_rest, mock.patch(
+                'keepercommander.commands.enterprise.api.query_enterprise',
+                side_effect=refresh_enterprise
+        ):
+            cmd.execute(
+                params,
+                node=[str(ent_env.node2_id)],
+                toggle_isolated=True,
+            )
+
+        request = communicate_rest.call_args.args[1]
+        self.assertEqual(request.nodeId, ent_env.node2_id)
+
+    def test_enterprise_node_toggle_isolation_reports_noop(self):
+        params = get_connected_params()
+        api.query_enterprise(params)
+        child = next(x for x in params.enterprise['nodes']
+                     if x['node_id'] == ent_env.node2_id)
+        child['restrict_visibility'] = True
+
+        cmd = enterprise.EnterpriseNodeCommand()
+        with mock.patch(
+                'keepercommander.commands.enterprise.api.communicate_rest'
+        ), mock.patch(
+                'keepercommander.commands.enterprise.api.query_enterprise'
+        ), self.assertLogs(level=logging.WARNING) as logs:
+            cmd.execute(
+                params,
+                node=[str(ent_env.node2_id)],
+                toggle_isolated=True,
+            )
+
+        self.assertTrue(any(
+            'server accepted the isolation toggle, but the state did not change' in message
+            for message in logs.output
+        ))
+
     def test_enterprise_add_user(self):
         params = get_connected_params()
         api.query_enterprise(params)