Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Commit

Permalink
Re-add windows ssh key (#390)
Browse files Browse the repository at this point in the history
Adds a scaleset specific setup script, which allows us to save the scaleset based SSH keys into the VM on setup.
  • Loading branch information
bmc-msft authored Jan 4, 2021
1 parent 3441790 commit 4c2679d
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 21 deletions.
45 changes: 26 additions & 19 deletions src/api-service/__app__/onefuzzlib/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
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 (
Expand All @@ -22,16 +22,6 @@
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)]
Expand Down Expand Up @@ -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(
Expand All @@ -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,
)
Expand Down Expand Up @@ -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]:
urls = [build_pool_config(pool_name)]
fuzz_extension = agent_config(region, vm_os, AgentMode.fuzz, urls=urls)
extensions = generic_extensions(region, vm_os)
def fuzz_extensions(pool: Pool, scaleset: Scaleset) -> List[Extension]:
urls = [build_pool_config(pool), build_scaleset_script(pool, scaleset)]
fuzz_extension = agent_config(scaleset.region, pool.os, AgentMode.fuzz, urls=urls)
extensions = generic_extensions(scaleset.region, pool.os)
extensions += [fuzz_extension]
return extensions

Expand Down
4 changes: 2 additions & 2 deletions src/api-service/__app__/onefuzzlib/pools.py
Original file line number Diff line number Diff line change
Expand Up @@ -701,7 +701,7 @@ def setup(self) -> None:
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,
Expand Down Expand Up @@ -1034,7 +1034,7 @@ def update_configs(self) -> None:
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:
Expand Down
14 changes: 14 additions & 0 deletions src/runtime-tools/linux/setup.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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}
Expand Down
4 changes: 4 additions & 0 deletions src/runtime-tools/win64/setup.ps1
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4c2679d

Please sign in to comment.