Add support for constrained extension policy

azurelinuxagent/ga/exthandlers.py

@@ -42,7 +42,7 @@
 from azurelinuxagent.common.utils.textutil import redact_sas_token
 from azurelinuxagent.ga.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.ga.policy.policy_engine import ExtensionPolicyEngine, ExtensionDisallowedError, \
-    ExtensionSignaturePolicyError
+    ExtensionSignaturePolicyError, ExtensionRuntimePolicyError
 from azurelinuxagent.common.datacontract import get_properties, set_properties
 from azurelinuxagent.common.errorstate import ErrorState
 from azurelinuxagent.common.event import add_event, elapsed_milliseconds, WALAEventOperation, \
@@ -730,6 +730,13 @@ def handle_ext_handler(self, ext_handler_i, extension, goal_state_id):
             ).format(operation, ext_handler_i.ext_handler.name, conf.get_policy_file_path())
             self.__handle_ext_disallowed_error(ext_handler_i, error_code, report_op=WALAEventOperation.ExtensionSignaturePolicy, message=msg,
                                                extension=extension)
+        except ExtensionRuntimePolicyError as error:
+            _, error_code = _EXT_DISALLOWED_ERROR_MAP.get(ext_handler_i.ext_handler.state)
+            msg = (
+                "Extension will not be processed: {0}"
+            ).format(ustr(error))
+            self.__handle_ext_disallowed_error(ext_handler_i, error_code, report_op=WALAEventOperation.ExtensionPolicy, message=msg,
+                                               extension=extension)
         except PackageValidationError as error:
             code = ExtensionErrorCodes.PluginInstallProcessingFailed   # Signature validation is only done during extension install
             self.__handle_ext_disallowed_error(ext_handler_i, code, report_op=error.operation,
@@ -829,6 +836,9 @@ def handle_enable(self, ext_handler_i, extension):
 
             self.__setup_new_handler(ext_handler_i, extension, self.__should_ignore_signature_validation_errors(ext_handler_i))
 
+            # Create runtime policy file for extension before enabling
+            self._policy_engine.create_runtime_policy_file(ext_handler_i)
+
             if old_ext_handler_i is None:
                 ext_handler_i.install(extension=extension)
             elif ext_handler_i.version_ne(old_ext_handler_i):
@@ -842,6 +852,9 @@ def handle_enable(self, ext_handler_i, extension):
             extension_is_signed = ext_handler_i.signature_validated
             self._policy_engine.check_extension_policy(ext_handler_i.ext_handler.name, extension_is_signed)
 
+            # Create runtime policy file for extension before enabling
+            self._policy_engine.create_runtime_policy_file(ext_handler_i)
+
             ext_handler_i.ensure_consistent_data_for_mc()
             ext_handler_i.update_settings(extension)
 
@@ -2455,6 +2468,9 @@ def get_env_file(self):
     def get_log_dir(self):
         return os.path.join(conf.get_ext_log_dir(), self.ext_handler.name)
 
+    def get_runtime_policy_file(self):
+        return os.path.join(self.get_conf_dir(), 'waagent_runtime_policy.json')
+
     @staticmethod
     def _read_status_file(ext_status_file):
         err_count = 0
@@ -2573,6 +2589,10 @@ def supports_multiple_extensions(self):
         value = self.data['handlerManifest'].get('supportsMultipleExtensions', False)
         return self._parse_boolean_value(value, default_val=False)
 
+    def supports_policy(self):
+        value = self.data['handlerManifest'].get('supportsPolicy', False)
+        return self._parse_boolean_value(value, default_val=False)
+
     def get_resource_limits(self):
         return ResourceLimits(self.data.get('resourceLimits', None))
 

azurelinuxagent/ga/policy/policy_engine.py

@@ -26,6 +26,7 @@
 from azurelinuxagent.common.protocol.extensions_goal_state_from_vm_settings import _CaseFoldedDict
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from azurelinuxagent.ga.confidential_vm_info import ConfidentialVMInfo
+from azurelinuxagent.common.utils import fileutil
 
 
 # Default policy values to be used when customer does not specify these attributes in the policy file.
@@ -77,6 +78,12 @@ def __init__(self):
         super(ExtensionDisallowedError, self).__init__()
 
 
+class ExtensionRuntimePolicyError(PolicyError):
+    """
+    Any error related to extension-specific runtime policy.
+    """
+
+
 class _PolicyEngine(object):
     """
     Implements base policy engine API.
@@ -406,3 +413,47 @@ def check_extension_policy(self, extension_name, extension_is_signed):
         enforce_signature = self.should_enforce_signature_validation(extension_name)
         if enforce_signature and not extension_is_signed:
             raise ExtensionSignaturePolicyError() # Caller sets message and error code, based on requested extension operation
+
+    def get_runtime_policy(self, extension_name):
+        """
+        Return runtime policy for the extension if specified. If not, return None.
+        """
+        extension_policies = self._policy.get("extensionPolicies")
+        if extension_policies is None:
+            return None
+        extensions = extension_policies.get("extensions")
+        if extensions is None:
+            return None
+        individual_policy = extensions.get(extension_name)
+        return individual_policy.get("runtimePolicy") if individual_policy is not None else None
+
+    def create_runtime_policy_file(self, ext_handler_i):
+        """
+        Create the runtime policy file for the extension. The extension manifest includes a "supportsPolicy" attribute
+        indicating whether the extension supports runtime policy enforcement.
+
+        | supportsPolicy | runtimePolicy specified | Result                              |
+        |----------------|-------------------------|-------------------------------------|
+        | True           | Yes                     | Write runtimePolicy to file         |
+        | True           | No                      | Write {} to file                    |
+        | False          | Yes                     | Raise ExtensionRuntimePolicyError   |
+        | False          | No                      | Do nothing                          |
+        """
+        if not self._policy_enforcement_enabled:
+            return
+
+        runtime_policy_file_path = ext_handler_i.get_runtime_policy_file()
+        supports_policy = ext_handler_i.load_manifest().supports_policy()
+        runtime_policy = self.get_runtime_policy(ext_handler_i.ext_handler.name)
+
+        if supports_policy:
+            data_to_write = runtime_policy if runtime_policy is not None else {}
+            try:
+                fileutil.write_file(runtime_policy_file_path, json.dumps(data_to_write))
+            except IOError as e:
+                raise ExtensionRuntimePolicyError("Failed to save runtime policy file: {0}. Error: {1}".format(runtime_policy_file_path, e))
+        elif runtime_policy is not None:
+            raise ExtensionRuntimePolicyError(
+                "Runtime policy is specified for extension '{0}', but this extension does not support policy enforcement. "
+                "To continue, remove the entry '{0}.runtimePolicy' from the policy file ({1}).".format(ext_handler_i.ext_handler.name, conf.get_policy_file_path())
+            )

tests/ga/test_extension.py

@@ -3554,6 +3554,8 @@ def setUp(self):
         self.patch_is_cvm.start()
         self.maxDiff = None     # When long error messages don't match, display the entire diff.
 
+        self.runtime_policy_path = os.path.join(conf.get_lib_dir(), "OSTCExtensions.ExampleHandlerLinux-1.0.0", "config", "waagent_runtime_policy.json")
+
     def tearDown(self):
         self.mock_sleep.stop()
         self.patch_policy_path.stop()
@@ -3875,6 +3877,143 @@ def test_should_save_policy_file_to_history_directory(self):
             with open(file_path, mode='r') as f:
                 self.assertEqual(policy, json.load(f))
 
+    def test_should_create_runtime_policy_file_before_enabling(self):
+        # If "runtimePolicy" is specified in policy file, and "supportsPolicy" is true in handler manifest, create runtime policy file.
+        policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {
+                        "runtimePolicy": {
+                            "allowDirectScripts": False
+                        }
+                    }
+                }
+            }
+        }
+
+        with patch("azurelinuxagent.ga.exthandlers.HandlerManifest.supports_policy", return_value=True):
+            self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=0,
+                                   expected_handler_status='Ready', expected_ext_count=1)
+            self.assertTrue(os.path.exists(self.runtime_policy_path), "Runtime policy file was not created")
+            with open(self.runtime_policy_path, mode='r') as f:
+                expected_policy = {
+                    "allowDirectScripts": False
+                }
+                self.assertEqual(expected_policy, json.load(f))
+
+    def test_should_create_empty_runtime_policy_file_if_not_specified(self):
+        # If "runtimePolicy" is not specified in policy file, and "supportsPolicy" is true in handler manifest, create empty runtime policy file.
+        policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {}
+                }
+            }
+        }
+
+        with patch("azurelinuxagent.ga.exthandlers.HandlerManifest.supports_policy", return_value=True):
+            self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=0,
+                                   expected_handler_status='Ready', expected_ext_count=1)
+            self.assertTrue(os.path.exists(self.runtime_policy_path), "Runtime policy file was not created")
+            with open(self.runtime_policy_path, mode='r') as f:
+                expected_policy = {}
+                self.assertEqual(expected_policy, json.load(f))
+
+    def test_enable_should_fail_if_supportsPolicy_false_but_runtime_policy_specified(self):
+        # If "runtimePolicy" is specified in policy file, and "supportsPolicy" is false in handler manifest, fail extension.
+        policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {
+                        "runtimePolicy": {
+                            "allowDirectScripts": False
+                        }
+                    }
+                }
+            }
+        }
+
+        with patch("azurelinuxagent.ga.exthandlers.HandlerManifest.supports_policy", return_value=False):
+            expected_msg = "Runtime policy is specified for extension 'OSTCExtensions.ExampleHandlerLinux', but this extension does not support policy enforcement."
+            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)
+            self.assertFalse(os.path.exists(self.runtime_policy_path), "Runtime policy file should not have been created")
+
+    def test_should_not_create_runtime_policy_file_if_supportsPolicy_false_and_not_specified(self):
+        # If "runtimePolicy" is not specified in policy file, and "supportsPolicy" is false in handler manifest, do not create runtime policy file.
+        policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {}
+                }
+            }
+        }
+
+        with patch("azurelinuxagent.ga.exthandlers.HandlerManifest.supports_policy", return_value=False):
+            self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=0,
+                                   expected_handler_status='Ready', expected_ext_count=1)
+            self.assertFalse(os.path.exists(self.runtime_policy_path), "Runtime policy file should not have been created")
+
+    def test_should_create_runtime_policy_file_on_each_enable(self):
+        # Runtime policy file should be created fresh on each enable, not just the first one.
+        # This ensures that if the policy changes between enables, the extension gets the updated policy.
+        initial_policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {
+                        "runtimePolicy": {
+                            "allowDirectScripts": False
+                        }
+                    }
+                }
+            }
+        }
+
+        updated_policy = {
+            "policyVersion": "0.1.0",
+            "extensionPolicies": {
+                "extensions": {
+                    "OSTCExtensions.ExampleHandlerLinux": {
+                        "runtimePolicy": {
+                            "allowDirectScripts": True,
+                            "maxExecutionTime": 300
+                        }
+                    }
+                }
+            }
+        }
+
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
+            protocol.aggregate_status = None
+            protocol.report_vm_status = MagicMock()
+            exthandlers_handler = get_exthandlers_handler(protocol)
+
+            with patch("azurelinuxagent.ga.exthandlers.HandlerManifest.supports_policy", return_value=True):
+                # First enable - initial policy
+                self._create_policy_file(initial_policy)
+                exthandlers_handler.run()
+
+                self.assertTrue(os.path.exists(self.runtime_policy_path), "Runtime policy file was not created on first enable")
+                with open(self.runtime_policy_path, mode='r') as f:
+                    self.assertEqual({"allowDirectScripts": False}, json.load(f), "Runtime policy file content does not match initial policy")
+
+                # Second enable - updated policy (extension already installed)
+                self._create_policy_file(updated_policy)
+                protocol.mock_wire_data.set_incarnation(2)
+                protocol.client.update_goal_state()
+                exthandlers_handler.run()
+
+                self.assertTrue(os.path.exists(self.runtime_policy_path), "Runtime policy file was not created on second enable")
+                with open(self.runtime_policy_path, mode='r') as f:
+                    self.assertEqual({"allowDirectScripts": True, "maxExecutionTime": 300}, json.load(f), "Runtime policy file was not updated on second enable")
+
 
 class _TestSignatureValidationBase(TestExtensionBase):
     def setUp(self):