Re-add windows ssh key (#390)

Adds a scaleset specific setup script, which allows us to save the scaleset based SSH keys into the VM on setup.

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

@@ -8,7 +8,7 @@ from typing import List, Optional
 from uuid import UUID
 
 from onefuzztypes.enums import OS, AgentMode
-from onefuzztypes.models import AgentConfig, ReproConfig
+from onefuzztypes.models import AgentConfig, Pool, ReproConfig, Scaleset
 from onefuzztypes.primitives import Extension, Region
 
 from .azure.containers import (
@@ -22,16 +22,6 @@ from .azure.monitor import get_monitor_settings
 from .azure.queue import get_queue_sas
 from .reports import get_report
 
-# TODO: figure out how to create VM specific SSH keys for Windows.
-#
-# Previously done via task specific scripts:
-
-# if is_windows and auth is not None:
-#     ssh_key = auth.public_key.strip()
-#     ssh_path = "$env:ProgramData/ssh/administrators_authorized_keys"
-#     commands += ['Set-Content -Path %s -Value "%s"' % (ssh_path, ssh_key)]
-#     return commands
-
 
 def generic_extensions(region: Region, vm_os: OS) -> List[Extension]:
     extensions = [monitor_extension(region, vm_os)]
@@ -94,9 +84,24 @@ def dependency_extension(region: Region, vm_os: OS) -> Optional[Extension]:
         return None
 
 
-def build_pool_config(pool_name: str) -> str:
+def build_scaleset_script(pool: Pool, scaleset: Scaleset) -> str:
+    commands = []
+    extension = "ps1" if pool.os == OS.windows else "sh"
+    filename = f"{scaleset.scaleset_id}/scaleset-setup.{extension}"
+    sep = "\r\n" if pool.os == OS.windows else "\n"
+
+    if pool.os == OS.windows and scaleset.auth is not None:
+        ssh_key = scaleset.auth.public_key.strip()
+        ssh_path = "$env:ProgramData/ssh/administrators_authorized_keys"
+        commands += [f'Set-Content -Path {ssh_path} -Value "{ssh_key}"']
+
+    save_blob("vm-scripts", filename, sep.join(commands) + sep, StorageType.config)
+    return get_file_sas_url("vm-scripts", filename, StorageType.config, read=True)
+
+
+def build_pool_config(pool: Pool) -> str:
     config = AgentConfig(
-        pool_name=pool_name,
+        pool_name=pool.name,
         onefuzz_url=get_instance_url(),
         instrumentation_key=os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY"),
         heartbeat_queue=get_queue_sas(
@@ -108,16 +113,18 @@ def build_pool_config(pool_name: str) -> str:
         instance_id=get_instance_id(),
     )
 
+    filename = f"{pool.name}/config.json"
+
     save_blob(
         "vm-scripts",
-        "%s/config.json" % pool_name,
+        filename,
         config.json(),
         StorageType.config,
     )
 
     return get_file_sas_url(
         "vm-scripts",
-        "%s/config.json" % pool_name,
+        filename,
         StorageType.config,
         read=True,
     )
@@ -240,10 +247,10 @@ def agent_config(
     raise NotImplementedError("unsupported OS: %s" % vm_os)
 
 
-def fuzz_extensions(region: Region, vm_os: OS, pool_name: str) -> List[Extension]:
+def fuzz_extensions(pool: Pool, scaleset: Scaleset) -> List[Extension]:
-    urls = [build_pool_config(pool_name)]
+    urls = [build_pool_config(pool), build_scaleset_script(pool, scaleset)]
-    fuzz_extension = agent_config(region, vm_os, AgentMode.fuzz, urls=urls)
+    fuzz_extension = agent_config(scaleset.region, pool.os, AgentMode.fuzz, urls=urls)
-    extensions = generic_extensions(region, vm_os)
+    extensions = generic_extensions(scaleset.region, pool.os)
     extensions += [fuzz_extension]
     return extensions
 

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

@@ -701,7 +701,7 @@ class Scaleset(BASE_SCALESET, ORMMixin):
                 return
 
             logging.info("creating scaleset: %s", self.scaleset_id)
-            extensions = fuzz_extensions(self.region, pool.os, self.pool_name)
+            extensions = fuzz_extensions(pool, self)
             result = create_vmss(
                 self.region,
                 self.scaleset_id,
@@ -1034,7 +1034,7 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             return
 
         logging.debug("updating scaleset configs: %s", self.scaleset_id)
-        extensions = fuzz_extensions(self.region, pool.os, self.pool_name)
+        extensions = fuzz_extensions(pool, self)
         try:
             update_extensions(self.scaleset_id, extensions)
         except UnableToUpdate:

src/runtime-tools/linux/setup.sh

@@ -11,6 +11,7 @@ INSTANCE_SETUP="/onefuzz/instance-specific-setup/setup.sh"
 USER_SETUP="/onefuzz/setup/setup.sh"
 TASK_SETUP="/onefuzz/bin/task-setup.sh"
 MANAGED_SETUP="/onefuzz/bin/managed.sh"
+SCALESET_SETUP="/onefuzz/bin/scaleset-setup.sh"
 export ONEFUZZ_ROOT=/onefuzz
 export ASAN_SYMBOLIZER_PATH=/onefuzz/bin/llvm-symbolizer
 
@@ -44,6 +45,10 @@ fi
 if [ -f /onefuzz/downloaded/repro-stdout.sh ]; then
     mv /onefuzz/downloaded/repro-stdout.sh /onefuzz/bin/
 fi
+if [ -f /onefuzz/downloaded/scaleset-setup.sh ]; then
+    mv /onefuzz/downloaded/scaleset-setup.sh /onefuzz/bin
+fi
+
 chmod -R a+rx /onefuzz/bin
 
 if [ -f ${MANAGED_SETUP} ]; then
@@ -55,6 +60,15 @@ else
     logger "onefuzz: no managed setup script"
 fi
 
+if [ -f ${SCALESET_SETUP} ]; then
+    logger "onefuzz: scaleset setup script start"
+    chmod +x ${SCALESET_SETUP}
+    ${SCALESET_SETUP} 2>&1 | logger -s -i -t 'onefuzz-scaleset-setup'
+    logger "onefuzz: scaleset setup script stop"
+else
+    logger "onefuzz: no scaleset setup script"
+fi
+
 if [ -f ${INSTANCE_SETUP} ]; then
     logger "onefuzz: instance setup script start"
     chmod +x ${INSTANCE_SETUP}

src/runtime-tools/win64/setup.ps1

@@ -41,6 +41,10 @@ function Install-OnefuzzSetup {
     log "onefuzz: executing managed-setup"
     ./managed.ps1
   }
+  if (Test-Path -Path scaleset-setup.ps1) {
+    log "onefuzz: executing scaleset-setup"
+    ./scaleset-setup.ps1
+  }
   if (Test-Path -Path task-setup.ps1) {
     log "onefuzz: executing task-setup"
     ./task-setup.ps1