Update NSGs after changes to instance config (#1385)

* Refactor set_admins into configuration.py and update deployment params with nsg_config

* Fixing arguments.

* Param takes in network config json

* Fixing Client in deploy

* removing import

* Adding onefuzztypes to reqs.txt

* Reverting to single list

* Removing imports.

* Retriggering build

* Setting specific pip version for local testing.

* Removing imports?

* More imports.

* Fixing formatting.

* Updating how to parse nsg param.

* Removing old logging statements.

* Fixing types.

* REmoving bad log

* Removing local pip version.

* Removing comments

* fixing

* Formatting

* Fixing .split()

* Adding NSG rule checks and type.

* Formatting.

* Formatting.

* Removing imports.

* Fixing formatting.

* Testing formatting.

* Retrigger?

* New InstanceConfigClient class.

* Retrigger.

* Cherry picked commit.

* Reformatting.

* Actually fixing formatting.

* Fixing table_service call.

* Fixing return statement and nsg_rule pass.

* Full config.

* Removing commented out code.

* Fixing logic.

* Adding wildcard check.

* Code for updating NSGs when instance_config updated.

* Updating argument to set_allowed_rules

* Updating model to no longer be optional.

* Fixing args for set_allowed_rules

* trying to fix calls to get_nsg

* Updating calls to nsg lib

* Fixing imports.

* Updating calls to set_allowed and creating constructor for NSGConfig type.

* Removing constructor and manually setting default ip

* Fixing models.

* Hopefully fixing docs.

* Fix set_allowed call

* Adding error handling for update config.

* Changing to error check.

* Fixing error call.

* Fixing imports.

* Adding empty() function on request.

* Removing empty function.

# Conflicts:
#	src/pytypes/onefuzztypes/models.py

* Fixing files for update.

* Fixing nsg.py.

* Fixing imports.

* removing commented code.

Co-authored-by: nharper285 <nharper285@gmail.com>

docs/webhook_events.md

@@ -650,7 +650,8 @@ Each event will be submitted via HTTP POST to the user provided URL.
             "subnet": "10.0.0.0/16"
         },
         "proxy_nsg_config": {
-            "allowed_ips": []
+            "allowed_ips": [],
+            "allowed_service_tags": []
         },
         "proxy_vm_sku": "Standard_B2s"
     }

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

@@ -8,9 +8,11 @@
 from onefuzztypes.models import Error
 from onefuzztypes.requests import InstanceConfigUpdate
 
+from ..onefuzzlib.azure.nsg import set_allowed
 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
+from ..onefuzzlib.workers.scalesets import Scaleset
 
 
 def get(req: func.HttpRequest) -> func.HttpResponse:
@@ -30,8 +32,33 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
             context="instance_config_update",
         )
 
+    update_nsg = False
+    if request.config.proxy_nsg_config and config.proxy_nsg_config:
+        request_config = request.config.proxy_nsg_config
+        current_config = config.proxy_nsg_config
+        if set(request_config.allowed_service_tags) != set(
+            current_config.allowed_service_tags
+        ) or set(request_config.allowed_ips) != set(current_config.allowed_ips):
+            update_nsg = True
+
     config.update(request.config)
     config.save()
+
+    # 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",
+                )
+
     return ok(config)
 
 

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

@@ -17,7 +17,7 @@
 )
 from msrestazure.azure_exceptions import CloudError
 from onefuzztypes.enums import ErrorCode
-from onefuzztypes.models import Error
+from onefuzztypes.models import Error, NetworkSecurityGroupConfig
 from onefuzztypes.primitives import Region
 from pydantic import BaseModel, validator
 
@@ -127,7 +127,7 @@ def delete_nsg(name: str) -> bool:
     return False
 
 
-def set_allowed(name: str, sources: List[str]) -> Union[None, Error]:
+def set_allowed(name: str, sources: NetworkSecurityGroupConfig) -> Union[None, Error]:
     resource_group = get_base_resource_group()
     nsg = get_nsg(name)
     if not nsg:
@@ -141,24 +141,25 @@ def set_allowed(name: str, sources: List[str]) -> Union[None, Error]:
         resource_group,
         name,
     )
+    all_sources = sources.allowed_ips + sources.allowed_service_tags
     security_rules = []
     # NSG security rule priority range defined here:
     # https://docs.microsoft.com/en-us/azure/virtual-network/network-security-groups-overview
     min_priority = 100
     # NSG rules per NSG limits:
     # https://docs.microsoft.com/en-us/azure/azure-resource-manager/management/azure-subscription-service-limits?toc=/azure/virtual-network/toc.json#networking-limits
     max_rule_count = 1000
-    if len(sources) > max_rule_count:
+    if len(all_sources) > max_rule_count:
         return Error(
             code=ErrorCode.INVALID_REQUEST,
             errors=[
                 "too many rules provided %d. Max allowed: %d"
-                % ((len(sources)), max_rule_count),
+                % ((len(all_sources)), max_rule_count),
             ],
         )
 
     priority = min_priority
-    for src in sources:
+    for src in all_sources:
         security_rules.append(
             SecurityRule(
                 name="Allow" + str(priority),
@@ -173,15 +174,15 @@ def set_allowed(name: str, sources: List[str]) -> Union[None, Error]:
             )
         )
         # Will not exceed `max_rule_count` or max NSG priority (4096)
-        # due to earlier check of `len(sources)`.
+        # due to earlier check of `len(all_sources)`.
         priority += 1
 
     nsg.security_rules = security_rules
     return update_nsg(nsg)
 
 
 def clear_all_rules(name: str) -> Union[None, Error]:
-    return set_allowed(name, [])
+    return set_allowed(name, NetworkSecurityGroupConfig())
 
 
 def get_all_rules(name: str) -> Union[Error, List[SecurityRule]]:
@@ -328,7 +329,9 @@ def delete(self) -> bool:
     def get(self) -> Optional[NetworkSecurityGroup]:
         return get_nsg(self.name)
 
-    def set_allowed_sources(self, sources: List[str]) -> Union[None, Error]:
+    def set_allowed_sources(
+        self, sources: NetworkSecurityGroupConfig
+    ) -> Union[None, Error]:
         return set_allowed(self.name, sources)
 
     def clear_all_rules(self) -> Union[None, Error]:

src/api-service/__app__/onefuzzlib/proxy.py

@@ -101,7 +101,9 @@ def init(self) -> None:
                 self.set_failed(result)
                 return
 
-            result = nsg.set_allowed_sources(["*"])
+            config = InstanceConfig.fetch()
+            nsg_config = config.proxy_nsg_config
+            result = nsg.set_allowed_sources(nsg_config)
             if isinstance(result, Error):
                 self.set_failed(result)
                 return

src/api-service/__app__/onefuzzlib/repro.py

@@ -21,6 +21,7 @@
 from .azure.nsg import NSG
 from .azure.storage import StorageType
 from .azure.vm import VM
+from .config import InstanceConfig
 from .extension import repro_extensions
 from .orm import ORMMixin, QueryFilter
 from .reports import get_report
@@ -95,7 +96,9 @@ def init(self) -> None:
                 self.set_failed(result)
                 return
 
-            result = nsg.set_allowed_sources(["*"])
+            config = InstanceConfig.fetch()
+            nsg_config = config.proxy_nsg_config
+            result = nsg.set_allowed_sources(nsg_config)
             if isinstance(result, Error):
                 self.set_failed(result)
                 return

src/deployment/configuration.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import ipaddress
+import json
+import logging
+from typing import List, Optional
+from uuid import UUID
+
+from azure.cosmosdb.table.tableservice import TableService
+
+storage_client_logger = logging.getLogger("azure.cosmosdb.table.common.storageclient")
+TABLE_NAME = "InstanceConfig"
+
+logger = logging.getLogger("deploy")
+
+
+class InstanceConfigClient:
+
+    table_service: TableService
+    resource_group: str
+
+    def __init__(self, table_service: TableService, resource_group: str):
+        self.resource_group = resource_group
+        self.table_service = table_service
+        self.create_if_missing(table_service)
+
+    ## Disable logging from storageclient. This module displays an error message
+    ## when a resource is not found even if the exception is raised and handled internally.
+    ## This happen when a table does not exist. An error message is displayed but the exception is
+    ## handled by the library.
+    def disable_storage_client_logging(self) -> None:
+        if storage_client_logger:
+            storage_client_logger.disabled = True
+
+    def enable_storage_client_logging(self) -> None:
+        if storage_client_logger:
+            storage_client_logger.disabled = False
+
+    def create_if_missing(self, table_service: TableService) -> None:
+        try:
+            self.disable_storage_client_logging()
+
+            if not table_service.exists(TABLE_NAME):
+                table_service.create_table(TABLE_NAME)
+        finally:
+            self.enable_storage_client_logging()
+
+
+class NsgRule:
+
+    rule: str
+    is_tag: bool
+
+    def __init__(self, rule: str):
+        try:
+            self.is_tag = False
+            self.check_rule(rule)
+            self.rule = rule
+        except Exception:
+            raise ValueError(
+                "Invalid rule. Please provide a valid rule or supply the wild card *."
+            )
+
+    def check_rule(self, value: str) -> None:
+        if value is None or len(value.strip()) == 0:
+            raise ValueError(
+                "Rule can not be None or empty string. Please provide a valid rule or supply the wild card *."
+            )
+        # Check Wild Card
+        if value == "*":
+            return
+        # Check if IP Address
+        try:
+            ipaddress.ip_address(value)
+            return
+        except ValueError:
+            pass
+        # Check if IP Range
+        try:
+            ipaddress.ip_network(value)
+            return
+        except ValueError:
+            pass
+
+        self.is_tag = True
+
+
+def update_allowed_aad_tenants(
+    config_client: InstanceConfigClient, tenants: List[UUID]
+) -> None:
+    as_str = [str(x) for x in tenants]
+    config_client.table_service.insert_or_merge_entity(
+        TABLE_NAME,
+        {
+            "PartitionKey": config_client.resource_group,
+            "RowKey": config_client.resource_group,
+            "allowed_aad_tenants": json.dumps(as_str),
+        },
+    )
+
+
+def update_admins(config_client: InstanceConfigClient, admins: List[UUID]) -> None:
+    admins_as_str: Optional[List[str]] = None
+    if admins:
+        admins_as_str = [str(x) for x in admins]
+
+    config_client.table_service.insert_or_merge_entity(
+        TABLE_NAME,
+        {
+            "PartitionKey": config_client.resource_group,
+            "RowKey": config_client.resource_group,
+            "admins": json.dumps(admins_as_str),
+        },
+    )
+
+
+def parse_rules(rules_str: str) -> List[NsgRule]:
+    rules_list = rules_str.split(",")
+
+    nsg_rules = []
+    for rule in rules_list:
+        try:
+            nsg_rule = NsgRule(rule)
+            nsg_rules.append(nsg_rule)
+        except Exception:
+            raise ValueError(
+                "One or more input rules was invalid. Please enter a comma-separted list if valid sources."
+            )
+    return nsg_rules
+
+
+def update_nsg(
+    config_client: InstanceConfigClient,
+    allowed_rules: List[NsgRule],
+) -> None:
+    tags_as_str = [x.rule for x in allowed_rules if x.is_tag]
+    ips_as_str = [x.rule for x in allowed_rules if not x.is_tag]
+    nsg_config = {"allowed_service_tags": tags_as_str, "allowed_ips": ips_as_str}
+    # create class initialized by table service/resource group outside function that's checked in deploy.py
+    config_client.table_service.insert_or_merge_entity(
+        TABLE_NAME,
+        {
+            "PartitionKey": config_client.resource_group,
+            "RowKey": config_client.resource_group,
+            "proxy_nsg_config": json.dumps(nsg_config),
+        },
+    )
+
+
+if __name__ == "__main__":
+    pass

src/deployment/deploy.py

@@ -47,6 +47,13 @@
 )
 from msrest.serialization import TZ_UTC
 
+from configuration import (
+    InstanceConfigClient,
+    parse_rules,
+    update_admins,
+    update_allowed_aad_tenants,
+    update_nsg,
+)
 from data_migration import migrate
 from registration import (
     GraphQueryError,
@@ -61,7 +68,6 @@
     set_app_audience,
     update_pool_registration,
 )
-from set_admins import update_admins, update_allowed_aad_tenants
 
 # Found by manually assigning the User.Read permission to application
 # registration in the admin portal. The values are in the manifest under
@@ -104,6 +110,7 @@ def __init__(
         location: str,
         application_name: str,
         owner: str,
+        nsg_config: str,
         client_id: Optional[str],
         client_secret: Optional[str],
         app_zip: str,
@@ -127,6 +134,7 @@ def __init__(
         self.location = location
         self.application_name = application_name
         self.owner = owner
+        self.nsg_config = nsg_config
         self.app_zip = app_zip
         self.tools = tools
         self.instance_specific = instance_specific
@@ -608,13 +616,19 @@ def set_instance_config(self) -> None:
         tenant = UUID(self.results["deploy"]["tenant_id"]["value"])
         table_service = TableService(account_name=name, account_key=key)
 
+        config_client = InstanceConfigClient(table_service, self.application_name)
+
+        if self.nsg_config:
+            rules = parse_rules(self.nsg_config)
+            update_nsg(config_client, rules)
+
         if self.admins:
-            update_admins(table_service, self.application_name, self.admins)
+            update_admins(config_client, self.admins)
 
         tenants = self.allowed_aad_tenants
         if tenant not in tenants:
             tenants.append(tenant)
-        update_allowed_aad_tenants(table_service, self.application_name, tenants)
+        update_allowed_aad_tenants(config_client, tenants)
 
     def create_eventgrid(self) -> None:
         logger.info("creating eventgrid subscription")
@@ -972,6 +986,7 @@ def main() -> None:
     parser.add_argument("resource_group")
     parser.add_argument("application_name")
     parser.add_argument("owner")
+    parser.add_argument("nsg_config")
     parser.add_argument(
         "--arm-template",
         type=arg_file,
@@ -1079,6 +1094,7 @@ def main() -> None:
         location=args.location,
         application_name=args.application_name,
         owner=args.owner,
+        nsg_config=args.nsg_config,
         client_id=args.client_id,
         client_secret=args.client_secret,
         app_zip=args.app_zip,

src/deployment/requirements.txt

@@ -8,4 +8,4 @@ azure-storage-blob==12.8.1
 pyfunctional==1.4.3
 pyopenssl==19.1.0
 adal~=1.2.5
-idna<3,>=2.5
+idna<3,>=2.5