use managed identity reader access for scaleset configs (#1060)

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

@@ -11,7 +11,12 @@ from onefuzztypes.enums import OS, AgentMode
 from onefuzztypes.models import AgentConfig, Pool, ReproConfig, Scaleset
 from onefuzztypes.primitives import Container, Extension, Region
 
-from .azure.containers import get_container_sas_url, get_file_sas_url, save_blob
+from .azure.containers import (
+    get_container_sas_url,
+    get_file_sas_url,
+    get_file_url,
+    save_blob,
+)
 from .azure.creds import get_instance_id, get_instance_url
 from .azure.monitor import get_monitor_settings
 from .azure.queue import get_queue_sas
@@ -94,9 +99,7 @@ def build_scaleset_script(pool: Pool, scaleset: Scaleset) -> str:
     save_blob(
         Container("vm-scripts"), filename, sep.join(commands) + sep, StorageType.config
     )
-    return get_file_sas_url(
-        Container("vm-scripts"), filename, StorageType.config, read=True
-    )
+    return get_file_url(Container("vm-scripts"), filename, StorageType.config)
 
 
 def build_pool_config(pool: Pool) -> str:
@@ -126,12 +129,7 @@ def build_pool_config(pool: Pool) -> str:
         StorageType.config,
     )
 
-    return get_file_sas_url(
-        Container("vm-scripts"),
-        filename,
-        StorageType.config,
-        read=True,
-    )
+    return config_url(Container("vm-scripts"), filename, False)
 
 
 def update_managed_scripts() -> None:
@@ -167,8 +165,20 @@ def update_managed_scripts() -> None:
     )
 
 
+def config_url(container: Container, filename: str, with_sas: bool) -> str:
+    if with_sas:
+        return get_file_sas_url(container, filename, StorageType.config, read=True)
+    else:
+        return get_file_url(container, filename, StorageType.config)
+
+
 def agent_config(
-    region: Region, vm_os: OS, mode: AgentMode, *, urls: Optional[List[str]] = None
+    region: Region,
+    vm_os: OS,
+    mode: AgentMode,
+    *,
+    urls: Optional[List[str]] = None,
+    with_sas: bool = False,
 ) -> Extension:
     update_managed_scripts()
 
@@ -177,29 +187,17 @@ def agent_config(
 
     if vm_os == OS.windows:
         urls += [
-            get_file_sas_url(
-                Container("vm-scripts"),
-                "managed.ps1",
-                StorageType.config,
-                read=True,
-            ),
-            get_file_sas_url(
-                Container("tools"),
-                "win64/azcopy.exe",
-                StorageType.config,
-                read=True,
-            ),
-            get_file_sas_url(
+            config_url(Container("vm-scripts"), "managed.ps1", with_sas),
+            config_url(Container("tools"), "win64/azcopy.exe", with_sas),
+            config_url(
                 Container("tools"),
                 "win64/setup.ps1",
-                StorageType.config,
-                read=True,
+                with_sas,
             ),
-            get_file_sas_url(
+            config_url(
                 Container("tools"),
                 "win64/onefuzz.ps1",
-                StorageType.config,
-                read=True,
+                with_sas,
             ),
         ]
         to_execute_cmd = (
@@ -214,29 +212,31 @@ def agent_config(
             "force_update_tag": uuid4(),
             "type_handler_version": "1.9",
             "auto_upgrade_minor_version": True,
-            "settings": {"commandToExecute": to_execute_cmd, "fileUris": urls},
-            "protectedSettings": {},
+            "settings": {
+                "commandToExecute": to_execute_cmd,
+                "fileUris": urls,
+            },
+            "protectedSettings": {
+                "managedIdentity": {},
+            },
         }
         return extension
     elif vm_os == OS.linux:
         urls += [
-            get_file_sas_url(
+            config_url(
                 Container("vm-scripts"),
                 "managed.sh",
-                StorageType.config,
-                read=True,
+                with_sas,
             ),
-            get_file_sas_url(
+            config_url(
                 Container("tools"),
                 "linux/azcopy",
-                StorageType.config,
-                read=True,
+                with_sas,
             ),
-            get_file_sas_url(
+            config_url(
                 Container("tools"),
                 "linux/setup.sh",
-                StorageType.config,
-                read=True,
+                with_sas,
             ),
         ]
         to_execute_cmd = "sh setup.sh %s" % (mode.name)
@@ -247,10 +247,15 @@ def agent_config(
             "type": "CustomScript",
             "typeHandlerVersion": "2.1",
             "location": region,
-            "autoUpgradeMinorVersion": True,
             "force_update_tag": uuid4(),
-            "settings": {"commandToExecute": to_execute_cmd, "fileUris": urls},
-            "protectedSettings": {},
+            "autoUpgradeMinorVersion": True,
+            "settings": {
+                "commandToExecute": to_execute_cmd,
+                "fileUris": urls,
+            },
+            "protectedSettings": {
+                "managedIdentity": {},
+            },
         }
         return extension
 
@@ -337,7 +342,9 @@ def repro_extensions(
             ),
         ]
 
-    base_extension = agent_config(region, repro_os, AgentMode.repro, urls=urls)
+    base_extension = agent_config(
+        region, repro_os, AgentMode.repro, urls=urls, with_sas=True
+    )
     extensions = generic_extensions(region, repro_os)
     extensions += [base_extension]
     return extensions
@@ -359,7 +366,9 @@ def proxy_manager_extensions(region: Region, proxy_id: UUID) -> List[Extension]:
         ),
     ]
 
-    base_extension = agent_config(region, OS.linux, AgentMode.proxy, urls=urls)
+    base_extension = agent_config(
+        region, OS.linux, AgentMode.proxy, urls=urls, with_sas=True
+    )
     extensions = generic_extensions(region, OS.linux)
     extensions += [base_extension]
     return extensions