Block extensions disallowed by policy

azurelinuxagent/ga/policy/policy_engine.py

@@ -36,12 +36,6 @@
 _MAX_SUPPORTED_POLICY_VERSION = "0.1.0"
 
 
-class PolicyError(AgentError):
-    """
-    Error raised during agent policy enforcement.
-    """
-
-
 class InvalidPolicyError(AgentError):
     """
     Error raised if user-provided policy is invalid.
@@ -50,7 +44,6 @@ def __init__(self, msg, inner=None):
         msg = "Customer-provided policy file ('{0}') is invalid, please correct the following error: {1}".format(conf.get_policy_file_path(), msg)
         super(InvalidPolicyError, self).__init__(msg, inner)
 
-
 class _PolicyEngine(object):
     """
     Implements base policy engine API.
@@ -61,6 +54,7 @@ def __init__(self):
         if not self.policy_enforcement_enabled:
             return
 
+        _PolicyEngine._log_policy_event("Policy enforcement is enabled.")
         self._policy = self._parse_policy(self.__read_policy())
 
     @staticmethod
@@ -98,8 +92,10 @@ def __read_policy():
         with open(conf.get_policy_file_path(), 'r') as f:
             try:
                 contents = f.read()
+                # TODO: Consider copying the policy file contents to the history folder, and only log the policy locally
+                # in the case of policy-related failure.
                 _PolicyEngine._log_policy_event(
-                    "Policy enforcement is enabled. Enforcing policy using policy file found at '{0}'. File contents:\n{1}"
+                    "Enforcing policy using policy file found at '{0}'. File contents:\n{1}"
                     .format(conf.get_policy_file_path(), contents))
                 # json.loads will raise error if file contents are not a valid json (including empty file).
                 custom_policy = json.loads(contents)

tests/data/test_waagent.conf

@@ -140,4 +140,4 @@ OS.SshDir=/notareal/path
 # - The default is false to protect the state of existing VMs
 OS.EnableFirewall=n
 
-Debug.EnableExtensionPolicy=y
+Debug.EnableExtensionPolicy=n

tests/ga/test_extension.py

@@ -427,7 +427,7 @@ def test_migration_ignores_tree_remove_errors(self, shutil_mock):  # pylint: dis
 class TestExtensionBase(AgentTestCase):
     def _assert_handler_status(self, report_vm_status, expected_status,
                                expected_ext_count, version,
-                               expected_handler_name="OSTCExtensions.ExampleHandlerLinux", expected_msg=None):
+                               expected_handler_name="OSTCExtensions.ExampleHandlerLinux", expected_msg=None, expected_code=None):
         self.assertTrue(report_vm_status.called)
         args, kw = report_vm_status.call_args  # pylint: disable=unused-variable
         vm_status = args[0]
@@ -443,6 +443,9 @@ def _assert_handler_status(self, report_vm_status, expected_status,
         if expected_msg is not None:
             self.assertIn(expected_msg, handler_status.message)
 
+        if expected_code is not None:
+            self.assertEqual(expected_code, handler_status.code)
+
 
 # Deprecated. New tests should be added to the TestExtension class
 @patch('time.sleep', side_effect=lambda _: mock_sleep(0.001))
@@ -1649,13 +1652,13 @@ def test_extensions_disabled(self, _, *args):
         vm_status = args[0]
         self.assertEqual(1, len(vm_status.vmAgent.extensionHandlers))
         exthandler = vm_status.vmAgent.extensionHandlers[0]
-        self.assertEqual(-1, exthandler.code)
+        self.assertEqual(ExtensionErrorCodes.PluginEnableProcessingFailed, exthandler.code)
         self.assertEqual('NotReady', exthandler.status)
-        self.assertEqual("Extension will not be processed since extension processing is disabled. To enable extension processing, set Extensions.Enabled=y in '/etc/waagent.conf'", exthandler.message)
+        self.assertEqual("Extension 'OSTCExtensions.ExampleHandlerLinux' will not be processed since extension processing is disabled. To enable extension processing, set Extensions.Enabled=y in '/etc/waagent.conf'", exthandler.message)
         ext_status = exthandler.extension_status
-        self.assertEqual(-1, ext_status.code)
+        self.assertEqual(ExtensionErrorCodes.PluginEnableProcessingFailed, ext_status.code)
         self.assertEqual('error', ext_status.status)
-        self.assertEqual("Extension will not be processed since extension processing is disabled. To enable extension processing, set Extensions.Enabled=y in '/etc/waagent.conf'", ext_status.message)
+        self.assertEqual("Extension 'OSTCExtensions.ExampleHandlerLinux' will not be processed since extension processing is disabled. To enable extension processing, set Extensions.Enabled=y in '/etc/waagent.conf'", ext_status.message)
 
     def test_extensions_deleted(self, *args):
         # Ensure initial enable is successful
@@ -3507,5 +3510,211 @@ def test_report_msg_if_handler_manifest_contains_invalid_values(self):
                 self.assertIn("'supportsMultipleExtensions' has a non-boolean value", kw_messages[2]['message'])
 
 
+class TestExtensionPolicy(TestExtensionBase):
+    def setUp(self):
+        AgentTestCase.setUp(self)
+        self.policy_path = os.path.join(self.tmp_dir, "waagent_policy.json")
+
+        # Patch attributes to enable policy feature
+        self.patch_policy_path = patch('azurelinuxagent.common.conf.get_policy_file_path',
+                                       return_value=str(self.policy_path))
+        self.patch_policy_path.start()
+        self.patch_conf_flag = patch('azurelinuxagent.ga.policy.policy_engine.conf.get_extension_policy_enabled',
+                                     return_value=True)
+        self.patch_conf_flag.start()
+        self.maxDiff = None     # When long error messages don't match, display the entire diff.
+
+    def tearDown(self):
+        patch.stopall()
+        AgentTestCase.tearDown(self)
+
+    def _create_policy_file(self, policy):
+        with open(self.policy_path, mode='w') as policy_file:
+            if isinstance(policy, dict):
+                json.dump(policy, policy_file, indent=4)
+            else:
+                policy_file.write(policy)
+            policy_file.flush()
+
+    def _test_policy_case(self, policy, op, expected_status_code, expected_handler_status, expected_ext_count,
+                             expected_status_msg=None):
+
+        # Set up a mock protocol instance.
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
+            if op == ExtensionRequestedState.Uninstall:
+                # Generate a new mock goal state to uninstall the extension - increment the incarnation
+                protocol.mock_wire_data.set_incarnation(2)
+                protocol.mock_wire_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
+                protocol.client.update_goal_state()
+            protocol.aggregate_status = None
+            protocol.report_vm_status = MagicMock()
+            exthandlers_handler = get_exthandlers_handler(protocol)
+
+            # Create policy file and process extensions.
+            self._create_policy_file(policy)
+            exthandlers_handler.run()
+            exthandlers_handler.report_ext_handlers_status()
+
+            # Assert that agent is reporting the expected handler status
+            report_vm_status = protocol.report_vm_status
+            self.assertTrue(report_vm_status.called)
+            self._assert_handler_status(report_vm_status, expected_handler_status, expected_ext_count=expected_ext_count,
+                                        version="1.0.0", expected_handler_name='OSTCExtensions.ExampleHandlerLinux',
+                                        expected_msg=expected_status_msg, expected_code=expected_status_code)
+
+    def test_should_fail_enable_if_extension_disallowed(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+               }
+            }
+        expected_msg = "failed to run extension 'OSTCExtensions.ExampleHandlerLinux' because it is not specified as an allowed extension."
+        self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                          expected_handler_status='NotReady', expected_ext_count=1, expected_status_msg=expected_msg)
+
+    def test_should_fail_enable_for_invalid_policy(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": "False"
+                }
+            }
+        expected_msg = "attribute 'extensionPolicies.allowListedExtensionsOnly'; must be 'boolean'"
+        self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                          expected_handler_status='NotReady', expected_ext_count=1, expected_status_msg=expected_msg)
+
+    def test_should_fail_extension_if_error_thrown_during_policy_engine_init(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0"
+            }
+        with patch('azurelinuxagent.ga.policy.policy_engine.ExtensionPolicyEngine.__init__',
+                   side_effect=Exception("mock exception")):
+            expected_msg = "Extension will not be processed: mock exception"
+            self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled,
+                                      expected_status_code=ExtensionErrorCodes.PluginEnableProcessingFailed,
+                                      expected_handler_status='NotReady', expected_ext_count=1, expected_status_msg=expected_msg)
+
+    def test_should_fail_uninstall_if_extension_disallowed(self):
+        policy = \
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {}
+                },
+            }
+        expected_msg = "failed to uninstall extension 'OSTCExtensions.ExampleHandlerLinux' because it is not specified as an allowed extension."
+        self._test_policy_case(policy=policy, op=ExtensionRequestedState.Uninstall, expected_status_code=ExtensionErrorCodes.PluginDisableProcessingFailed,
+                                  expected_handler_status='NotReady', expected_ext_count=1, expected_status_msg=expected_msg)
+
+    def test_should_fail_enable_if_dependent_extension_disallowed(self):
+        self._create_policy_file({
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "extensions": {
+                        "OSTCExtensions.ExampleHandlerLinux": {}
+                    }
+                }
+            })
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE_EXT_SEQUENCING) as protocol:
+            protocol.aggregate_status = None
+            protocol.report_vm_status = MagicMock()
+            exthandlers_handler = get_exthandlers_handler(protocol)
+            dep_ext_level_2 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux")
+            dep_ext_level_1 = extension_emulator(name="OSTCExtensions.OtherExampleHandlerLinux")
+
+            exthandlers_handler.run()
+            exthandlers_handler.report_ext_handlers_status()
+
+            # OtherExampleHandlerLinux should be disallowed by policy, ExampleHandlerLinux should be skipped because
+            # dependent extension failed
+            self._assert_handler_status(protocol.report_vm_status, expected_status="NotReady", expected_ext_count=1,
+                                        version="1.0.0", expected_handler_name="OSTCExtensions.OtherExampleHandlerLinux",
+                                        expected_msg=("failed to run extension 'OSTCExtensions.OtherExampleHandlerLinux' "
+                                                      "because it is not specified as an allowed extension."))
+
+            self._assert_handler_status(protocol.report_vm_status, expected_status="NotReady", expected_ext_count=0,
+                                        version="1.0.0", expected_handler_name="OSTCExtensions.ExampleHandlerLinux",
+                                        expected_msg="Skipping processing of extensions since execution of dependent "
+                                                     "extension OSTCExtensions.OtherExampleHandlerLinux failed")
+
+            # check handler list and dependency levels
+            self.assertTrue(exthandlers_handler.ext_handlers is not None)
+            self.assertTrue(exthandlers_handler.ext_handlers is not None)
+            self.assertEqual(len(exthandlers_handler.ext_handlers), 2)
+            self.assertEqual(1, next(handler for handler in exthandlers_handler.ext_handlers if
+                                     handler.name == dep_ext_level_1.name).settings[0].dependencyLevel)
+            self.assertEqual(2, next(handler for handler in exthandlers_handler.ext_handlers if
+                                     handler.name == dep_ext_level_2.name).settings[0].dependencyLevel)
+
+    def test_enable_should_succeed_if_extension_allowed(self):
+        policy_cases = [
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": False,
+                }
+            },
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "extensions": {
+                        "OSTCExtensions.ExampleHandlerLinux": {}
+                    }
+                }
+            }
+        ]
+        for policy in policy_cases:
+            self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=0,
+                                   expected_handler_status='Ready', expected_ext_count=1)
+
+    def test_uninstall_should_succeed_if_extension_allowed(self):
+        policy_cases = [
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": False,
+                }
+            },
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "extensions": {
+                        "OSTCExtensions.ExampleHandlerLinux": {}
+                    }
+                }
+            }
+        ]
+        for policy in policy_cases:
+            with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
+                # Generate a new mock goal state to uninstall the extension - increment the incarnation
+                protocol.mock_wire_data.set_incarnation(2)
+                protocol.mock_wire_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
+                protocol.client.update_goal_state()
+                protocol.aggregate_status = None
+                protocol.report_vm_status = MagicMock()
+                exthandlers_handler = get_exthandlers_handler(protocol)
+
+                # Create policy file and process extensions.
+                self._create_policy_file(policy)
+                exthandlers_handler.run()
+                exthandlers_handler.report_ext_handlers_status()
+
+                # Assert that no status is being reported for the extension, to confirm that uninstall was successful.
+                report_vm_status = protocol.report_vm_status
+                self.assertTrue(report_vm_status.called)
+                args, kw = report_vm_status.call_args  # pylint: disable=unused-variable
+                vm_status = args[0]
+                self.assertEqual(0, len(vm_status.vmAgent.extensionHandlers))
+
+
 if __name__ == '__main__':
     unittest.main()