address comments

tests_e2e/test_suites/agent_firewall.yml

@@ -10,6 +10,5 @@ name: "AgentFirewall"
 tests:
   - "agent_firewall/agent_firewall.py"
 images:
-  - "endorsed"
-  - "endorsed-arm64"
+  - "ubuntu_2004"
 owns_vm: true  # This vm cannot be shared with other tests because it modifies the firewall rules and agent status.

tests_e2e/tests/agent_firewall/agent_firewall.py

@@ -31,16 +31,10 @@ def __init__(self, context: AgentTestContext):
         self._ssh_client = self._context.create_ssh_client()
 
     def run(self):
-        self._prepare_agent()
         log.info("Checking iptable rules added by the agent")
         self._run_remote_test(f"agent_firewall-verify_all_firewall_rules.py --user {self._context.username}", use_sudo=True)
         log.info("Successfully verified all rules present and working as expected.")
 
-    def _prepare_agent(self) -> None:
-        log.info("Executing remote script update-waagent-conf to enable agent firewall config flag")
-        self._run_remote_test("update-waagent-conf OS.EnableFirewall=y", use_sudo=True)
-        log.info("Successfully enabled agent firewall config flag")
-
 
 if __name__ == "__main__":
     AgentFirewall.run_from_command_line()

tests_e2e/tests/lib/add_network_security_group.py

@@ -17,11 +17,13 @@
 
 import json
 import http.client
+from subprocess import PIPE, Popen
 
 from typing import Any, Dict, List
 
 from tests_e2e.tests.lib.logging import log
 from tests_e2e.tests.lib.retry import retry
+from tests_e2e.tests.lib.shell import CommandError
 from tests_e2e.tests.lib.update_arm_template import UpdateArmTemplate
 
 # Name of the security group added by this class
@@ -140,10 +142,13 @@ def _my_ip_address(self) -> str:
         """
         if self.__my_ip_address is None:
             def get_my_address():
-                connection = http.client.HTTPSConnection("ifconfig.io")
-                connection.request("GET", "/ip")
-                response = connection.getresponse()
-                return response.read().decode().strip()
+                # Forcing -4 option to fetch the ipv4 address
+                cmd = ["curl", "-4", "ifconfig.io/ip"]
+                process = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=False, text=True)
+                stdout, stderr = process.communicate()
+                if process.returncode != 0:
+                    raise CommandError(cmd, process.returncode, stdout, stderr)
+                return stdout.strip()
             self.__my_ip_address = retry(get_my_address, attempts=3, delay=10)
         return self.__my_ip_address
 

tests_e2e/tests/lib/retry.py

@@ -72,9 +72,9 @@ def retry_if_false(operation: Callable[[], bool], attempts: int = 5, delay: int
             log.warning("Error in operation: %s", e)
             if attempts == 0:
                 raise
-        if not success:
+        if not success and attempts != 0:
             log.info("Current operation failed, retrying in %s secs.", delay)
-        time.sleep(delay)
+            time.sleep(delay)
     return success
 
 

tests_e2e/tests/scripts/agent_firewall-verify_all_firewall_rules.py

@@ -22,34 +22,40 @@
 import os
 import pwd
 import socket
-import subprocess
-import sys
-import time
 from typing import List, Tuple
 
 from assertpy import fail
 
 from azurelinuxagent.common.utils import shellutil
+from azurelinuxagent.common.utils.shellutil import CommandError
+from azurelinuxagent.common.utils.textutil import format_exception
 from tests_e2e.tests.lib.logging import log
 from tests_e2e.tests.lib.remote_test import run_remote_test
+import http.client as httpclient
 
-if sys.version_info[0] == 3:
-    import http.client as httpclient
-elif sys.version_info[0] == 2:
-    import httplib as httpclient  # pylint: disable=E0401
+from tests_e2e.tests.lib.retry import retry_if_false, retry
 
 ROOT_USER = 'root'
 WIRESERVER_ENDPOINT_FILE = '/var/lib/waagent/WireServerEndpoint'
-WIRESERVER_URL = '168.63.129.16'
+WIRESERVER_IP = '168.63.129.16'
 VERSIONS_PATH = '/?comp=versions'
+FIREWALL_PERIOD = 60
+
+
+class FirewallRules(object):
+    # -D deletes the specific rule in the iptable chain
+    DELETE_COMMAND = "-D"
+
+    # -C checks if a specific rule exists
+    CHECK_COMMAND = "-C"
 
 
 def get_wireserver_ip() -> str:
     try:
         with open(WIRESERVER_ENDPOINT_FILE, 'r') as f:
             wireserver_ip = f.read()
     except Exception:
-        wireserver_ip = WIRESERVER_URL
+        wireserver_ip = WIRESERVER_IP
     return wireserver_ip
 
 
@@ -65,35 +71,34 @@ def switch_user(user: str) -> None:
         raise Exception("Error -- failed to switch user to {0} : Failed with exception {1}".format(user, e))
 
 
-def get_root_accept_rule(command: str) -> List[str]:
+def get_root_accept_rule_command(command: str) -> List[str]:
     return ['sudo', 'iptables', '-t', 'security', command, 'OUTPUT', '-d', get_wireserver_ip(), '-p', 'tcp', '-m',
             'owner',
             '--uid-owner',
             '0', '-j', 'ACCEPT', '-w']
 
 
-def get_non_root_accept_rule(command: str) -> List[str]:
+def get_non_root_accept_rule_command(command: str) -> List[str]:
     return ['sudo', 'iptables', '-t', 'security', command, 'OUTPUT', '-d', get_wireserver_ip(), '-p', 'tcp',
             '--destination-port', '53', '-j',
             'ACCEPT', '-w']
 
 
-def get_non_root_drop_rule(command: str) -> List[str]:
+def get_non_root_drop_rule_command(command: str) -> List[str]:
     return ['sudo', 'iptables', '-t', 'security', command, 'OUTPUT', '-d', get_wireserver_ip(), '-p', 'tcp', '-m',
             'conntrack', '--ctstate',
             'INVALID,NEW', '-j', 'DROP', '-w']
 
 
 def execute_cmd(cmd: List[str]):
     """
-    Note: The shellutil.run_command don't return the exit code of the command. So using subprocess.Popen
+    Note: The shellutil.run_command return stdout if exit_code=0, otherwise returns commanderror
     """
-    proc = subprocess.Popen(cmd,
-                            stdout=subprocess.PIPE,
-                            stderr=subprocess.PIPE,
-                            shell=False)
-    stdout, stderr = proc.communicate()
-    return proc.returncode, stdout, stderr
+    try:
+        stdout = shellutil.run_command(cmd)
+    except CommandError as e:
+        return e.returncode, e.stdout, e.stderr
+    return 0, stdout, ""
 
 
 def check_if_iptable_rule_is_available(full_command: List[str]) -> bool:
@@ -124,25 +129,22 @@ def print_current_iptable_rules():
 
 
 def get_all_iptable_rule_commands(command: str) -> Tuple[List[str], List[str], List[str]]:
-    return get_root_accept_rule(command), get_non_root_accept_rule(command), get_non_root_drop_rule(command)
+    return get_root_accept_rule_command(command), get_non_root_accept_rule_command(command), get_non_root_drop_rule_command(command)
 
 
-def verify_all_rules_exist(max_retry: int = 5) -> None:
+def verify_all_rules_exist() -> None:
     """
     This function is used to verify all the iptable rules are present in the rule set
     """
-    log.info("-----Verifying all ip table rules are present in rule set")
-    retry: int = 0
-    found: bool = False
-    while retry < max_retry and not found:
-        root_accept, non_root_accept, non_root_drop = get_all_iptable_rule_commands("-C")
-        found = check_if_iptable_rule_is_available(root_accept) and check_if_iptable_rule_is_available(
+    def check_all_iptables() -> bool:
+        root_accept, non_root_accept, non_root_drop = get_all_iptable_rule_commands(FirewallRules.CHECK_COMMAND)
+        found: bool = check_if_iptable_rule_is_available(root_accept) and check_if_iptable_rule_is_available(
             non_root_accept) and check_if_iptable_rule_is_available(non_root_drop)
+        return found
 
-        if not found:
-            log.info("Not all ip table rules are present, retrying in 30 secs.\n")
-            time.sleep(30)
-        retry += 1
+    log.info("-----Verifying all ip table rules are present in rule set")
+    # Agent will re-add rules within OS.EnableFirewallPeriod, So waiting that time + some buffer
+    found: bool = retry_if_false(check_all_iptables, attempts=2, delay=FIREWALL_PERIOD+30)
 
     if not found:
         fail("IP table rules missing in rule set.\n Current iptable rules:\n {0}".format(
@@ -181,14 +183,14 @@ def delete_iptable_rules(commands: List[List[str]] = None) -> None:
     if commands is None:
         commands = []
     if not commands:
-        root_accept, non_root_accept, non_root_drop = get_all_iptable_rule_commands("-C")
+        root_accept, non_root_accept, non_root_drop = get_all_iptable_rule_commands(FirewallRules.CHECK_COMMAND)
         commands.extend([root_accept, non_root_accept, non_root_drop])
 
     log.info("-----Deleting ip table rules \n %s", commands)
 
     try:
         for command in commands:
-            execute_cmd(command)
+            retry(lambda: execute_cmd(command), attempts=3)
     except Exception as e:
         raise Exception("Error -- Failed to Delete the ip table rule set {0}".format(e))
 
@@ -202,7 +204,7 @@ def verify_dns_tcp_to_wireserver_is_allowed(user: str) -> None:
     log.info("-----Verifying DNS tcp to wireserver is allowed")
     switch_user(user)
     try:
-        socket.create_connection((get_wireserver_ip(), 53), timeout=10)
+        socket.create_connection((get_wireserver_ip(), 53), timeout=30)
     except Exception as e:
         raise Exception(
             "Error -- while using DNS TCP request as user:({0}), make sure the firewall rules are set correctly {1}".format(user,
@@ -220,8 +222,12 @@ def verify_dns_tcp_to_wireserver_is_blocked(user: str) -> None:
     try:
         socket.create_connection((get_wireserver_ip(), 53), timeout=10)
         raise Exception("Error -- unprivileged user:({0}) could connect to wireserver port 53 using TCP".format(user))
-    except Exception:
-        log.info("Success -- unprivileged user:(%s) access to wireserver port 53 using TCP is blocked", user)
+    except Exception as e:
+        # Expected timeout if unprivileged user reaches wireserver
+        if isinstance(e, socket.timeout):
+            log.info("Success -- unprivileged user:(%s) access to wireserver port 53 using TCP is blocked", user)
+        else:
+            raise Exception("Unexpected error while connecting to wireserver: {0}".format(format_exception(e)))
 
 
 def verify_http_to_wireserver_blocked(user: str) -> None:
@@ -231,16 +237,19 @@ def verify_http_to_wireserver_blocked(user: str) -> None:
     log.info("-----Verifying http request to wireserver is blocked")
     switch_user(user)
     try:
-        client = httpclient.HTTPConnection(get_wireserver_ip(), timeout=1)
+        client = httpclient.HTTPConnection(get_wireserver_ip(), timeout=10)
     except Exception as e:
         raise Exception("Error -- failed to create HTTP connection with user: {0} \n {1}".format(user, e))
 
     try:
         blocked = False
         client.request('GET', VERSIONS_PATH)
-    except Exception:
-        # if we get exception, it means the request is blocked
-        blocked = True
+    except Exception as e:
+        # if we get timeout exception, it means the request is blocked
+        if isinstance(e, socket.timeout):
+            blocked = True
+        else:
+            raise Exception("Unexpected error while connecting to wireserver: {0}".format(format_exception(e)))
 
     if not blocked:
         raise Exception("Error -- unprivileged user:({0}) could connect to wireserver, make sure the firewall rules are set correctly".format(user))
@@ -255,7 +264,7 @@ def verify_http_to_wireserver_allowed(user: str) -> None:
     log.info("-----Verifying http request to wireserver is allowed")
     switch_user(user)
     try:
-        client = httpclient.HTTPConnection(get_wireserver_ip(), timeout=1)
+        client = httpclient.HTTPConnection(get_wireserver_ip(), timeout=30)
     except Exception as e:
         raise Exception("Error -- failed to create HTTP connection with user:{0} \n {1}".format(user, e))
 
@@ -286,10 +295,10 @@ def verify_non_root_accept_rule():
     shellutil.run_command(stop_agent)
 
     # deleting non root accept rule
-    non_root_accept_delete_cmd = get_non_root_accept_rule("-D")
+    non_root_accept_delete_cmd = get_non_root_accept_rule_command(FirewallRules.DELETE_COMMAND)
     delete_iptable_rules([non_root_accept_delete_cmd])
     # verifying deletion successful
-    non_root_accept_check_cmd = get_non_root_accept_rule("-C")
+    non_root_accept_check_cmd = get_non_root_accept_rule_command(FirewallRules.CHECK_COMMAND)
     verify_rules_deleted_successfully([non_root_accept_check_cmd])
 
     log.info("** Current IP table rules\n")
@@ -313,6 +322,15 @@ def verify_non_root_accept_rule():
     verify_dns_tcp_to_wireserver_is_allowed(NON_ROOT_USER)
     verify_http_to_wireserver_blocked(NON_ROOT_USER)
 
+    log.info("Ensuring missing rules are re-added by the running agent")
+    # deleting non root accept rule
+    non_root_accept_delete_cmd = get_non_root_accept_rule_command(FirewallRules.DELETE_COMMAND)
+    delete_iptable_rules([non_root_accept_delete_cmd])
+
+    verify_all_rules_exist()
+    log.info("** Current IP table rules \n")
+    print_current_iptable_rules()
+
     log.info("non root accept rule verified successfully\n")
 
 
@@ -334,13 +352,13 @@ def verify_root_accept_rule():
     shellutil.run_command(stop_agent)
 
     # deleting root accept rule
-    root_accept_delete_cmd = get_root_accept_rule("-D")
+    root_accept_delete_cmd = get_root_accept_rule_command(FirewallRules.DELETE_COMMAND)
     # deleting drop rule too otherwise after restart, the daemon will go into loop since it cannot connect to wireserver. This would block the agent initialization
-    drop_delete_cmd = get_non_root_drop_rule("-D")
+    drop_delete_cmd = get_non_root_drop_rule_command(FirewallRules.DELETE_COMMAND)
     delete_iptable_rules([root_accept_delete_cmd, drop_delete_cmd])
     # verifying deletion successful
-    root_accept_check_cmd = get_root_accept_rule("-C")
-    drop_check_cmd = get_non_root_drop_rule("-C")
+    root_accept_check_cmd = get_root_accept_rule_command(FirewallRules.CHECK_COMMAND)
+    drop_check_cmd = get_non_root_drop_rule_command(FirewallRules.CHECK_COMMAND)
     verify_rules_deleted_successfully([root_accept_check_cmd, drop_check_cmd])
 
     log.info("** Current IP table rules\n")
@@ -361,6 +379,15 @@ def verify_root_accept_rule():
     verify_http_to_wireserver_blocked(NON_ROOT_USER)
     verify_http_to_wireserver_allowed(ROOT_USER)
 
+    log.info("Ensuring missing rules are re-added by the running agent")
+    # deleting root accept rule
+    root_accept_delete_cmd = get_root_accept_rule_command(FirewallRules.DELETE_COMMAND)
+    delete_iptable_rules([root_accept_delete_cmd])
+
+    verify_all_rules_exist()
+    log.info("** Current IP table rules \n")
+    print_current_iptable_rules()
+
     log.info("root accept rule verified successfully\n")
 
 
@@ -378,10 +405,10 @@ def verify_non_root_dcp_rule():
     shellutil.run_command(stop_agent)
 
     # deleting non root delete rule
-    non_root_drop_delete_cmd = get_non_root_drop_rule("-D")
+    non_root_drop_delete_cmd = get_non_root_drop_rule_command(FirewallRules.DELETE_COMMAND)
     delete_iptable_rules([non_root_drop_delete_cmd])
     # verifying deletion successful
-    non_root_drop_check_cmd = get_non_root_drop_rule("-C")
+    non_root_drop_check_cmd = get_non_root_drop_rule_command(FirewallRules.CHECK_COMMAND)
     verify_rules_deleted_successfully([non_root_drop_check_cmd])
 
     log.info("** Current IP table rules\n")
@@ -405,10 +432,27 @@ def verify_non_root_dcp_rule():
     verify_http_to_wireserver_blocked(NON_ROOT_USER)
     verify_http_to_wireserver_allowed(ROOT_USER)
 
+    log.info("Ensuring missing rules are re-added by the running agent")
+    # deleting non root delete rule
+    non_root_drop_delete_cmd = get_non_root_drop_rule_command(FirewallRules.DELETE_COMMAND)
+    delete_iptable_rules([non_root_drop_delete_cmd])
+
+    verify_all_rules_exist()
+    log.info("** Current IP table rules\n")
+    print_current_iptable_rules()
+
     log.info("non root drop rule verified successfully\n")
 
 
+def prepare_agent():
+    log.info("Executing script update-waagent-conf to enable agent firewall config flag")
+    # Changing the firewall period from default 5 mins to 1 min, so that test won't wait for that long to verify rules
+    shellutil.run_command(["update-waagent-conf", "OS.EnableFirewall=y", f"OS.EnableFirewallPeriod={FIREWALL_PERIOD}"])
+    log.info("Successfully enabled agent firewall config flag")
+
+
 def main():
+    prepare_agent()
     log.info("** Current IP table rules\n")
     print_current_iptable_rules()