Newly created scaleset proxy are behind NSG microsoft#1325

Apply existing rules to newly created scaleset in a new region microsoft#1326

src/api-service/__app__/instance_config/__init__.py

@@ -3,12 +3,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import logging
 import azure.functions as func
 from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error
 from onefuzztypes.requests import InstanceConfigUpdate
 
-from ..onefuzzlib.azure.nsg import set_allowed
+from ..onefuzzlib.azure.nsg import set_allowed, is_one_fuzz_nsg, list_nsgs
 from ..onefuzzlib.config import InstanceConfig
 from ..onefuzzlib.endpoint_authorization import call_if_user, can_modify_config
 from ..onefuzzlib.request import not_ok, ok, parse_request
@@ -46,18 +47,19 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
 
     # Update All NSGs
     if update_nsg:
-        scalesets = Scaleset.search()
-        regions = set(x.region for x in scalesets)
-        for region in regions:
-            result = set_allowed(region, request.config.proxy_nsg_config)
-            if isinstance(result, Error):
-                return not_ok(
-                    Error(
-                        code=ErrorCode.UNABLE_TO_CREATE,
-                        errors=["Unable to update nsg %s due to %s" % (region, result)],
-                    ),
-                    context="instance_config_update",
-                )
+        nsgs = list_nsgs()
+        for nsg in nsgs:
+            logging.info("Checking if nsg: %s (%s) owned by OneFuzz" % (nsg.location, nsg.name))
+            if is_one_fuzz_nsg(nsg.location, nsg.name):
+                result = set_allowed(nsg.location, request.config.proxy_nsg_config)
+                if isinstance(result, Error):
+                    return not_ok(
+                        Error(
+                            code=ErrorCode.UNABLE_TO_CREATE,
+                            errors=["Unable to update nsg %s due to %s" % (nsg.location, result)],
+                        ),
+                        context="instance_config_update",
+                    )
 
     return ok(config)
 

src/api-service/__app__/onefuzzlib/azure/nsg.py

@@ -13,7 +13,6 @@
     NetworkSecurityGroup,
     SecurityRule,
     SecurityRuleAccess,
-    SecurityRuleProtocol,
     Subnet,
     VirtualNetwork,
 )
@@ -106,6 +105,8 @@ def update_nsg(nsg: NetworkSecurityGroup) -> Union[None, Error]:
 def ok_to_delete(active_regions: Set[Region], nsg_region: str, nsg_name: str) -> bool:
     return nsg_region not in active_regions and nsg_region == nsg_name
 
+def is_one_fuzz_nsg(nsg_region: str, nsg_name: str) -> bool:
+    return nsg_region == nsg_name
 
 def delete_nsg(name: str) -> bool:
     # NSG can be only deleted if no other resource is associated with it
@@ -165,7 +166,7 @@ def set_allowed(name: str, sources: NetworkSecurityGroupConfig) -> Union[None, E
         security_rules.append(
             SecurityRule(
                 name="Allow" + str(priority),
-                protocol=SecurityRuleProtocol.ANY,
+                protocol="*",
                 source_port_range="*",
                 destination_port_range="*",
                 source_address_prefix=src,

src/api-service/functional_tests/__init__.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.

src/api-service/functional_tests/nsg_test.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import subprocess
+import sys
+import uuid
+
+from onefuzz.api import Onefuzz
+from onefuzztypes.enums import OS
+
+from __app__.onefuzzlib.proxy import Proxy
+
+
+class NsgTests:
+    def __init__(self):
+        self.onefuzz = Onefuzz()
+
+    def allow_all(self):
+        instance_config = self.onefuzz.instance_config.get()
+        instance_config.proxy_nsg_config.allowed_ips = ["*"]
+        self.onefuzz.instance_config.update(instance_config)
+
+    def block_all(self):
+        instance_config = self.onefuzz.instance_config.get()
+        instance_config.proxy_nsg_config.allowed_ips = []
+        self.onefuzz.instance_config.update(instance_config)
+
+    def get_proxy_for_region(self, region):
+        proxies = Proxy.search()
+        for proxy in proxies:
+            if proxy.region == region:
+                return proxy
+        return None
+
+    def ping(self, proxy):
+
+        if proxy.ip is None:
+            print("Proxy IP is missing")
+            return False
+        cmd = ["ping", proxy.ip]
+        if sys.platform == "linux":
+            cmd.append("-w")
+            cmd.append("5")
+
+        return 0 == subprocess.call(cmd)
+
+    def nsg_test_name(self):
+        return "nsg-test-%s" % uuid.uuid4().hex
+
+    def test_proxy_access(self, region):
+        pool_name = self.nsg_test_name()
+        pool = self.onefuzz.pools.create(pool_name, OS.linux)
+        scaleset = self.onefuzz.scalesets.create(pool_name, 1, region=region)
+        try:
+            # expects existing deployment with a debug proxy and a scaleset in
+            # specified region
+            proxy = self.get_proxy_for_region(region)
+            print("Got proxy: %s for region %s" % (proxy.proxy_id, region))
+            #self.onefuzz.scalesets.
+
+
+            self.allow_all()
+            # can ping since all allowed
+            if not self.ping(proxy):
+                raise Exception("Failed to reach proxy in %s region" % region)
+
+            self.block_all()
+            # should not be able to ping since all blocked
+            if self.ping(proxy):
+                raise Exception("Managed to ping proxy in %s region" % region)
+        finally:
+            #self.onefuzz.scalesets.shutdown(scaleset.scaleset_id, now=True)
+            pass

src/api-service/functional_tests/nsg_tst.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import subprocess  # nosec
+import sys
+import time
+import uuid
+
+from onefuzz.api import Onefuzz
+from onefuzztypes.enums import OS, ScalesetState
+
+
+class NsgTests:
+    def __init__(self):
+        self.onefuzz = Onefuzz()
+
+    def allow_all(self):
+        instance_config = self.onefuzz.instance_config.get()
+        instance_config.proxy_nsg_config.allowed_ips = ["*"]
+        self.onefuzz.instance_config.update(instance_config)
+
+    def block_all(self):
+        instance_config = self.onefuzz.instance_config.get()
+        instance_config.proxy_nsg_config.allowed_ips = []
+        self.onefuzz.instance_config.update(instance_config)
+
+    def generate_name(self):
+        return "nsg-test-%s" % uuid.uuid4().hex
+
+    def get_proxy_ip(self, scaleset):
+        machine_id = scaleset.nodes[0].machine_id
+        scaleset_id = scaleset.scaleset_id
+        timeout_seconds = 60
+        time_waited = 0
+        wait_poll = 5
+
+        ip = None
+        print("Retrieving proxy IP for region: %s" % scaleset.region)
+
+        while time_waited < timeout_seconds and ip is None:
+            proxy = self.onefuzz.scaleset_proxy.create(
+                scaleset_id, machine_id, dst_port=1
+            )
+            ip = proxy.ip
+            if ip is None:
+                time.sleep(wait_poll)
+
+        if ip is None:
+            raise Exception("Failed to get proxy IP for region: %s" % scaleset.region)
+        return ip
+
+    def test_connection(self, ip):
+        cmd = ["ping", ip]
+        if sys.platform == "linux":
+            cmd.append("-w")
+            cmd.append("5")
+        return 0 == subprocess.call(cmd)  # nosec
+
+    def wait_for_scaleset(self, scaleset):
+        timeout_seconds = 600
+        wait_poll = 10
+        # wait for scaleset creation to finish
+        time_waited = 0
+        tmp_scaleset = scaleset
+        while (
+            time_waited < timeout_seconds
+            and tmp_scaleset.error is None
+            and tmp_scaleset.state != ScalesetState.running
+        ):
+            tmp_scaleset = self.onefuzz.scalesets.get(tmp_scaleset.scaleset_id)
+            print(
+                "Waiting for scaleset creation... Current scaleset state: %s"
+                % tmp_scaleset.state
+            )
+            time.sleep(wait_poll)
+            time_waited = time_waited + wait_poll
+
+        if tmp_scaleset.error:
+            raise Exception(
+                "Failed to provision scaleset %s" % (tmp_scaleset.scaleset_id)
+            )
+
+        return tmp_scaleset
+
+    def create_pool(self, pool_name):
+        self.onefuzz.pools.create(pool_name, OS.linux)
+
+    def create_scaleset(self, pool_name, region):
+        scaleset = self.onefuzz.scalesets.create(pool_name, 1, region=region)
+        return self.wait_for_scaleset(scaleset)
+
+    def wait_for_nsg_rules_to_apply(self):
+        time.sleep(120)
+
+    def test_proxy_access(self, pool_name, region):
+        scaleset = self.create_scaleset(pool_name, region)
+        ip = self.get_proxy_ip(scaleset)
+        print("Allow connection")
+
+        self.allow_all()
+        self.wait_for_nsg_rules_to_apply()
+        # can ping since all allowed
+        result = self.test_connection(ip)
+        if not result:
+            raise Exception("Failed to connect to proxy")
+
+        print("Block connection")
+        self.block_all()
+        self.wait_for_nsg_rules_to_apply()
+        # should not be able to ping since all blocked
+        result = self.test_connection(ip)
+        if result:
+            raise Exception("Connected to proxy")
+
+        print("Allow connection")
+        self.allow_all()
+        self.wait_for_nsg_rules_to_apply()
+        # can ping since all allowed
+        result = self.test_connection(ip)
+        if not result:
+            raise Exception("Failed to connect to proxy")
+
+        print("Block connection")
+        self.block_all()
+        self.wait_for_nsg_rules_to_apply()
+        # should not be able to ping since all blocked
+        result = self.test_connection(ip)
+        if result:
+            raise Exception("Connected to proxy")
+
+    def test_new_scaleset_region(self, pool_name, region1, region2):
+        if region1 == region2:
+            raise Exception(
+                (
+                    "Test input parameter validation failure.",
+                    " Scalesets expted to be in different regions",
+                )
+            )
+
+        scaleset1 = self.create_scaleset(pool_name, region1)
+        ip1 = self.get_proxy_ip(scaleset1)
+
+        print("Block connection")
+        self.block_all()
+        self.wait_for_nsg_rules_to_apply()
+        # should not be able to ping since all blocked
+        print("Attempting connection for region %s" % region1)
+        result = self.test_connection(ip1)
+        if result:
+            raise Exception("Connected to proxy1 in region %s" % region1)
+
+        print("Allow connection")
+        self.allow_all()
+        self.wait_for_nsg_rules_to_apply()
+        # can ping since all allowed
+        print("Attempting connection for region %s" % region1)
+        result = self.test_connection(ip1)
+        if not result:
+            raise Exception("Failed to connect to proxy1 in region %s" % region1)
+
+        print("Creating scaleset in region %s" % region2)
+        scaleset2 = self.create_scaleset(pool_name, region2)
+
+        ip2 = self.get_proxy_ip(scaleset2)
+        # should not be able to ping since all blocked
+        print("Attempting connection for region %s" % region2)
+        result = self.test_connection(ip2)
+        if not result:
+            raise Exception("Failed to connect to proxy2 in region %s" % region2)
+
+        print("Block connection")
+        self.block_all()
+        self.wait_for_nsg_rules_to_apply()
+        # should not be able to ping since all blocked
+        print("Attempting connection for region %s" % region1)
+        result = self.test_connection(ip1)
+        if result:
+            raise Exception("Connected to proxy1 in region" % region1)
+
+        # should not be able to ping since all blocked
+        print("Attempting connection for region %s" % region2)
+        result = self.test_connection(ip2)
+        if result:
+            raise Exception("Connected to proxy2 in region %s" % region2)
+
+        # self.onefuzz.scalesets.shutdown(scaleset1.scaleset_id, now=True)
+
+
+def main():
+    t = NsgTests()
+    pool_name = t.generate_name()
+    t.create_pool(pool_name)
+    print("Test basic proxy access")
+    t.test_proxy_access(pool_name, "eastus")
+    print("Test new region addition access")
+    t.test_new_scaleset_region(pool_name, "westus2", "westus3")
+
+
+if __name__ == "__main__":
+    main()