Repository: Azure/WALinuxAgent
Pull request #3246 (2024-10-21 20:13:41)

Alternative implementation of the policy document parser
Discussed with @mgunnala an alternative approach to write the parse for policy documents. It is similar in concept to a recursive descendent parser and produces code that is simpler to write, read, and debug.

We also decided to make the document version a required attribute.

Lastly, I added code to also parse the runtimePolicy for extensions.

narrieta: @mgunnala  - FYI: as part of my latest changes, I changed ExtensionPolicyEngine.should_allow_extension() and ExtensionPolicyEngine.should_enforce_signature_validation() to take an extension name, rather than an Extension object. This is in order to avoid the implicit dependency on azurelinuxagent.common.protocol.restapi; there is no need to depend on that module.

Patch:
diff --git a/azurelinuxagent/ga/policy/policy_engine.py b/azurelinuxagent/ga/policy/policy_engine.py
index 8d93e132..065589f6 100644
--- a/azurelinuxagent/ga/policy/policy_engine.py
+++ b/azurelinuxagent/ga/policy/policy_engine.py
@@ -15,8 +15,8 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
-import copy
 import json
+import re
 import os
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common import logger
@@ -26,25 +26,10 @@ from azurelinuxagent.common.exception import AgentError
 from azurelinuxagent.common.protocol.extensions_goal_state_from_vm_settings import _CaseFoldedDict
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 
-
-# Schema for policy file.
-_POLICY_SCHEMA = \
-    {
-        "policyVersion": ustr,
-        "extensionPolicies": {
-            "allowListedExtensionsOnly": bool,
-            "signatureRequired": bool,
-            "extensions": {
-                "<extensionName>": {
-                    "signatureRequired": bool
-                }
-            }
-        }
-    }
-
 # Default policy values to be used when customer does not specify these attributes in the policy file.
 _DEFAULT_ALLOW_LISTED_EXTENSIONS_ONLY = False
 _DEFAULT_SIGNATURE_REQUIRED = False
+_DEFAULT_EXTENSIONS = {}
 
 # Agent supports up to this version of the policy file ("policyVersion" in schema).
 # Increment this number when any new attributes are added to the policy schema.
@@ -72,25 +57,11 @@ class _PolicyEngine(object):
     """
     def __init__(self):
         # Set defaults for policy
-        self._policy = \
-        {
-            "policyVersion": _MAX_SUPPORTED_POLICY_VERSION,
-            "extensionPolicies": {
-                "allowListedExtensionsOnly": _DEFAULT_ALLOW_LISTED_EXTENSIONS_ONLY,
-                "signatureRequired": _DEFAULT_SIGNATURE_REQUIRED,
-                "extensions": {}
-            }
-        }
-
         self._policy_enforcement_enabled = self.__get_policy_enforcement_enabled()
         if not self.policy_enforcement_enabled:
             return
 
-        # Use a copy of the policy as a template. Update the template as we parse the custom policy.
-        template = copy.deepcopy(self._policy)
-        custom_policy = self.__read_policy()
-        self.__parse_policy(template, custom_policy)
-        self._policy = template
+        self._policy = self._parse_policy(self.__read_policy())
 
     @staticmethod
     def _log_policy_event(msg, is_success=True, op=WALAEventOperation.Policy, send_event=True):
@@ -142,148 +113,173 @@ class _PolicyEngine(object):
             return custom_policy
 
     @staticmethod
-    def __parse_policy(template, custom_policy):
+    def _parse_policy(policy):
         """
-        Update template with attributes specified in custom_policy:
-            - attributes provided in custom_policy override the default values in template
-            - if an attribute is not provided, use the default value
-            - if an unrecognized attribute is present in custom_policy (not defined in _POLICY_SCHEMA), raise an error
-            - if an attribute does not match the type specified in the schema, raise an error
+        Parses the given policy document and returns an equivalent document that has been populated with default values and verified for correctness, i.e.
+        that conforms the following schema:
+
+            {
+                 "policyVersion": "0.1.0",
+                 "extensionPolicies": {
+                     "allowListedExtensionsOnly": <true, false>, [Optional; default: false]
+                     "signatureRequired": <true, false>, [Optional; default: false]
+                     "extensions": { [Optional; default: {} (empty)]
+                         "<extension_name>": {
+                             "signatureRequired": <true, false> [Optional; no default]
+                             "runtimePolicy": <extension-specific policy> [Optional; no default]
+                         }
+                     },
+                 }
+             }
+
+        Raises InvalidPolicyError if the policy document is invalid.
         """
-        # Validate top level attributes and then parse each section of the custom policy.
-        # Individual parsing functions are responsible for validating schema of that section (any nested dicts).
-        # Note that validation must happen before parsing.
-        _PolicyEngine.__validate_schema(custom_policy, _POLICY_SCHEMA)
-        _PolicyEngine.__parse_version(template, custom_policy)
-        _PolicyEngine.__parse_extension_policies(template, custom_policy)
+        if not isinstance(policy, dict):
+            raise InvalidPolicyError("expected an object describing a Policy; got {0}.".format(type(policy).__name__))
+
+        _PolicyEngine._check_attributes(policy, object_name="policy", valid_attributes=["policyVersion", "extensionPolicies"])
+
+        return {
+            "policyVersion": _PolicyEngine._parse_policy_version(policy),
+            "extensionPolicies": _PolicyEngine._parse_extension_policies(policy)
+        }
 
     @staticmethod
-    def __parse_version(template, policy):
+    def _parse_policy_version(policy):
         """
-        Validate and return "policyVersion" attribute. If not a string in the format "x.y.z", raise InvalidPolicyError.
+        Validate and return "policyVersion" attribute. If not a string in the format "major[.minor[.patch]]", raise InvalidPolicyError.
         If policy_version is greater than maximum supported version, raise InvalidPolicyError.
         """
-        version = policy.get("policyVersion")
-        if version is None:
-            return
+        version = _PolicyEngine._get_string(policy, attribute="policyVersion")
 
-        try:
-            flexible_version = FlexibleVersion(version)
-        except ValueError:
-            raise InvalidPolicyError(
-                "invalid value for attribute 'policyVersion' attribute 'policyVersion' is expected to be in format 'major.minor.patch' "
-                                         "(e.g., '1.0.0'). Please change to a valid value.")
+        if not re.match(r"^\d+(\.\d+(\.\d+)?)?$", version):
+            raise InvalidPolicyError("invalid value for attribute 'policyVersion': '{0}'; it should be in format 'major[.minor[.patch]]' (e.g., '1', '1.0', '1.0.0')".format(version))
 
-        if FlexibleVersion(_MAX_SUPPORTED_POLICY_VERSION) < flexible_version:
-            raise InvalidPolicyError("policy version '{0}' is not supported. The agent supports policy versions up to '{1}'. Please provide a compatible policy version."
-                                     .format(version, _MAX_SUPPORTED_POLICY_VERSION))
+        if FlexibleVersion(_MAX_SUPPORTED_POLICY_VERSION) < FlexibleVersion(version):
+            raise InvalidPolicyError("policy version '{0}' is not supported. The agent supports policy versions up to '{1}'.".format(version, _MAX_SUPPORTED_POLICY_VERSION))
 
-        template["policyVersion"] = version
+        return version
 
     @staticmethod
-    def __parse_extension_policies(template, policy):
-        extension_policies = policy.get("extensionPolicies")
-        if extension_policies is not None:
-            _PolicyEngine.__validate_schema(extension_policies, _POLICY_SCHEMA["extensionPolicies"], "extensionPolicies")
-
-            # Parse allowlist policy
-            allowlist_policy = extension_policies.get("allowListedExtensionsOnly")
-            if allowlist_policy is not None:
-                template["extensionPolicies"]["allowListedExtensionsOnly"] = allowlist_policy
+    def _parse_extension_policies(policy):
+        """
+        Parses the "extensionPolicies" attribute of the policy document. See _parse_policy() for schema.
+        """
+        extension_policies = _PolicyEngine._get_dictionary(policy, attribute="extensionPolicies", optional=True, default={})
 
-            # Parse global signature policy
-            signature_policy = extension_policies.get("signatureRequired")
-            if signature_policy is not None:
-                template["extensionPolicies"]["signatureRequired"] = signature_policy
+        _PolicyEngine._check_attributes(extension_policies, object_name="extensionPolicies", valid_attributes=["allowListedExtensionsOnly", "signatureRequired", "extensions"])
 
-            # Parse individual extension policies
-            _PolicyEngine.__parse_extensions(template, extension_policies)
+        return {
+            "allowListedExtensionsOnly": _PolicyEngine._get_boolean(extension_policies, attribute="allowListedExtensionsOnly", name_prefix="extensionPolicies.", optional=True, default=_DEFAULT_ALLOW_LISTED_EXTENSIONS_ONLY),
+            "signatureRequired": _PolicyEngine._get_boolean(extension_policies, attribute="signatureRequired", name_prefix="extensionPolicies.", optional=True, default=_DEFAULT_SIGNATURE_REQUIRED),
+            "extensions": _PolicyEngine._parse_extensions(
+                _PolicyEngine._get_dictionary(extension_policies, attribute="extensions", name_prefix="extensionPolicies.", optional=True, default=_DEFAULT_EXTENSIONS)
+            )
+        }
 
     @staticmethod
-    def __parse_extensions(template, extensions_policy):
+    def _parse_extensions(extensions):
         """
-        Parse "extensions" dict and update in template.
-        "extensions" is expected to be in the format:
-        {
-            "extensions": {
-                "<extensionName>": {
-                    "signatureRequired": bool
-                }
-            }
-        }
+        Parses the "extensions" attribute. See _parse_policy() for schema.
+        The return value is a case-folded dict. CRP allows extensions to be any case, so we allow for case-insensitive lookup of individual extension policies.
+        """
+        parsed = _CaseFoldedDict.from_dict({})
 
-        If "signatureRequired" isn't provided, the global "signatureRequired" value will be used instead.
-        The "extensions" attribute will be converted to a case-folded dict. CRP allows extensions to be any
-        case, so we use case-folded dict to allow for case-insensitive lookup of individual extension policies.
+        for extension, extension_policy in extensions.items():
+            if not isinstance(extension_policy, dict):
+                raise InvalidPolicyError("invalid type {0} for attribute 'extensionPolicies.extensions.{1}'; must be 'object'".format(type(extension_policy).__name__, extension))
+            parsed[extension] = _PolicyEngine._parse_extension(extension_policy)
+
+        return parsed
+
+    @staticmethod
+    def _parse_extension(extension):
         """
-        extensions = extensions_policy.get("extensions")
-        if extensions is None:
-            return
+        Parses an individual extension. See _parse_policy() for schema.
+        """
+        extension_attribute_name = "extensionPolicies.extensions.{0}".format(extension)
 
-        parsed_extensions_dict = {}
-        for extension_name, individual_policy in extensions.items():
+        _PolicyEngine._check_attributes(extension, object_name=extension_attribute_name, valid_attributes=["signatureRequired", "runtimePolicy"])
 
-            # We don't validate "extensions" against the schema, because the attributes (individual extension names)
-            # are dynamic and not defined in the schema. We do validate that individual_policy is a dict, and validate
-            # the schema of individual_policy.
-            individual_policy_schema = _POLICY_SCHEMA["extensionPolicies"]["extensions"]["<extensionName>"]
-            if not isinstance(individual_policy, dict):
-                raise InvalidPolicyError("invalid type {0} for attribute '{1}', please change to object."
-                                         .format(type(individual_policy).__name__, extension_name))
-            _PolicyEngine.__validate_schema(individual_policy, individual_policy_schema, extension_name)
+        return_value = {}
 
-            extension_signature_policy = individual_policy.get("signatureRequired")
-            if extension_signature_policy is None:
-                extension_signature_policy = template["extensionPolicies"]["signatureRequired"]
+        signature_required = _PolicyEngine._get_boolean(extension, attribute="signatureRequired", name_prefix=extension_attribute_name, optional=True, default=None)
+        if signature_required is not None:
+            return_value["signatureRequired"] = signature_required
 
-            parsed_extensions_dict[extension_name] = \
-                {
-                    "signatureRequired": extension_signature_policy
-                }
+        # The runtimePolicy is an arbitrary object.
+        runtime_policy = extension.get("runtimePolicy")
+        if runtime_policy is not None:
+            return_value["runtimePolicy"] = runtime_policy
 
-        # Convert "extensions" to a case-folded dict for case-insensitive lookup
-        case_folded_extensions_dict = _CaseFoldedDict.from_dict(parsed_extensions_dict)
-        template["extensionPolicies"]["extensions"] = case_folded_extensions_dict
+        return return_value
 
     @staticmethod
-    def __validate_schema(policy, schema, section_name=None):
+    def _check_attributes(object_, object_name, valid_attributes):
         """
-        Validate the provided policy against the schema - we only do a shallow check (no recursion into nested dicts).
-        If there is an unrecognized attribute, raise an error.
+        Check that the given object, which should be a dictionary, has only the specified attributes.
+        If any other attributes are present, raise InvalidPolicyError.
+        The object_name is used in the error message.
         """
-        for key, value in policy.items():
-            if key not in schema:
-                raise InvalidPolicyError("attribute '{0}' is not defined in the policy schema. Please refer to the policy documentation "
-                                         "and change or remove this attribute accordingly.".format(key))
-
-            expected_type = schema.get(key)
-            if isinstance(expected_type, dict):
-                expected_type = dict
-            type_in_err_msg = {
-                dict: "object",
-                ustr: "string",
-                bool: "boolean"
-            }
-
-            if not isinstance(value, expected_type):
-                if section_name is None:
-                    msg = ("invalid type {0} for attribute '{1}', please change to {2}."
-                           .format(type(value).__name__, key, type_in_err_msg.get(expected_type)))
-                else:
-                    msg = ("invalid type {0} for attribute '{1}' in section '{2}', please change to {3}."
-                           .format(type(value).__name__, key, section_name,
-                                   type_in_err_msg.get(expected_type)))
-
-                raise InvalidPolicyError(msg)
+        for k in object_.keys():
+            if k not in valid_attributes:
+                raise InvalidPolicyError("unrecognized attribute '{0}' in {1}".format(k, object_name))
+
+    @staticmethod
+    def _get_dictionary(object_, attribute, name_prefix="", optional=False, default=None):
+        """
+        Returns object[attribute] if it exists, verifying that it is a dictionary.
+        If object_[attribute] does not exist and 'optional' is True, returns 'default'; if 'optional' is False raises InvalidPolicyError.
+        If object_[attribute] is not a dictionary, raises InvalidPolicyError.
+        The name_prefix indicates the path of the attribute within the policy document and is used in the error message.
+        """
+        return _PolicyEngine._get_value(object_, attribute, name_prefix, dict, "object", optional=optional, default=default)
+
+    @staticmethod
+    def _get_string(object_, attribute, name_prefix="", optional=False, default=None):
+        """
+        Returns object[attribute] if it exists, verifying that it is a string, else returns default.
+        If object_[attribute] does not exist and 'optional' is True, returns 'default'; if 'optional' is False raises InvalidPolicyError.
+        If object_[attribute] is not a string, raises InvalidPolicyError.
+        The name_prefix indicates the path of the attribute within the policy document and is used in the error message.
+        """
+        return _PolicyEngine._get_value(object_, attribute, name_prefix, ustr, "string", optional=optional, default=default)
+
+    @staticmethod
+    def _get_boolean(object_, attribute, name_prefix="", optional=False, default=None):
+        """
+        Returns object[attribute] if it exists, verifying that it is a boolean, else returns default.
+        If object_[attribute] does not exist and 'optional' is True, returns 'default'; if 'optional' is False raises InvalidPolicyError.
+        If object_[attribute] is not a boolean, raises InvalidPolicyError.
+        The name_prefix indicates the path of the attribute within the policy document and is used in the error message.
+        """
+        return _PolicyEngine._get_value(object_, attribute, name_prefix, bool, "boolean", optional=optional, default=default)
+
+    @staticmethod
+    def _get_value(object_, attribute, name_prefix, type_, type_name, optional, default):
+        """
+        Returns object[attribute] if it exists, verifying that it is of the given type_, else returns default.
+        If object_[attribute] does not exist and 'optional' is True, returns 'default'; if 'optional' is False raises InvalidPolicyError.
+        If the type of object_[attribute] is not 'type_', raises InvalidPolicyError.
+        The name_prefix indicates the path of the attribute within the policy document, the type_name indicates a user-friendly name for type_; both are used in the error message.
+        """
+        if default is not None and not optional:
+            raise ValueError("default value should only be provided for optional attributes")
+        value = object_.get(attribute)
+        if value is None:
+            if not optional:
+                raise InvalidPolicyError("missing required attribute '{0}{1}'".format(name_prefix, attribute))
+            return default
+        if not isinstance(value, type_):
+            raise InvalidPolicyError("invalid type {0} for attribute '{1}{2}'; must be '{3}'".format(type(value).__name__, name_prefix, attribute, type_name))
+        return value
 
 
 class ExtensionPolicyEngine(_PolicyEngine):
 
-    def should_allow_extension(self, extension_to_check):
+    def should_allow_extension(self, extension_name):
         """
         Return whether we should allow extension download based on policy.
-        extension_to_check is expected to be an Extension object.
-
         If policy feature not enabled, return True.
         If allowListedExtensionsOnly=true, return true only if extension present in "extensions" allowlist.
         If allowListedExtensions=false, return true always.
@@ -294,24 +290,20 @@ class ExtensionPolicyEngine(_PolicyEngine):
         allow_listed_extension_only = self._policy.get("extensionPolicies").get("allowListedExtensionsOnly")
         extension_allowlist = self._policy.get("extensionPolicies").get("extensions")
 
-        should_allow = not allow_listed_extension_only or extension_allowlist.get(extension_to_check.name) is not None
+        should_allow = not allow_listed_extension_only or extension_allowlist.get(extension_name) is not None
         return should_allow
 
-    def should_enforce_signature_validation(self, extension_to_check):
+    def should_enforce_signature_validation(self, extension_name):
         """
         Return whether we should enforce signature based on policy.
-        extension_to_check is expected to be an Extension object.
-
         If policy feature not enabled, return False.
-        Individual policy takes precedence over global - if individual signing policy present, return true/false based on
-        individual policy. Else, return true/false based on global policy.
+        Individual policy takes precedence over global.
         """
         if not self.policy_enforcement_enabled:
             return False
 
         global_signature_required = self._policy.get("extensionPolicies").get("signatureRequired")
-        individual_policy = self._policy.get("extensionPolicies").get("extensions").get(extension_to_check.name)
-        if individual_policy is None:
-            return global_signature_required
-        else:
-            return individual_policy.get("signatureRequired")
+        individual_policy = self._policy.get("extensionPolicies").get("extensions").get(extension_name)
+        individual_signature_required = individual_policy.get("signatureRequired") if individual_policy is not None else None
+
+        return individual_signature_required if individual_signature_required is not None else global_signature_required


Test patch:
diff --git a/tests/ga/test_policy_engine.py b/tests/ga/test_policy_engine.py
index 3a7e9329..c55a5f5c 100644
--- a/tests/ga/test_policy_engine.py
+++ b/tests/ga/test_policy_engine.py
@@ -22,7 +22,6 @@ from tests.lib.tools import AgentTestCase
 from azurelinuxagent.ga.policy.policy_engine import ExtensionPolicyEngine, InvalidPolicyError, \
     _PolicyEngine, _DEFAULT_ALLOW_LISTED_EXTENSIONS_ONLY, _DEFAULT_SIGNATURE_REQUIRED
 from tests.lib.tools import patch
-from azurelinuxagent.common.protocol.restapi import Extension
 
 TEST_EXTENSION_NAME = "Microsoft.Azure.ActiveDirectory.AADSSHLoginForLinux"
 
@@ -77,7 +76,10 @@ class TestPolicyEngine(_TestPolicyBase):
         When conf flag is set to true and policy file is present at expected location, feature should be enabled.
         """
         # Create policy file with empty policy object at the expected path to enable feature.
-        self._create_policy_file({})
+        self._create_policy_file(
+        {
+            "policyVersion": "0.0.1"
+        })
         engine = _PolicyEngine()
         self.assertTrue(engine.policy_enforcement_enabled,
                         msg="Conf flag is set to true so policy enforcement should be enabled.")
@@ -115,7 +117,8 @@ class TestPolicyEngine(_TestPolicyBase):
                     "signatureRequired": True,
                     "extensions": {
                         TEST_EXTENSION_NAME: {
-                            "signatureRequired": False
+                            "signatureRequired": False,
+                            "runtimePolicy": True
                         }
                     }
                 }
@@ -128,7 +131,8 @@ class TestPolicyEngine(_TestPolicyBase):
                     "signatureRequired": False,
                     "extensions": {
                         TEST_EXTENSION_NAME: {
-                            "signatureRequired": True
+                            "signatureRequired": True,
+                            "runtimePolicy": [ True, None, { "bar": "baz" } ]
                         }
                     }
                 }
@@ -147,6 +151,21 @@ class TestPolicyEngine(_TestPolicyBase):
             actual_individual_policy = actual_extension_policy.get("extensions").get(TEST_EXTENSION_NAME)
             expected_individual_policy = expected_extension_policy.get("extensions").get(TEST_EXTENSION_NAME)
             self.assertEqual(actual_individual_policy.get("signatureRequired"), expected_individual_policy.get("signatureRequired"))
+            self.assertEqual(actual_individual_policy.get("runtimePolicy"), expected_individual_policy.get("runtimePolicy"))
+
+    def test_it_should_verify_policy_version_is_required(self):
+        self._create_policy_file({
+                "extensionPolicies": {}
+            })
+        with self.assertRaises(InvalidPolicyError):
+            _PolicyEngine()
+
+    def test_it_should_accept_partially_specified_policy_versions(self):
+        for policy_version in ['0', '0.1', '0.1.0']:
+            self._create_policy_file({
+                    "policyVersion": policy_version,
+                })
+            self.assertEqual(policy_version, _PolicyEngine()._policy["policyVersion"])
 
     def test_should_raise_error_if_policy_file_is_invalid_json(self):
         cases = [
@@ -295,11 +314,10 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         When custom policy file not present, should allow all extensions and not enforce signature.
         """
         # No policy file is present - feature is disabled.
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         engine = ExtensionPolicyEngine()
-        should_allow = engine.should_allow_extension(test_extension)
+        should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME)
         self.assertTrue(should_allow, msg="Policy feature is disabled because no policy file present, so all extensions should be allowed.")
-        should_enforce = engine.should_enforce_signature_validation(test_extension)
+        should_enforce = engine.should_enforce_signature_validation(TEST_EXTENSION_NAME)
         self.assertFalse(should_enforce, msg="Policy feature is disabled because no policy file present, so signature should not be enforced.")
 
     def test_should_allow_and_should_not_enforce_signature_if_conf_flag_false(self):
@@ -308,20 +326,17 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         """
         self.patch_conf_flag.stop()
         self._create_policy_file({})
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         engine = ExtensionPolicyEngine()
-        should_allow = engine.should_allow_extension(test_extension)
+        should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME)
         self.assertTrue(should_allow, msg="Policy feature is disabled because conf flag false, so all extensions should be allowed.")
-        should_enforce = engine.should_enforce_signature_validation(test_extension)
+        should_enforce = engine.should_enforce_signature_validation(TEST_EXTENSION_NAME)
         self.assertFalse(should_enforce, msg="Policy feature is disabled because conf flag false, so signature should not be enforced.")
 
     def test_should_use_default_policy_if_no_extension_policy_specified(self):
         """
         Test that default policy is used when policy file does not specify the extension policy.
         """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         policy_cases = [
-            {},
             {
                 "policyVersion": "0.1.0"
             },
@@ -333,10 +348,10 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         for policy in policy_cases:
             self._create_policy_file(policy)
             engine = ExtensionPolicyEngine()
-            should_allow = engine.should_allow_extension(test_extension)
+            should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME)
             self.assertEqual(should_allow, not _DEFAULT_ALLOW_LISTED_EXTENSIONS_ONLY,
                              msg="Extension policy is not specified, so should use default policy.")
-            should_enforce = engine.should_enforce_signature_validation(test_extension)
+            should_enforce = engine.should_enforce_signature_validation(TEST_EXTENSION_NAME)
             self.assertEqual(should_enforce, _DEFAULT_SIGNATURE_REQUIRED,
                              msg="Extension policy is not specified, so should use default policy.")
 
@@ -344,9 +359,7 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         """
         If allowListedExtensionsOnly is true and extension in list, should_allow = True.
         """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         TEST_EXTENSION_NAME_2 = "Test.Extension.Name"
-        test_extension_2 = Extension(name=TEST_EXTENSION_NAME_2)
         policy = \
             {
                 "policyVersion": "0.1.0",
@@ -363,16 +376,15 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
             }
         self._create_policy_file(policy)
         engine = ExtensionPolicyEngine()
-        should_allow = engine.should_allow_extension(test_extension)
+        should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME)
         self.assertTrue(should_allow, msg="Extension is in allowlist, so should be allowed.")
-        should_allow = engine.should_allow_extension(test_extension_2)
+        should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME_2)
         self.assertTrue(should_allow, msg="Extension is in allowlist, so should be allowed.")
 
     def test_should_not_allow_if_allowListedExtensionsOnly_true_and_extension_not_in_list(self):
         """
         If allowListedExtensionsOnly is true and extension not in list, should_allow = False.
         """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         policy = \
             {
                 "policyVersion": "0.1.0",
@@ -384,7 +396,7 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
             }
         self._create_policy_file(policy)
         engine = ExtensionPolicyEngine()
-        should_allow = engine.should_allow_extension(test_extension)
+        should_allow = engine.should_allow_extension(TEST_EXTENSION_NAME)
         self.assertFalse(should_allow,
                             msg="allowListedExtensionsOnly is true and extension is not in allowlist, so should not be allowed.")
 
@@ -393,8 +405,6 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         If allowListedExtensionsOnly is false, should_allow = True (whether extension in list or not).
         """
         # Test an extension in the allowlist, and an extension not in the allowlist. Both should be allowed.
-        test_ext_in_list = Extension(name=TEST_EXTENSION_NAME)
-        test_ext_not_in_list = Extension(name="Random.Ext")
         policy = \
             {
                 "policyVersion": "0.1.0",
@@ -410,16 +420,15 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
             }
         self._create_policy_file(policy)
         engine = ExtensionPolicyEngine()
-        self.assertTrue(engine.should_allow_extension(test_ext_in_list),
+        self.assertTrue(engine.should_allow_extension(TEST_EXTENSION_NAME),
                         msg="allowListedExtensionsOnly is false, so extension should be allowed.")
-        self.assertTrue(engine.should_allow_extension(test_ext_not_in_list),
+        self.assertTrue(engine.should_allow_extension("Random.Ext"),
                         msg="allowListedExtensionsOnly is false, so extension should be allowed.")
 
     def test_should_enforce_signature_if_individual_signatureRequired_true(self):
         """
         If signatureRequired is true for individual extension, should_enforce_signature_validation = True (whether global signatureRequired is true or false).
         """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         for global_rule in [True, False]:
             policy = \
                 {
@@ -436,7 +445,7 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
                 }
             self._create_policy_file(policy)
             engine = ExtensionPolicyEngine()
-            should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
+            should_enforce_signature = engine.should_enforce_signature_validation(TEST_EXTENSION_NAME)
             self.assertTrue(should_enforce_signature,
                             msg="Individual signatureRequired policy is true, so signature should be enforced.")
 
@@ -444,7 +453,6 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
         """
         If signatureRequired is false for individual extension policy, should_enforce_signature_validation = False (whether global signatureRequired is true or false).
         """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
         for global_rule in [True, False]:
             policy = \
                 {
@@ -461,93 +469,41 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
                 }
             self._create_policy_file(policy)
             engine = ExtensionPolicyEngine()
-            should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
+            should_enforce_signature = engine.should_enforce_signature_validation(TEST_EXTENSION_NAME)
             self.assertFalse(should_enforce_signature,
                                 msg="Individual signatureRequired policy is false, so signature should be not enforced.")
 
-    def test_should_enforce_signature_if_global_signatureRequired_true_and_no_individual_policy(self):
-        """
-        If signatureRequired is true globally and no individual extension signature policy, should_enforce_signature_validation = True.
-        """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
-        policy = \
-            {
-                "policyVersion": "0.1.0",
-                "extensionPolicies": {
-                    "allowListedExtensionsOnly": True,
-                    "signatureRequired": True,
-                    "extensions": {}
-                }
-            }
-        self._create_policy_file(policy)
-        engine = ExtensionPolicyEngine()
-        should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
-        self.assertTrue(should_enforce_signature,
-                        msg="Global signatureRequired policy is true, so signature should be enforced.")
-
-    def test_should_not_enforce_signature_if_global_signatureRequired_false_and_no_individual_policy(self):
-        """
-        If signatureRequired is false globally and no individual extension signature policy, should_enforce_signature_validation = False.
-        """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
-        policy = \
-            {
-                "policyVersion": "0.1.0",
-                "extensionPolicies": {
-                    "allowListedExtensionsOnly": True,
-                    "signatureRequired": False,
-                    "extensions": {}
-                }
-            }
-        self._create_policy_file(policy)
-        engine = ExtensionPolicyEngine()
-        should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
-        self.assertFalse(should_enforce_signature,
-                            msg="Global signatureRequired policy is false, so signature should not be enforced.")
-
-    def test_should_enforce_signature_if_global_signatureRequired_true_and_individual_signatureRequired_not_specified(self):
-        """
-        If individual policy is present, but signatureRequired is not specified for that policy, use global signatureRequired.
-        """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
-        policy = \
-            {
-                "policyVersion": "0.1.0",
-                "extensionPolicies": {
-                    "allowListedExtensionsOnly": True,
-                    "signatureRequired": True,
-                    "extensions": {
-                        TEST_EXTENSION_NAME: {}
+    def test_should_use_global_signatureRequired_when_an_individual_policy_is_not_specified(self):
+        for global_policy in [True, False]:
+            extensions_test_cases = [
+                None,
+                {},
+                {
+                    TEST_EXTENSION_NAME: {}
+                },
+                {
+                    TEST_EXTENSION_NAME: {
+                        "runtimePolicy": "an arbitrary object"
                     }
                 }
-            }
-        self._create_policy_file(policy)
-        engine = ExtensionPolicyEngine()
-        should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
-        self.assertTrue(should_enforce_signature,
-                            msg="Individual signatureRequired policy is not set, so should use global policy and enforce signature.")
-
-    def test_should_not_enforce_signature_if_global_signatureRequired_false_and_individual_signatureRequired_not_specified(self):
-        """
-        If individual policy is present, but signatureRequired is not specified for that policy, use global signatureRequired.
-        """
-        test_extension = Extension(name=TEST_EXTENSION_NAME)
-        policy = \
-            {
-                "policyVersion": "0.1.0",
-                "extensionPolicies": {
-                    "allowListedExtensionsOnly": True,
-                    "signatureRequired": False,
-                    "extensions": {
-                        TEST_EXTENSION_NAME: {}
+            ]
+            for extensions in extensions_test_cases:
+                policy = {
+                    "policyVersion": "0.1.0",
+                    "extensionPolicies": {
+                        "allowListedExtensionsOnly": True,
+                        "signatureRequired": global_policy,
                     }
                 }
-            }
-        self._create_policy_file(policy)
-        engine = ExtensionPolicyEngine()
-        should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
-        self.assertFalse(should_enforce_signature,
-                            msg="Individual signatureRequired policy is not set, so should use global policy and enforce signature.")
+                if extensions is not None:
+                    policy["extensionPolicies"]["extensions"] = extensions
+
+                self._create_policy_file(policy)
+
+                self.assertEqual(
+                    global_policy,
+                    ExtensionPolicyEngine().should_enforce_signature_validation(TEST_EXTENSION_NAME),
+                    "The global signatureRequired ({0}) should have been used. Policy:\n{1}".format(global_policy, policy))
 
     def test_extension_name_in_policy_should_be_case_insensitive(self):
         """
@@ -559,7 +515,6 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
             "MicrOsoft.aZure.activedirectory.aaDsShloginFORlinux",
             "microsoft.azure.activedirectory.aadsshloginforlinux"
         ]:
-            test_extension = Extension(name=ext_name_to_test)
             policy = \
                 {
                     "policyVersion": "0.1.0",
@@ -576,9 +531,9 @@ class TestExtensionPolicyEngine(_TestPolicyBase):
 
             self._create_policy_file(policy)
             engine = ExtensionPolicyEngine()
-            should_allow = engine.should_allow_extension(test_extension)
-            should_enforce_signature = engine.should_enforce_signature_validation(test_extension)
+            should_allow = engine.should_allow_extension(ext_name_to_test)
+            should_enforce_signature = engine.should_enforce_signature_validation(ext_name_to_test)
             self.assertTrue(should_allow,
                             msg="Extension should have been found in allowlist regardless of extension name case.")
             self.assertTrue(should_enforce_signature,
-                            msg="Individual signatureRequired policy should have been found and used, regardless of extension name case.")
\ No newline at end of file
+                            msg="Individual signatureRequired policy should have been found and used, regardless of extension name case.")