Enable User assigned managed identity for scalesets (#219)

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

@@ -14,6 +14,7 @@ from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error
 from pydantic import BaseModel
 
+from .azure.creds import get_scaleset_principal_id
 from .pools import Scaleset
 from .request import not_ok
 
@@ -54,6 +55,12 @@ def try_get_token_auth_header(request: func.HttpRequest) -> Union[Error, TokenDa
 
 @cached(ttl=60)
 def is_authorized(token_data: TokenData) -> bool:
+    # verify object_id against the user assigned managed identity
+    if get_scaleset_principal_id() == token_data.object_id:
+        return True
+
+    # backward compatibility case for scalesets deployed before the migration
+    # to user assigned managed id
     scalesets = Scaleset.get_by_object_id(token_data.object_id)
     return len(scalesets) > 0
 

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

@@ -6,6 +6,7 @@
 import logging
 import os
 from typing import Any, List, Optional, Tuple
+from uuid import UUID
 
 from azure.cli.core import CLIError
 from azure.common.client_factory import get_client_from_cli_profile
@@ -123,3 +124,24 @@ def is_member_of(group_id: str, member_id: str) -> bool:
             CheckGroupMembershipParameters(group_id=group_id, member_id=member_id)
         ).value
     )
+
+
+@cached
+def get_scaleset_identity_resource_path() -> str:
+    scaleset_id_name = "%s-scalesetid" % get_instance_name()
+    resource_group_path = "/subscriptions/%s/resourceGroups/%s/providers" % (
+        get_subscription(),
+        get_base_resource_group(),
+    )
+    return "%s/Microsoft.ManagedIdentity/userAssignedIdentities/%s" % (
+        resource_group_path,
+        scaleset_id_name,
+    )
+
+
+@cached
+def get_scaleset_principal_id() -> UUID:
+    api_version = "2018-11-30"  # matches the apiversion in the deplyoment template
+    client = mgmt_client_factory(ResourceManagementClient)
+    uid = client.resources.get_by_id(get_scaleset_identity_resource_path(), api_version)
+    return UUID(uid.properties["principalId"])

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

@@ -16,7 +16,11 @@ from onefuzztypes.enums import OS, ErrorCode
 from onefuzztypes.models import Error
 from onefuzztypes.primitives import Region
 
-from .creds import get_base_resource_group, mgmt_client_factory
+from .creds import (
+    get_base_resource_group,
+    get_scaleset_identity_resource_path,
+    mgmt_client_factory,
+)
 from .image import get_os
 
 
@@ -234,7 +238,10 @@ def create_vmss(
         "do_not_run_extensions_on_overprovisioned_vms": True,
         "upgrade_policy": {"mode": "Manual"},
         "sku": sku,
-        "identity": {"type": "SystemAssigned"},
+        "identity": {
+            "type": "userAssigned",
+            "userAssignedIdentities": {get_scaleset_identity_resource_path(): {}},
+        },
         "virtual_machine_profile": {
             "priority": "Regular",
             "storage_profile": {"image_reference": image_ref},

src/deployment/azuredeploy.json

@@ -38,6 +38,7 @@
         }
     },
     "variables": {
+        "scaleset_identity": "[concat(parameters('name'), '-scalesetid')]",
         "telemetry": "d7a73cf4-5a1a-4030-85e1-e5b25867e45a",
         "signalr-name": "[concat('onefuzz-', uniquestring(resourceGroup().id))]",
         "monitorAccountName": "[concat('logs-wb-', uniquestring(resourceGroup().id))]",
@@ -47,6 +48,7 @@
         "Storage Account Contributor": "17d1049b-9a84-46fb-8f53-869881c3d3ab",
         "Virtual Machine Contributor": "9980e02c-c2be-4d73-94e8-173b1dc7cf3c",
         "Log Analytics Contributor": "92aaf0da-9dab-42b6-94a3-d43ce8d16293",
+        "Managed Identity Operator": "f1a07417-d97a-45cb-824c-7a7467783830",
         "storage_account_sas": {
             "signedServices": "bfqt",
             "signedPermission": "rwdlacup",
@@ -93,6 +95,12 @@
         }
     ],
     "resources": [
+        {
+            "type": "Microsoft.ManagedIdentity/userAssignedIdentities",
+            "name": "[variables('scaleset_identity')]",
+            "apiVersion": "2018-11-30",
+            "location": "[resourceGroup().location]"
+        },
         {
             "apiVersion": "2018-11-01",
             "name": "[parameters('name')]",
@@ -580,6 +588,21 @@
                 "OWNER": "[parameters('owner')]"
             }
         },
+        {
+            "type": "Microsoft.Authorization/roleAssignments",
+            "apiVersion": "2017-09-01",
+            "name": "[guid(concat(resourceGroup().id, '-user_managed_idenity'))]",
+            "properties": {
+                "roleDefinitionId": "[concat('/subscriptions/', subscription().subscriptionId, '/providers/Microsoft.Authorization/roleDefinitions/', variables('Managed Identity Operator'))]",
+                "principalId": "[reference(resourceId('Microsoft.Web/sites', parameters('name')), '2018-02-01', 'Full').identity.principalId]"
+            },
+            "DependsOn": [
+                "[resourceId('Microsoft.Web/sites', parameters('name'))]"
+            ],
+            "tags": {
+                "OWNER": "[parameters('owner')]"
+            }
+        },
         {
             "type": "Microsoft.SignalRService/SignalR",
             "apiVersion": "2018-10-01",
@@ -640,6 +663,10 @@
         "func-key": {
             "type": "string",
             "value": "[listKeys(resourceId('Microsoft.Storage/storageAccounts', variables('storageAccountNameFunc')), '2019-06-01').keys[0].value]"
+        },
+        "scaleset-identity": {
+            "type": "string",
+            "value": "[variables('scaleset_identity')]"
         }
     }
 }

src/deployment/deploy.py

@@ -62,6 +62,7 @@ from msrest.serialization import TZ_UTC
 from data_migration import migrate
 from registration import (
     add_application_password,
+    assign_scaleset_role,
     authorize_application,
     get_application,
     register_application,
@@ -395,6 +396,11 @@ class Client:
             sys.exit(1)
         self.results["deploy"] = result.properties.outputs
 
+        logger.info("assigning the user managed identity role")
+        assign_scaleset_role(
+            self.application_name, self.results["deploy"]["scaleset-identity"]["value"]
+        )
+
     def apply_migrations(self):
         self.results["deploy"]["func-storage"]["value"]
         name = self.results["deploy"]["func-name"]["value"]

src/deployment/registration.py

@@ -298,28 +298,39 @@ def assign_scaleset_role(onefuzz_instance_name: str, scaleset_name: str):
         raise Exception("scaleset service principal not found")
     scaleset_service_principal = scaleset_service_principals["value"][0]
 
-    lab_machine_role = (
+    managed_node_role = (
         seq(onefuzz_service_principal["appRoles"])
         .filter(lambda x: x["value"] == "ManagedNode")
         .head_option()
     )
 
-    if not lab_machine_role:
+    if not managed_node_role:
         raise Exception(
             "ManagedNode role not found int the onefuzz application registration. Please redeploy the instance"
         )
 
-    query_microsoft_graph(
+    assignments = query_microsoft_graph(
-        method="POST",
+        method="GET",
-        resource="servicePrincipals/%s/appRoleAssignedTo"
+        resource="servicePrincipals/%s/appRoleAssignments"
         % scaleset_service_principal["id"],
-        body={
-            "principalId": scaleset_service_principal["id"],
-            "resourceId": onefuzz_service_principal["id"],
-            "appRoleId": lab_machine_role["id"],
-        },
     )
 
+    # check if the role is already assigned
+    role_assigned = seq(assignments["value"]).find(
+        lambda assignment: assignment["appRoleId"] == managed_node_role["id"]
+    )
+    if not role_assigned:
+        query_microsoft_graph(
+            method="POST",
+            resource="servicePrincipals/%s/appRoleAssignedTo"
+            % scaleset_service_principal["id"],
+            body={
+                "principalId": scaleset_service_principal["id"],
+                "resourceId": onefuzz_service_principal["id"],
+                "appRoleId": managed_node_role["id"],
+            },
+        )
+
 
 def main():
     formatter = argparse.ArgumentDefaultsHelpFormatter