Create Network Security Group for test VMs (#2882)

* Create Network Security Group for test VMs

* error handling

---------

Co-authored-by: narrieta <narrieta>

tests_e2e/orchestrator/lib/update_arm_template_hook.py

@@ -15,10 +15,11 @@
 # limitations under the License.
 #
 
-import importlib
+import importlib.util
 import logging
+
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any
 
 # Disable those warnings, since 'lisa' is an external, non-standard, dependency
 # E0401: Unable to import 'lisa.*' (import-error)
@@ -29,6 +30,8 @@
 # pylint: enable=E0401
 
 import tests_e2e
+from tests_e2e.tests.lib.add_network_security_group import AddNetworkSecurityGroup
+from tests_e2e.tests.lib.update_arm_template import UpdateArmTemplate
 
 
 class UpdateArmTemplateHook:
@@ -37,30 +40,45 @@ class UpdateArmTemplateHook:
     """
     @hookimpl
     def azure_update_arm_template(self, template: Any, environment: Environment) -> None:
+        log: logging.Logger = logging.getLogger("lisa")
+
+        #
+        # Add the network security group for the test VM. This group includes a rule allowing SSH access from the current machine.
+        #
+        log.info("******** Waagent: Adding network security rule to the ARM template")
+        AddNetworkSecurityGroup().update(template)
+
+        #
+        # Apply any template customizations provided by the tests.
+        #
         azure_runbook: AzurePlatformSchema = environment.platform.runbook.get_extended_runbook(AzurePlatformSchema)
         vm_tags = azure_runbook.vm_tags
-        templates = vm_tags.get("templates")
-        if templates is not None:
-            log: logging.Logger = logging.getLogger("lisa")
-            log.info("******** Waagent: Applying custom templates '%s' to environment '%s'", templates, environment.name)
+        # The "templates" tag is a comma-separated list of the template customizations provided by the tests
+        test_templates = vm_tags.get("templates")
+        if test_templates is not None:
+            log.info("******** Waagent: Applying custom templates '%s' to environment '%s'", test_templates, environment.name)
 
-            for t in templates.split(","):
+            for t in test_templates.split(","):
                 update_arm_template = self._get_update_arm_template(t)
-                update_arm_template(template)
+                update_arm_template().update(template)
 
     _SOURCE_CODE_ROOT: Path = Path(tests_e2e.__path__[0])
 
     @staticmethod
-    def _get_update_arm_template(template_path: str) -> Callable:
-        source_file: Path = UpdateArmTemplateHook._SOURCE_CODE_ROOT/"tests"/template_path
+    def _get_update_arm_template(test_template: str) -> UpdateArmTemplate:
+        """
+        Returns the UpdateArmTemplate class that implements the template customization for the test.
+        """
+        source_file: Path = UpdateArmTemplateHook._SOURCE_CODE_ROOT/"tests"/test_template
 
         spec = importlib.util.spec_from_file_location(f"tests_e2e.tests.templates.{source_file.name}", str(source_file))
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
 
-        matches = [v for v in module.__dict__.values() if callable(v) and v.__name__ == "update_arm_template"]
+        # find all the classes in the module that are subclasses of UpdateArmTemplate but are not UpdateArmTemplate itself.
+        matches = [v for v in module.__dict__.values() if isinstance(v, type) and issubclass(v, UpdateArmTemplate) and v != UpdateArmTemplate]
         if len(matches) != 1:
-            raise Exception(f"Could not find update_arm_template in {source_file}")
+            raise Exception(f"Error in {source_file}: template files must contain exactly one class derived from UpdateArmTemplate)")
         return matches[0]
 
 

tests_e2e/test_suites/agent_not_provisioned.yml

@@ -6,7 +6,7 @@ name: "AgentNotProvisioned"
 tests:
   - "agent_not_provisioned/agent_not_provisioned.py"
 images: "random(endorsed)"
-template: "agent_not_provisioned/template.py"
+template: "agent_not_provisioned/disable_agent_provisioning.py"
 owns_vm: true
 install_test_agent: false
 

tests_e2e/test_suites/no_outbound_connections.yml

@@ -16,5 +16,5 @@ tests:
   - "agent_bvt/vm_access.py"
   - "no_outbound_connections/check_fallback_to_hgap.py"
 images: "random(endorsed)"
-template: "no_outbound_connections/template.py"
+template: "no_outbound_connections/deny_outbound_connections.py"
 owns_vm: true

tests_e2e/tests/agent_not_provisioned/disable_agent_provisioning.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+# Microsoft Azure Linux Agent
+#
+# Copyright 2018 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Any, Dict
+
+from tests_e2e.tests.lib.update_arm_template import UpdateArmTemplate
+
+
+class DenyOutboundConnections(UpdateArmTemplate):
+    """
+    Updates the ARM template to set osProfile.linuxConfiguration.provisionVMAgent to false.
+    """
+    def update(self, template: Dict[str, Any]) -> None:
+        #
+        # NOTE: LISA's template uses this function to generate the value for osProfile.linuxConfiguration. The function is
+        #       under the 'lisa' namespace.
+        #
+        #     "getLinuxConfiguration": {
+        #         "parameters": [
+        #             {
+        #                 "name": "keyPath",
+        #                 "type": "string"
+        #             },
+        #             {
+        #                 "name": "publicKeyData",
+        #                 "type": "string"
+        #             }
+        #         ],
+        #         "output": {
+        #             "type": "object",
+        #             "value": {
+        #                 "disablePasswordAuthentication": true,
+        #                 "ssh": {
+        #                     "publicKeys": [
+        #                         {
+        #                             "path": "[parameters('keyPath')]",
+        #                             "keyData": "[parameters('publicKeyData')]"
+        #                         }
+        #                     ]
+        #                 },
+        #                 "provisionVMAgent": true
+        #             }
+        #         }
+        #     }
+        #
+        # The code below sets template['functions'][i]['members']['getLinuxConfiguration']['output']['value']['provisionVMAgent'] to True,
+        # where template['functions'][i] is the 'lisa' namespace.
+        #
+        functions = template.get("functions")
+        if functions is None:
+            raise Exception('Cannot find "functions" in the LISA template.')
+        for namespace in functions:
+            name = namespace.get("namespace")
+            if name is None:
+                raise Exception(f'Cannot find "namespace" in the LISA template: {namespace}')
+            if name == "lisa":
+                members = namespace.get('members')
+                if members is None:
+                    raise Exception(f'Cannot find the members of the lisa namespace in the LISA template: {namespace}')
+                get_linux_configuration = members.get('getLinuxConfiguration')
+                if get_linux_configuration is None:
+                    raise Exception(f'Cannot find the "getLinuxConfiguration" function the lisa namespace in the LISA template: {namespace}')
+                output = get_linux_configuration.get('output')
+                if output is None:
+                    raise Exception(f'Cannot find the "output" of the getLinuxConfiguration function in the LISA template: {get_linux_configuration}')
+                value = output.get('value')
+                if value is None:
+                    raise Exception(f"Cannot find the output's value of the getLinuxConfiguration function in the LISA template: {get_linux_configuration}")
+                value['provisionVMAgent'] = False
+                break
+        else:
+            raise Exception(f'Cannot find the "lisa" namespace in the LISA template: {functions}')
+