Repository: Azure/WALinuxAgent
Pull request #3234 (2024-09-25 13:30:46)

Replace Regorus with Python-based policy engine
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

Add Python-based policy handing logic to replace Regorus.
Extension handling code is not changed in this PR, so agent functionality is not affected. No e2e tests were added.

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

narrieta: Did not go thru the implementation details, but approach looks good to me

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3234?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
Attention: Patch coverage is `94.89796%` with `5 lines` in your changes missing coverage. Please review.
> Project coverage is 72.40%. Comparing base [(`3aebcdd`)](https://app.codecov.io/gh/Azure/WALinuxAgent/commit/3aebcdd781549b55356d0682d135eec6d53b7831?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) to head [(`e038560`)](https://app.codecov.io/gh/Azure/WALinuxAgent/commit/e03856012ca4961327a5da46e83a1b374954dfa2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).
> Report is 314 commits behind head on develop.

| [Files with missing lines](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3234?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Patch % | Lines |
|---|---|---|
| [azurelinuxagent/ga/policy/policy\_engine.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3234?src=pr&el=tree&filepath=azurelinuxagent%2Fga%2Fpolicy%2Fpolicy_engine.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL3BvbGljeS9wb2xpY3lfZW5naW5lLnB5) | 94.68% | [3 Missing and 2 partials :warning: ](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3234?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop    #3234      +/-   ##
===========================================
+ Coverage    71.97%   72.40%   +0.42%     
===========================================
  Files          103      113      +10     
  Lines        15692    16956    +1264     
  Branches      2486     2452      -34     
===========================================
+ Hits         11295    12277     +982     
- Misses        3881     4121     +240     
- Partials       516      558      +42     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3234?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).


Patch:
diff --git a/azurelinuxagent/common/conf.py b/azurelinuxagent/common/conf.py
index cd5c9a2d..cf38286a 100644
--- a/azurelinuxagent/common/conf.py
+++ b/azurelinuxagent/common/conf.py
@@ -170,7 +170,8 @@ __STRING_OPTIONS__ = {
     "ResourceDisk.MountPoint": "/mnt/resource",
     "ResourceDisk.MountOptions": None,
     "ResourceDisk.Filesystem": "ext3",
-    "AutoUpdate.GAFamily": "Prod"
+    "AutoUpdate.GAFamily": "Prod",
+    "Policy.PolicyFilePath": "/etc/waagent_policy.json"
 }
 
 
@@ -315,6 +316,10 @@ def get_agent_log_file():
     return "/var/log/waagent.log"
 
 
+def get_policy_file_path(conf=__conf__):
+    return conf.get("Policy.PolicyFilePath", "/etc/waagent_policy.json")
+
+
 def get_fips_enabled(conf=__conf__):
     return conf.get_switch("OS.EnableFIPS", False)
 
diff --git a/azurelinuxagent/ga/policy/policy_engine.py b/azurelinuxagent/ga/policy/policy_engine.py
index 20350dac..8d93e132 100644
--- a/azurelinuxagent/ga/policy/policy_engine.py
+++ b/azurelinuxagent/ga/policy/policy_engine.py
@@ -15,10 +15,40 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+import copy
+import json
+import os
+from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.event import WALAEventOperation, add_event
 from azurelinuxagent.common import conf
 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
+
+# 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.
+_MAX_SUPPORTED_POLICY_VERSION = "0.1.0"
 
 
 class PolicyError(AgentError):
@@ -27,12 +57,43 @@ class PolicyError(AgentError):
     """
 
 
-class PolicyEngine(object):
+class InvalidPolicyError(AgentError):
+    """
+    Error raised if user-provided policy is invalid.
+    """
+    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.
     """
-    @classmethod
-    def _log_policy(cls, msg, is_success=True, op=WALAEventOperation.Policy, send_event=True):
+    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
+
+    @staticmethod
+    def _log_policy_event(msg, is_success=True, op=WALAEventOperation.Policy, send_event=True):
         """
         Log information to console and telemetry.
         """
@@ -41,15 +102,216 @@ class PolicyEngine(object):
         else:
             logger.error(msg)
         if send_event:
-            add_event(op=op, message=msg, is_success=is_success)
+            add_event(op=op, message=msg, is_success=is_success, log_event=False)
 
     @staticmethod
-    def is_policy_enforcement_enabled():
+    def __get_policy_enforcement_enabled():
+        """
+        Policy will be enabled if (1) policy file exists at the expected location and (2) the conf flag "Debug.EnableExtensionPolicy" is true.
         """
-        Check whether user has opted into policy enforcement feature.
-        Caller function should check this before performing any operations.
+        return conf.get_extension_policy_enabled() and os.path.exists(conf.get_policy_file_path())
+
+    @property
+    def policy_enforcement_enabled(self):
+        return self._policy_enforcement_enabled
+
+    @staticmethod
+    def __read_policy():
         """
-        # TODO: Add check for policy file present at /etc/waagent_policy.json.
-        # Policy should only be enabled if conf flag is true AND policy file is present.
-        return conf.get_extension_policy_enabled()
+        Read customer-provided policy JSON file, load and return as a dict.
+        Policy file is expected to be at conf.get_policy_file_path(). Note that this method should only be called
+        after verifying that the file exists (currently done in __init__).
+
+        Raise InvalidPolicyError if JSON is invalid, or any exceptions are thrown while reading the file.
+        """
+        with open(conf.get_policy_file_path(), 'r') as f:
+            try:
+                contents = f.read()
+                _PolicyEngine._log_policy_event(
+                    "Policy enforcement is enabled. 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)
+            except ValueError as ex:
+                msg = "policy file does not conform to valid json syntax"
+                raise InvalidPolicyError(msg=msg, inner=ex)
+            except Exception as ex:
+                msg = "unable to read policy file"
+                raise InvalidPolicyError(msg=msg, inner=ex)
+
+            return custom_policy
+
+    @staticmethod
+    def __parse_policy(template, custom_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
+        """
+        # 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)
+
+    @staticmethod
+    def __parse_version(template, policy):
+        """
+        Validate and return "policyVersion" attribute. If not a string in the format "x.y.z", raise InvalidPolicyError.
+        If policy_version is greater than maximum supported version, raise InvalidPolicyError.
+        """
+        version = policy.get("policyVersion")
+        if version is None:
+            return
+
+        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 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))
+
+        template["policyVersion"] = 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
 
+            # Parse global signature policy
+            signature_policy = extension_policies.get("signatureRequired")
+            if signature_policy is not None:
+                template["extensionPolicies"]["signatureRequired"] = signature_policy
+
+            # Parse individual extension policies
+            _PolicyEngine.__parse_extensions(template, extension_policies)
+
+    @staticmethod
+    def __parse_extensions(template, extensions_policy):
+        """
+        Parse "extensions" dict and update in template.
+        "extensions" is expected to be in the format:
+        {
+            "extensions": {
+                "<extensionName>": {
+                    "signatureRequired": bool
+                }
+            }
+        }
+
+        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.
+        """
+        extensions = extensions_policy.get("extensions")
+        if extensions is None:
+            return
+
+        parsed_extensions_dict = {}
+        for extension_name, individual_policy in extensions.items():
+
+            # 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)
+
+            extension_signature_policy = individual_policy.get("signatureRequired")
+            if extension_signature_policy is None:
+                extension_signature_policy = template["extensionPolicies"]["signatureRequired"]
+
+            parsed_extensions_dict[extension_name] = \
+                {
+                    "signatureRequired": extension_signature_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
+
+    @staticmethod
+    def __validate_schema(policy, schema, section_name=None):
+        """
+        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.
+        """
+        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)
+
+
+class ExtensionPolicyEngine(_PolicyEngine):
+
+    def should_allow_extension(self, extension_to_check):
+        """
+        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.
+        """
+        if not self.policy_enforcement_enabled:
+            return True
+
+        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
+        return should_allow
+
+    def should_enforce_signature_validation(self, extension_to_check):
+        """
+        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.
+        """
+        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")


Test patch:
diff --git a/tests/data/test_waagent.conf b/tests/data/test_waagent.conf
index 8fb05155..8abbd55e 100644
--- a/tests/data/test_waagent.conf
+++ b/tests/data/test_waagent.conf
@@ -140,3 +140,4 @@ OS.SshDir=/notareal/path
 # - The default is false to protect the state of existing VMs
 OS.EnableFirewall=n
 
+Debug.EnableExtensionPolicy=y
diff --git a/tests/ga/test_policy_engine.py b/tests/ga/test_policy_engine.py
index 7f9a8bcc..3a7e9329 100644
--- a/tests/ga/test_policy_engine.py
+++ b/tests/ga/test_policy_engine.py
@@ -15,20 +15,570 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+import os
+import json
+
 from tests.lib.tools import AgentTestCase
-from azurelinuxagent.ga.policy.policy_engine import PolicyEngine
+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"
+
+
+class _TestPolicyBase(AgentTestCase):
+    """
+    Define common methods for policy engine test classes.
+    """
+    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()
+
+    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 _run_test_cases_should_fail_to_parse(self, cases, assert_msg):
+        """
+        Cases should be a list of policies.
+        For each policy in the list, we create a policy file, initialize policy engine, and assert that InvalidPolicyError
+        is raised.
+        """
+        for policy in cases:
+            self._create_policy_file(policy)
+            msg = "invalid policy should not have parsed successfully: {0}.\nPolicy: \n{1}".format(assert_msg, policy)
+            with self.assertRaises(InvalidPolicyError, msg=msg):
+                _PolicyEngine()
+
+
+class TestPolicyEngine(_TestPolicyBase):
+    """
+    Test policy enablement and parsing logic for _PolicyEngine.
+    """
+    def test_policy_enforcement_should_be_enabled_when_policy_file_exists_and_conf_flag_true(self):
+        """
+        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({})
+        engine = _PolicyEngine()
+        self.assertTrue(engine.policy_enforcement_enabled,
+                        msg="Conf flag is set to true so policy enforcement should be enabled.")
 
+    def test_policy_enforcement_should_be_disabled_when_conf_flag_false_or_no_policy_file(self):
 
-class TestPolicyEngine(AgentTestCase):
-    def test_policy_enforcement_should_be_enabled(self):
-        with patch('azurelinuxagent.ga.policy.policy_engine.conf.get_extension_policy_enabled', return_value=True):
-            engine = PolicyEngine()
-            self.assertTrue(engine.is_policy_enforcement_enabled(),
-                            msg="Conf flag is set to true so policy enforcement should be enabled.")
+        # Test when conf flag is turned off - feature should be disabled.
+        self.patch_conf_flag.stop()
+        engine1 = _PolicyEngine()
+        self.assertFalse(engine1.policy_enforcement_enabled,
+                         msg="Conf flag is set to false and policy file missing so policy enforcement should be disabled.")
 
-    def test_policy_enforcement_should_be_disabled(self):
-        engine = PolicyEngine()
-        self.assertFalse(engine.is_policy_enforcement_enabled(),
+        # Turn on conf flag - feature should still be disabled, because policy file is not present.
+        self.patch_conf_flag.start()
+        engine2 = _PolicyEngine()
+        self.assertFalse(engine2.policy_enforcement_enabled,
+                         msg="Policy file is not present so policy enforcement should be disabled.")
+
+        # Create a policy file, but turn off conf flag - feature should be disabled due to flag.
+        self.patch_conf_flag.stop()
+        self._create_policy_file({})
+        engine3 = _PolicyEngine()
+        self.assertFalse(engine3.policy_enforcement_enabled,
                          msg="Conf flag is set to false so policy enforcement should be disabled.")
 
+    def test_should_parse_policy_successfully(self):
+        """
+        Values provided in custom policy should override any defaults.
+        """
+        policy1 = \
+            {
+                "policyVersion": "0.0.1",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": True,
+                    "extensions": {
+                        TEST_EXTENSION_NAME: {
+                            "signatureRequired": False
+                        }
+                    }
+                }
+            }
+        policy2 = \
+            {
+                "policyVersion": "0.0.1",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {
+                        TEST_EXTENSION_NAME: {
+                            "signatureRequired": True
+                        }
+                    }
+                }
+            }
+        for expected_policy in [policy1, policy2]:
+            self._create_policy_file(expected_policy)
+            engine = _PolicyEngine()
+            actual_policy = engine._policy
+            self.assertEqual(actual_policy.get("policyVersion"), expected_policy.get("policyVersion"))
+
+            actual_extension_policy = actual_policy.get("extensionPolicies")
+            expected_extension_policy = expected_policy.get("extensionPolicies")
+            self.assertEqual(actual_extension_policy.get("allowListedExtensionsOnly"), expected_extension_policy.get("allowListedExtensionsOnly"))
+            self.assertEqual(actual_extension_policy.get("signatureRequired"), expected_extension_policy.get("signatureRequired"))
+
+            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"))
+
+    def test_should_raise_error_if_policy_file_is_invalid_json(self):
+        cases = [
+            '''
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+            ''',
+            "",
+            " ",
+            "policy",
+            '''
+            { not_a_string: ""}
+            '''
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "not a valid json")
+
+    def test_should_raise_error_for_invalid_policy_version(self):
+        cases = [
+            {"policyVersion": "1.2.a"},
+            {"policyVersion": 0},
+            {"policyVersion": None}
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "policy version invalid")
+
+    def test_should_raise_error_for_unsupported_policy_version(self):
+        cases = [
+            {"policyVersion": "9.9.9"},
+            {"policyVersion": "9"}
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "agent does not support policy version")
+
+    def test_should_raise_error_if_extensions_policy_is_not_dict(self):
+        cases = [
+            {
+                "extensionPolicies": ""
+            },
+            {
+                "extensionPolicies": None
+            }
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "extensionPolicies is not a dict")
+
+    def test_should_raise_error_if_allowListedExtensionsOnly_is_not_bool(self):
+        cases = [
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": "True",  # Should be bool
+                    "signatureRequired": False,
+                    "extensions": {}
+                }
+            }
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "allowListedExtensionsOnly is not a bool")
+
+    def test_should_raise_error_if_signatureRequired_is_not_bool(self):
+        cases = [
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {
+                        TEST_EXTENSION_NAME: {
+                            "signatureRequired": "False"  # Should be bool
+                        }
+                    }
+                }
+            },
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": "False",  # Should be bool
+                    "extensions": {}
+                }
+            }
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "signatureRequired is not a bool")
+
+    def test_should_raise_error_if_extensions_is_not_dict(self):
+        cases = [
+            {
+                "extensionPolicies": {
+                    "extensions": []
+                }
+            },
+            {
+                "extensionPolicies": {
+                    "extensions": 0
+                }
+            },
+            {
+                "extensionPolicies": {
+                    "extensions": None
+                }
+            }
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "'extensions' is not a dict")
+
+    def test_should_raise_error_if_individual_extension_policy_is_not_dict(self):
+        cases = [
+            {
+                "extensionPolicies": {
+                    "extensions": {
+                        "Ext.Name": 0
+                    }
+                }
+            },
+            {
+                "extensionPolicies": {
+                    "extensions": {
+                        "Ext.Name": []
+                    }
+                }
+            }
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "individual extension policy is not a dict")
+
+    def test_should_raise_error_for_unrecognized_attribute(self):
+        # All cases below have either a typo or a random additional attribute.
+        cases = [
+            {"policyVerion": "0.0.1"},
+            {"extentionPolicies": {}},
+            {"extensionPolicies": {
+                "signingRequired": {}
+            }},
+            {"extensionPolicies": {
+                "extensions": {
+                    TEST_EXTENSION_NAME: {
+                        "randomAttribute": ""
+                    }
+                }
+            }}
+        ]
+        self._run_test_cases_should_fail_to_parse(cases, "unrecognized attribute in policy")
+
+
+class TestExtensionPolicyEngine(_TestPolicyBase):
+    """
+    Test ExtensionPolicyEngine should_allow() and should_enforce_signature_validation().
+    """
+    def test_should_allow_and_should_not_enforce_signature_if_no_custom_policy_file(self):
+        """
+        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)
+        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)
+        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):
+        """
+        When conf flag turned off, should allow all extensions and not enforce signature.
+        """
+        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)
+        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)
+        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"
+            },
+            {
+                "policyVersion": "0.1.0",
+                "extensionPolicies": {}
+            }
+        ]
+        for policy in policy_cases:
+            self._create_policy_file(policy)
+            engine = ExtensionPolicyEngine()
+            should_allow = engine.should_allow_extension(test_extension)
+            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)
+            self.assertEqual(should_enforce, _DEFAULT_SIGNATURE_REQUIRED,
+                             msg="Extension policy is not specified, so should use default policy.")
+
+    def test_should_allow_if_allowListedExtensionsOnly_true_and_extension_in_list(self):
+        """
+        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",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {
+                        TEST_EXTENSION_NAME: {},
+                        TEST_EXTENSION_NAME_2: {
+                            "signatureRequired": False
+                        }
+                    }
+                }
+            }
+        self._create_policy_file(policy)
+        engine = ExtensionPolicyEngine()
+        should_allow = engine.should_allow_extension(test_extension)
+        self.assertTrue(should_allow, msg="Extension is in allowlist, so should be allowed.")
+        should_allow = engine.should_allow_extension(test_extension_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",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": True,
+                    "signatureRequired": False,
+                    "extensions": {}  # Extension not in allowed list.
+                }
+            }
+        self._create_policy_file(policy)
+        engine = ExtensionPolicyEngine()
+        should_allow = engine.should_allow_extension(test_extension)
+        self.assertFalse(should_allow,
+                            msg="allowListedExtensionsOnly is true and extension is not in allowlist, so should not be allowed.")
+
+    def test_should_allow_if_allowListedExtensionsOnly_false(self):
+        """
+        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",
+                "extensionPolicies": {
+                    "allowListedExtensionsOnly": False,
+                    "signatureRequired": False,
+                    "extensions": {
+                        TEST_EXTENSION_NAME: {
+                            "signatureRequired": False
+                        }
+                    }
+                }
+            }
+        self._create_policy_file(policy)
+        engine = ExtensionPolicyEngine()
+        self.assertTrue(engine.should_allow_extension(test_ext_in_list),
+                        msg="allowListedExtensionsOnly is false, so extension should be allowed.")
+        self.assertTrue(engine.should_allow_extension(test_ext_not_in_list),
+                        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 = \
+                {
+                    "policyVersion": "0.1.0",
+                    "extensionPolicies": {
+                        "allowListedExtensionsOnly": False,
+                        "signatureRequired": global_rule,
+                        "extensions": {
+                            TEST_EXTENSION_NAME: {
+                                "signatureRequired": True
+                            }
+                        }
+                    }
+                }
+            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 true, so signature should be enforced.")
+
+    def test_should_not_enforce_signature_if_individual_signatureRequired_false(self):
+        """
+        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 = \
+                {
+                    "policyVersion": "0.1.0",
+                    "extensionPolicies": {
+                        "allowListedExtensionsOnly": False,
+                        "signatureRequired": global_rule,
+                        "extensions": {
+                            TEST_EXTENSION_NAME: {
+                                "signatureRequired": False,
+                            }
+                        }
+                    }
+                }
+            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 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: {}
+                    }
+                }
+            }
+        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: {}
+                    }
+                }
+            }
+        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.")
+
+    def test_extension_name_in_policy_should_be_case_insensitive(self):
+        """
+        Extension name is allowed to be any case. Test that should_allow() and should_enforce_signature_validation() return expected
+        results, even when the extension name does not match the case of the name specified in policy.
+        """
+        ext_name_in_policy = "Microsoft.Azure.ActiveDirectory.AADSSHLoginForLinux"
+        for ext_name_to_test in [
+            "MicrOsoft.aZure.activedirectory.aaDsShloginFORlinux",
+            "microsoft.azure.activedirectory.aadsshloginforlinux"
+        ]:
+            test_extension = Extension(name=ext_name_to_test)
+            policy = \
+                {
+                    "policyVersion": "0.1.0",
+                    "extensionPolicies": {
+                        "allowListedExtensionsOnly": True,
+                        "signatureRequired": False,
+                        "extensions": {
+                            ext_name_in_policy: {
+                                "signatureRequired": True
+                            }
+                        }
+                    }
+                }
+
+            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)
+            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
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 89dff1ce..abd8f661 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -47,7 +47,7 @@ Debug.CgroupDisableOnQuotaCheckFailure = True
 Debug.CgroupLogMetrics = False
 Debug.EnableAgentMemoryUsageCheck = False
 Debug.EnableCgroupV2ResourceLimiting = False
-Debug.EnableExtensionPolicy = False
+Debug.EnableExtensionPolicy = True
 Debug.EnableFastTrack = True
 Debug.EnableGAVersioning = True
 Debug.EtpCollectionPeriod = 300
@@ -86,6 +86,7 @@ OS.SshDir = /notareal/path
 OS.SudoersDir = /etc/sudoers.d
 OS.UpdateRdmaDriver = False
 Pid.File = /var/run/waagent.pid
+Policy.PolicyFilePath = /etc/waagent_policy.json
 Provisioning.Agent = auto
 Provisioning.AllowResetSysUser = False
 Provisioning.DecodeCustomData = False