Add ssh keys to nodes on demand (#411)

Our existing model has a per-scaleset SSH key.  This update moves towards using user provided SSH keys when they need to connect to a given node.

src/agent/Cargo.lock

@@ -1649,6 +1649,7 @@ dependencies = [
  "structopt",
  "tokio",
  "url",
+ "users",
  "uuid",
 ]
 
@@ -2778,6 +2779,16 @@ version = "0.7.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "110352d4e9076c67839003c7788d8604e24dcded13e0b375af3efaa8cf468517"
 
+[[package]]
+name = "users"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24cc0f6d6f267b73e5a2cadf007ba8f9bc39c6a6f9666f8cf25ea809a153b032"
+dependencies = [
+ "libc",
+ "log",
+]
+
 [[package]]
 name = "uuid"
 version = "0.8.1"

src/agent/onefuzz-supervisor/Cargo.toml

@@ -25,3 +25,6 @@ url = { version = "2.1.1", features = ["serde"] }
 uuid = { version = "0.8.1", features = ["serde", "v4"] }
 clap = "2.33"
 reqwest-retry = { path = "../reqwest-retry" }
+
+[target.'cfg(target_os = "linux")'.dependencies]
+users = "0.11"

src/agent/onefuzz-supervisor/src/agent.rs

@@ -264,7 +264,7 @@ impl Agent {
 
         if let Some(cmd) = cmd {
             verbose!("agent received node command: {:?}", cmd);
-            self.scheduler()?.execute_command(cmd)?;
+            self.scheduler()?.execute_command(cmd).await?;
         }
 
         Ok(())

src/agent/onefuzz-supervisor/src/commands.rs

@@ -0,0 +1,176 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::auth::Secret;
+use anyhow::Result;
+use onefuzz::machine_id::get_scaleset_name;
+use std::process::Stdio;
+use tokio::{fs, io::AsyncWriteExt, process::Command};
+
+#[cfg(target_os = "windows")]
+use std::{env, path::PathBuf};
+
+#[cfg(target_os = "linux")]
+use users::{get_user_by_name, os::unix::UserExt};
+
+#[cfg(target_os = "linux")]
+const ONEFUZZ_SERVICE_USER: &str = "onefuzz";
+
+#[derive(Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub struct SshKeyInfo {
+    pub public_key: Secret<String>,
+}
+
+#[cfg(target_os = "windows")]
+pub async fn add_ssh_key(key_info: SshKeyInfo) -> Result<()> {
+    if get_scaleset_name().await?.is_none() {
+        warn!("adding ssh keys only supported on managed nodes");
+        return Ok(());
+    }
+
+    let mut ssh_path =
+        PathBuf::from(env::var("ProgramData").unwrap_or_else(|_| "c:\\programdata".to_string()));
+    ssh_path.push("ssh");
+
+    let host_key_path = ssh_path.join("ssh_host_dsa_key");
+    let admin_auth_keys_path = ssh_path.join("administrators_authorized_keys");
+
+    {
+        let mut file = fs::OpenOptions::new()
+            .append(true)
+            .open(&admin_auth_keys_path)
+            .await?;
+        file.write_all(key_info.public_key.expose_ref().as_bytes())
+            .await?;
+    }
+
+    verbose!("removing Authenticated Users permissions from administrators_authorized_keys");
+    let result = Command::new("icacls.exe")
+        .arg(&admin_auth_keys_path)
+        .arg("/remove")
+        .arg("NT AUTHORITY/Authenticated Users")
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    if !result.status.success() {
+        bail!(
+            "set authorized_keys ({}) permissions failed: {:?}",
+            admin_auth_keys_path.display(),
+            result
+        );
+    }
+
+    verbose!("removing inheritance");
+    let result = Command::new("icacls.exe")
+        .arg(&admin_auth_keys_path)
+        .arg("/inheritance:r")
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    if !result.status.success() {
+        bail!(
+            "set authorized_keys ({}) permissions failed: {:?}",
+            admin_auth_keys_path.display(),
+            result
+        );
+    }
+
+    verbose!("copying ACL from ssh_host_dsa_key");
+    let result = Command::new("powershell.exe")
+        .args(&["-ExecutionPolicy", "Unrestricted", "-Command"])
+        .arg(format!(
+            "Get-Acl \"{}\" | Set-Acl \"{}\"",
+            admin_auth_keys_path.display(),
+            host_key_path.display()
+        ))
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    if !result.status.success() {
+        bail!(
+            "set authorized_keys ({}) permissions failed: {:?}",
+            admin_auth_keys_path.display(),
+            result
+        );
+    }
+
+    info!("ssh key written: {}", admin_auth_keys_path.display());
+
+    Ok(())
+}
+
+#[cfg(target_os = "linux")]
+pub async fn add_ssh_key(key_info: SshKeyInfo) -> Result<()> {
+    if get_scaleset_name().await?.is_none() {
+        warn!("adding ssh keys only supported on managed nodes");
+        return Ok(());
+    }
+
+    let user =
+        get_user_by_name(ONEFUZZ_SERVICE_USER).ok_or_else(|| format_err!("unable to find user"))?;
+    info!("adding sshkey:{:?} to user:{:?}", key_info, user);
+
+    let home_path = user.home_dir().to_owned();
+    if !home_path.exists() {
+        bail!("unable to add SSH key to missing home directory");
+    }
+
+    let mut ssh_path = home_path.join(".ssh");
+    if !ssh_path.exists() {
+        verbose!("creating ssh directory: {}", ssh_path.display());
+        fs::create_dir_all(&ssh_path).await?;
+    }
+
+    verbose!("setting ssh permissions");
+    let result = Command::new("chmod")
+        .arg("700")
+        .arg(&ssh_path)
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    if !result.status.success() {
+        bail!("set $HOME/.ssh permissions failed: {:?}", result);
+    }
+
+    ssh_path.push("authorized_keys");
+
+    {
+        let mut file = fs::OpenOptions::new().append(true).open(&ssh_path).await?;
+        file.write_all(key_info.public_key.expose_ref().as_bytes())
+            .await?;
+    }
+
+    verbose!("setting authorized_keys permissions");
+    let result = Command::new("chmod")
+        .arg("600")
+        .arg(&ssh_path)
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    if !result.status.success() {
+        bail!(
+            "set authorized_keys ({}) permissions failed: {:?}",
+            ssh_path.display(),
+            result
+        );
+    }
+
+    info!("ssh key written: {}", ssh_path.display());
+
+    Ok(())
+}

src/agent/onefuzz-supervisor/src/coordinator.rs

@@ -9,29 +9,31 @@ use serde::Serialize;
 use uuid::Uuid;
 
 use crate::auth::AccessToken;
+use crate::commands::SshKeyInfo;
 use crate::config::Registration;
 use crate::work::{TaskId, WorkSet};
 use crate::worker::WorkerEvent;
 
-#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[derive(Debug, Deserialize, Eq, PartialEq, Serialize)]
 pub struct StopTask {
     pub task_id: TaskId,
 }
 
-#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[derive(Debug, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum NodeCommand {
+    AddSshKey(SshKeyInfo),
     StopTask(StopTask),
     Stop {},
 }
 
-#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[derive(Debug, Deserialize, Eq, PartialEq, Serialize)]
 pub struct NodeCommandEnvelope {
     pub message_id: String,
     pub command: NodeCommand,
 }
 
-#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[derive(Debug, Deserialize, Eq, PartialEq, Serialize)]
 pub struct PendingNodeCommand {
     envelope: Option<NodeCommandEnvelope>,
 }

src/agent/onefuzz-supervisor/src/coordinator/double.rs

@@ -3,7 +3,7 @@
 
 use super::*;
 
-#[derive(Clone, Debug, Default)]
+#[derive(Debug, Default)]
 pub struct CoordinatorDouble {
     pub commands: Vec<NodeCommand>,
     pub events: Vec<NodeEvent>,

src/agent/onefuzz-supervisor/src/main.rs

@@ -11,6 +11,8 @@ extern crate onefuzz;
 extern crate serde;
 #[macro_use]
 extern crate clap;
+#[macro_use]
+extern crate anyhow;
 
 use crate::{
     config::StaticConfig, coordinator::StateUpdateEvent, heartbeat::*, work::WorkSet,
@@ -28,6 +30,7 @@ use structopt::StructOpt;
 
 pub mod agent;
 pub mod auth;
+pub mod commands;
 pub mod config;
 pub mod coordinator;
 pub mod debug;

src/agent/onefuzz-supervisor/src/scheduler.rs

@@ -6,6 +6,7 @@ use std::fmt;
 use anyhow::Result;
 use onefuzz::process::Output;
 
+use crate::commands::add_ssh_key;
 use crate::coordinator::{NodeCommand, NodeState};
 use crate::reboot::RebootContext;
 use crate::setup::ISetupRunner;
@@ -53,8 +54,11 @@ impl Scheduler {
         Self::default()
     }
 
-    pub fn execute_command(&mut self, cmd: NodeCommand) -> Result<()> {
+    pub async fn execute_command(&mut self, cmd: NodeCommand) -> Result<()> {
         match cmd {
+            NodeCommand::AddSshKey(ssh_key_info) => {
+                add_ssh_key(ssh_key_info).await?;
+            }
             NodeCommand::StopTask(stop_task) => {
                 if let Scheduler::Busy(state) = self {
                     state.stop(stop_task.task_id)?;

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

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import azure.functions as func
+from onefuzztypes.enums import ErrorCode
+from onefuzztypes.models import Error
+from onefuzztypes.requests import NodeAddSshKey
+from onefuzztypes.responses import BoolResult
+
+from ..onefuzzlib.pools import Node
+from ..onefuzzlib.request import not_ok, ok, parse_request
+
+
+def post(req: func.HttpRequest) -> func.HttpResponse:
+    request = parse_request(NodeAddSshKey, req)
+    if isinstance(request, Error):
+        return not_ok(request, context="NodeAddSshKey")
+
+    node = Node.get_by_machine_id(request.machine_id)
+    if not node:
+        return not_ok(
+            Error(code=ErrorCode.UNABLE_TO_FIND, errors=["unable to find node"]),
+            context=request.machine_id,
+        )
+    result = node.add_ssh_public_key(public_key=request.public_key)
+    if isinstance(result, Error):
+        return not_ok(result, context="NodeAddSshKey")
+
+    return ok(BoolResult(result=True))
+
+
+def main(req: func.HttpRequest) -> func.HttpResponse:
+    if req.method == "POST":
+        return post(req)
+    else:
+        raise Exception("invalid method")

src/api-service/__app__/node_add_ssh_key/function.json

@@ -0,0 +1,20 @@
+{
+    "scriptFile": "__init__.py",
+    "bindings": [
+        {
+            "authLevel": "anonymous",
+            "type": "httpTrigger",
+            "direction": "in",
+            "name": "req",
+            "methods": [
+                "post"
+            ],
+            "route": "node/add_ssh_key"
+        },
+        {
+            "type": "http",
+            "direction": "out",
+            "name": "$return"
+        }
+    ]
+}

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

@@ -18,9 +18,10 @@ from onefuzztypes.enums import (
 )
 from onefuzztypes.models import AutoScaleConfig, Error
 from onefuzztypes.models import Node as BASE_NODE
-from onefuzztypes.models import NodeAssignment, NodeCommand
+from onefuzztypes.models import NodeAssignment, NodeCommand, NodeCommandAddSshKey
 from onefuzztypes.models import NodeTasks as BASE_NODE_TASK
 from onefuzztypes.models import Pool as BASE_POOL
+from onefuzztypes.models import Result
 from onefuzztypes.models import Scaleset as BASE_SCALESET
 from onefuzztypes.models import (
     ScalesetNodeState,
@@ -257,11 +258,10 @@ class Node(BASE_NODE, ORMMixin):
         return self.version != __version__
 
     def send_message(self, message: NodeCommand) -> None:
-        stop_message = NodeMessage(
+        NodeMessage(
             agent_id=self.machine_id,
             message=message,
-        )
+        ).save()
-        stop_message.save()
 
     def to_reimage(self, done: bool = False) -> None:
         if done:
@@ -273,6 +273,21 @@ class Node(BASE_NODE, ORMMixin):
             self.reimage_requested = True
         self.save()
 
+    def add_ssh_public_key(self, public_key: str) -> Result[None]:
+        if self.scaleset_id is None:
+            return Error(
+                code=ErrorCode.INVALID_REQUEST,
+                errors=["only able to add ssh keys to scaleset nodes"],
+            )
+
+        if not public_key.endswith("\n"):
+            public_key += "\n"
+
+        self.send_message(
+            NodeCommand(add_ssh_key=NodeCommandAddSshKey(public_key=public_key))
+        )
+        return None
+
     def stop(self) -> None:
         self.to_reimage()
         self.send_message(NodeCommand(stop=StopNodeCommand()))

src/cli/onefuzz/api.py

@@ -1187,6 +1187,26 @@ class Node(Endpoint):
             ),
         )
 
+    def add_ssh_key(
+        self, machine_id: UUID_EXPANSION, *, public_key: str
+    ) -> responses.BoolResult:
+        self.logger.debug("add ssh public key to node: %s", machine_id)
+        machine_id_expanded = self._disambiguate_uuid(
+            "machine_id",
+            machine_id,
+            lambda: [str(x.machine_id) for x in self.list()],
+        )
+
+        return self._req_model(
+            "POST",
+            responses.BoolResult,
+            data=requests.NodeAddSshKey(
+                machine_id=machine_id_expanded,
+                public_key=public_key,
+            ),
+            alternate_endpoint="node/add_ssh_key",
+        )
+
 
 class Scaleset(Endpoint):
     """ Interact with managed scaleset pools """

src/pytypes/onefuzztypes/models.py

@@ -680,9 +680,14 @@ class StopTaskNodeCommand(BaseModel):
     task_id: UUID
 
 
+class NodeCommandAddSshKey(BaseModel):
+    public_key: str
+
+
 class NodeCommand(EnumModel):
     stop: Optional[StopNodeCommand]
     stop_task: Optional[StopTaskNodeCommand]
+    add_ssh_key: Optional[NodeCommandAddSshKey]
 
 
 class NodeCommandEnvelope(BaseModel):

src/pytypes/onefuzztypes/requests.py

@@ -233,3 +233,8 @@ class WebhookUpdate(BaseModel):
     event_types: Optional[List[WebhookEventType]]
     url: Optional[AnyHttpUrl]
     secret_token: Optional[str]
+
+
+class NodeAddSshKey(BaseModel):
+    machine_id: UUID
+    public_key: str