adding workload identity support for vn2

src/confcom/azext_confcom/config.py

@@ -87,6 +87,8 @@
 VIRTUAL_NODE_YAML_METADATA = "metadata"
 VIRTUAL_NODE_YAML_NAME = "name"
 VIRTUAL_NODE_YAML_ANNOTATIONS = "annotations"
+VIRTUAL_NODE_YAML_LABELS = "labels"
+VIRTUAL_NODE_YAML_LABEL_WORKLOAD_IDENTITY = "azure.workload.identity/use"
 VIRTUAL_NODE_YAML_POLICY = "microsoft.containerinstance.virtualnode.ccepolicy"
 VIRTUAL_NODE_YAML_LIFECYCLE = "lifecycle"
 VIRTUAL_NODE_YAML_LIFECYCLE_POST_START = "postStart"
@@ -173,6 +175,8 @@
 MANAGED_IDENTITY_ENV_RULES = _config["managedIdentity"]["environmentVariables"]
 # VN2 environment variables
 VIRTUAL_NODE_ENV_RULES = _config["default_envs_virtual_node"]["environmentVariables"]
+# VN2 environment variables for workload identities
+VIRTUAL_NODE_ENV_RULES_WORKLOAD_IDENTITY = _config["workload_identity_virtual_node"]["environmentVariables"]
 # Enable container restart environment variable for all containers
 ENABLE_RESTART_ENV_RULE = _config["enableRestart"]["environmentVariables"]
 # default mounts image for customer containers

src/confcom/azext_confcom/data/internal_config.json

@@ -101,6 +101,34 @@
             }
         ]
     },
+    "workload_identity_virtual_node": {
+        "environmentVariables": [
+            {
+                "name": "AZURE_CLIENT_ID",
+                "value": ".+",
+                "strategy": "re2",
+                "required": false
+            },
+            {
+                "name": "AZURE_TENANT_ID",
+                "value": ".+",
+                "strategy": "re2",
+                "required": false
+            },
+            {
+                "name": "AZURE_FEDERATED_TOKEN_FILE",
+                "value": ".+",
+                "strategy": "re2",
+                "required": false
+            },
+            {
+                "name": "AZURE_AUTHORITY_HOST",
+                "value": ".+",
+                "strategy": "re2",
+                "required": false
+            }
+        ]
+    },
     "managedIdentity": {
         "environmentVariables": [
             {

src/confcom/azext_confcom/security_policy.py

@@ -908,6 +908,10 @@ def load_policy_from_virtual_node_yaml_str(
         # extract existing policy and fragments for diff mode
         metadata = case_insensitive_dict_get(yaml, "metadata")
         annotations = case_insensitive_dict_get(metadata, config.VIRTUAL_NODE_YAML_ANNOTATIONS)
+        labels = case_insensitive_dict_get(metadata, config.VIRTUAL_NODE_YAML_LABELS) or []
+        use_workload_identity = (
+            config.VIRTUAL_NODE_YAML_LABEL_WORKLOAD_IDENTITY in labels
+            and labels.get(config.VIRTUAL_NODE_YAML_LABEL_WORKLOAD_IDENTITY) == "true")
         existing_policy = case_insensitive_dict_get(annotations, config.VIRTUAL_NODE_YAML_POLICY)
         try:
             if existing_policy:
@@ -951,6 +955,8 @@ def load_policy_from_virtual_node_yaml_str(
                 secrets_data,
                 approve_wildcards=approve_wildcards
             )
+            if use_workload_identity:
+                envs += config.VIRTUAL_NODE_ENV_RULES_WORKLOAD_IDENTITY
 
             # command
             command = case_insensitive_dict_get(container, "command") or []

src/confcom/azext_confcom/tests/latest/README.md

@@ -116,6 +116,7 @@ test_invalid_many_input_types | Makes sure we're only getting input from one sou
 test_diff_wrong_input_type | Makes sure we're only doing the diff command if we're using a ARM Template as the input type
 test_parameters_without_template | Makes sure we error out if a parameter file is getting passed in without an ARM Template
 test_input_and_virtual_node | Error out if both input and virtual node are specified
+test_workload_identity | Make sure env vars are injected if workload identity is used
 
 ## Tar File [test file](test_confcom_tar.py)
 

src/confcom/azext_confcom/tests/latest/test_confcom_virtual_node.py

@@ -218,6 +218,8 @@ class PolicyGeneratingVirtualNode(unittest.TestCase):
 kind: Pod
 metadata:
   name: simple-container-pod
+  labels:
+    azure.workload.identity/use: "true"
 spec:
   initContainers:
     - name: init-container
@@ -314,3 +316,15 @@ def test_init_containers(self):
         self.assertEqual(containers[0][config.POLICY_FIELD_CONTAINERS_NAME], "simple-container")
         # see if the main container is in the policy
         self.assertEqual(containers[1][config.POLICY_FIELD_CONTAINERS_NAME], "init-container")
+
+    def test_workload_identity(self):
+        virtual_node_policy = load_policy_from_virtual_node_yaml_str(self.custom_yaml_init_containers)[0]
+        virtual_node_policy.populate_policy_content_for_all_images()
+        container_start = "containers := "
+        containers = json.loads(extract_containers_from_text(virtual_node_policy.get_serialized_output(OutputType.PRETTY_PRINT), container_start))
+
+        # have to extract the name from the pattern
+        env_rule_names = [(env_rule['pattern']).split("=")[0] for env_rule in containers[0][config.POLICY_FIELD_CONTAINERS_ELEMENTS_ENVS]]
+
+        for var in config.VIRTUAL_NODE_ENV_RULES_WORKLOAD_IDENTITY:
+          self.assertTrue(var['name'] in env_rule_names)