diff --git a/src/kata-opa/allow-all-except-exec-process.rego b/src/kata-opa/allow-all-except-exec-process.rego
index ec3bf15a9e..7905f4e592 100644
--- a/src/kata-opa/allow-all-except-exec-process.rego
+++ b/src/kata-opa/allow-all-except-exec-process.rego
@@ -1,3 +1,8 @@
+# Copyright (c) 2023 Microsoft Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
 package agent_policy
 
 default AddARPNeighborsRequest := true
diff --git a/src/kata-opa/allow-all.rego b/src/kata-opa/allow-all.rego
index 7ac8134f03..7714e3e917 100644
--- a/src/kata-opa/allow-all.rego
+++ b/src/kata-opa/allow-all.rego
@@ -1,3 +1,8 @@
+# Copyright (c) 2023 Microsoft Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
 package agent_policy
 
 default AddARPNeighborsRequest := true
diff --git a/src/kata-opa/allow-set-policy.rego b/src/kata-opa/allow-set-policy.rego
new file mode 100644
index 0000000000..c253c60d80
--- /dev/null
+++ b/src/kata-opa/allow-set-policy.rego
@@ -0,0 +1,8 @@
+# Copyright (c) 2024 Microsoft Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+package agent_policy
+
+default SetPolicyRequest := true
diff --git a/tests/integration/kubernetes/setup.sh b/tests/integration/kubernetes/setup.sh
index eb13079696..1f829d1ef0 100644
--- a/tests/integration/kubernetes/setup.sh
+++ b/tests/integration/kubernetes/setup.sh
@@ -29,6 +29,7 @@ else
 	K8S_TEST_POLICY_FILES=( \
 		"allow-all.rego" \
 		"allow-all-except-exec-process.rego" \
+		"allow-set-policy.rego" \
     )
 fi
 
diff --git a/tools/osbuilder/rootfs-builder/rootfs.sh b/tools/osbuilder/rootfs-builder/rootfs.sh
index 0defcf8ec1..a38811dfe6 100755
--- a/tools/osbuilder/rootfs-builder/rootfs.sh
+++ b/tools/osbuilder/rootfs-builder/rootfs.sh
@@ -30,6 +30,7 @@ LIBC=${LIBC:-musl}
 SECCOMP=${SECCOMP:-"yes"}
 SELINUX=${SELINUX:-"no"}
 AGENT_POLICY=${AGENT_POLICY:-no}
+AGENT_POLICY_FILE=${AGENT_POLICY_FILE:-"${script_dir}/../../../src/kata-opa/allow-all.rego"}
 AGENT_SOURCE_BIN=${AGENT_SOURCE_BIN:-""}
 AGENT_TARBALL=${AGENT_TARBALL:-""}
 COCO_GUEST_COMPONENTS_TARBALL=${COCO_GUEST_COMPONENTS_TARBALL:-""}
@@ -39,6 +40,8 @@ PAUSE_IMAGE_TARBALL=${PAUSE_IMAGE_TARBALL:-""}
 lib_file="${script_dir}/../scripts/lib.sh"
 source "$lib_file"
 
+agent_policy_file="$(readlink -f "${AGENT_POLICY_FILE}")"
+
 #For cross build
 CROSS_BUILD=${CROSS_BUILD:-false}
 BUILDX=""
@@ -130,6 +133,11 @@ AGENT_INIT          When set to "yes", use ${AGENT_BIN} as init process in place
                     of systemd.
                     Default value: no
 
+AGENT_POLICY_FILE   Path to the agent policy rego file to be set in the rootfs.
+                    If defined, this overwrites the default setting of the
+                    permissive policy file.
+                    Default value: allow-all.rego
+
 AGENT_SOURCE_BIN    Path to the directory of agent binary.
                     If set, use the binary as agent but not build agent package.
                     AGENT_SOURCE_BIN and AGENT_TARBALL should never be used toghether.
@@ -355,6 +363,8 @@ check_env_variables()
 
 	[ -n "${KERNEL_MODULES_DIR}" ] && [ ! -d "${KERNEL_MODULES_DIR}" ] && die "KERNEL_MODULES_DIR defined but is not an existing directory"
 
+	[ ! -f "${agent_policy_file}" ] && die "agent policy file not found in '${agent_policy_file}'"
+
 	[ -n "${OSBUILDER_VERSION}" ] || die "need osbuilder version"
 }
 
@@ -494,6 +504,7 @@ build_rootfs_distro()
 			--env ROOTFS_DIR="/rootfs" \
 			--env AGENT_BIN="${AGENT_BIN}" \
 			--env AGENT_INIT="${AGENT_INIT}" \
+			--env AGENT_POLICY_FILE="${AGENT_POLICY_FILE}" \
 			--env ARCH="${ARCH}" \
 			--env MEASURED_ROOTFS="${MEASURED_ROOTFS}" \
 			--env KERNEL_MODULES_DIR="${KERNEL_MODULES_DIR}" \
@@ -701,13 +712,12 @@ EOF
 	if [ "${AGENT_POLICY}" == "yes" ]; then
 		info "Install the default policy"
 		# Install default settings for the kata-opa service.
-		local kata_opa_in_dir="${script_dir}/../../../src/kata-opa"
 		local opa_settings_dir="/etc/kata-opa"
-		local policy_file="allow-all.rego"
+		local policy_file_name="$(basename ${agent_policy_file})"
 		local policy_dir="${ROOTFS_DIR}/${opa_settings_dir}"
 		mkdir -p "${policy_dir}"
-		install -D -o root -g root -m 0644 "${kata_opa_in_dir}/${policy_file}" -T "${policy_dir}/${policy_file}"
-		ln -sf "${policy_file}" "${policy_dir}/default-policy.rego"
+		install -D -o root -g root -m 0644 "${agent_policy_file}" -T "${policy_dir}/${policy_file_name}"
+		ln -sf "${policy_file_name}" "${policy_dir}/default-policy.rego"
 	fi
 
 	info "Check init is installed"