✨ shell

src/api.rs

@@ -40,6 +40,19 @@ pub struct JoinResponse {
     pub mqtt: String,
 }
 
+#[derive(Deserialize, Serialize)]
+pub struct ShellParam {
+    pub cmd: String,
+    pub cwd: String,
+}
+
+#[derive(Deserialize, Serialize, Debug)]
+pub struct ShellResponse {
+    pub stdout: String,
+    pub stderr: String,
+    pub status: i32,
+}
+
 /// Setup a MQTT connection for `machine_id` on `mqtt`.
 ///
 /// This will connect either over `ws://` or `wss://` depending on the scheme of `mqtt`. By default it will use `wss://`.
@@ -175,6 +188,7 @@ impl ClientAction {
 #[derive(Debug, Serialize, Deserialize)]
 pub enum ClientActions {
     OSQuery(String),
+    Shell(String, String),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -195,4 +209,5 @@ impl ServerResponse {
 #[derive(Debug, Serialize, Deserialize)]
 pub enum ServerResponses {
     OSQuery(String),
+    Shell(ShellResponse),
 }

src/herd.rs

@@ -1,3 +1,4 @@
+use api::ServerResponse;
 use axum::{
     Router,
     routing::{get, post},
@@ -13,7 +14,7 @@ mod herd_core;
 use crate::herd_core::mqtt::{handle_mqtt, listen_to_devices};
 use herd_core::{
     config::Config,
-    route::{join_device, login_user},
+    route::{device_get_api, device_shell_cmd, join_device, login_user},
 };
 use herd_core::{model::Machine, route::devices_list};
 use sage::Identity;
@@ -21,8 +22,11 @@ use tokio::sync::OnceCell;
 
 pub static IDENTITY: OnceCell<Identity> = OnceCell::const_new();
 pub static CONFIG: OnceCell<Config> = OnceCell::const_new();
+pub static MQTT: OnceCell<AsyncClient> = OnceCell::const_new();
 
 pub static ONLINE: OnceCell<DashMap<String, chrono::DateTime<chrono::Utc>>> = OnceCell::const_new();
+pub static DISPATCH: OnceCell<DashMap<String, crossbeam::channel::Sender<ServerResponse>>> =
+    OnceCell::const_new();
 
 fn generate_token() -> String {
     let mut rng = rand::rng();
@@ -48,6 +52,7 @@ async fn main() {
     let _ = crate::CONFIG.set(config);
 
     crate::ONLINE.set(DashMap::new()).unwrap();
+    crate::DISPATCH.set(DashMap::new()).unwrap();
 
     let db = Database::filesystem("./herd/db");
     set_global_db!(db);
@@ -61,6 +66,8 @@ async fn main() {
 
     let user = Router::new()
         .route("/login", post(login_user))
+        .route("/device/{device_id}", get(device_get_api))
+        .route("/device/{device_id}/shell", post(device_shell_cmd))
         .route("/devices", get(devices_list));
 
     let app = Router::new().merge(device).merge(user);
@@ -70,6 +77,7 @@ async fn main() {
     let (client, eventloop) = api::mqtt_connect("server", &crate::CONFIG.get().unwrap().mqtt);
 
     listen_to_devices(&client).await;
+    crate::MQTT.set(client);
 
     tokio::spawn(async {
         let listener = tokio::net::TcpListener::bind("0.0.0.0:8000").await.unwrap();

src/herd_core/mqtt.rs

@@ -42,15 +42,33 @@ pub async fn handle_mqtt(topic: String, data: Vec<u8>) {
         }
         "respond" => {
             let resp: ServerResponse = serde_json::from_slice(&dec.payload).unwrap();
-
             log::info!("Got response {:?}", resp);
+
+            let entry = crate::DISPATCH.get().unwrap().get(&resp.id).unwrap();
+            entry.send(resp);
         }
         _ => {}
     }
 }
 
+pub struct TaskWaiter {
+    pub id: ulid::Ulid,
+    pub recv: crossbeam::channel::Receiver<ServerResponse>,
+}
+
+impl TaskWaiter {
+    pub async fn wait_for(&self, timeout: std::time::Duration) -> Option<ServerResponse> {
+        // TODO tokio spawn blocking?
+        self.recv.recv_timeout(timeout).ok()
+    }
+}
+
 /// Send a message to a registered `machine`
-pub async fn send_msg<T: Serialize>(client: &AsyncClient, machine: &Model<Machine>, request: T) {
+pub async fn send_msg(
+    client: &AsyncClient,
+    machine: &Model<Machine>,
+    request: ClientAction,
+) -> TaskWaiter {
     let data = serde_json::to_string(&request).unwrap();
     let pk = &machine.read().identity;
     let rec = pk.enc_key().unwrap();
@@ -66,6 +84,14 @@ pub async fn send_msg<T: Serialize>(client: &AsyncClient, machine: &Model<Machin
         .publish(topic, rumqttc::QoS::AtMostOnce, true, payload)
         .await
         .unwrap();
+
+    let (sender, recv) = tokio::sync::mpsc::channel(100);
+    crate::DISPATCH.get().unwrap().insert(request.id, sender);
+
+    TaskWaiter {
+        id: request.id,
+        recv,
+    }
 }
 
 /// Subscribe to all `device->server` topics

src/herd_core/route.rs

@@ -1,9 +1,13 @@
 use std::ops::Deref;
 
 use crate::api;
+use crate::api::ClientAction;
 use crate::api::JoinResponse;
 use crate::api::LoginParam;
+use crate::api::MachineAPI;
+use crate::api::ShellResponse;
 use crate::herd_core::model::Machine;
+use crate::herd_core::mqtt::listen_to_device;
 use axum::Json;
 use axum::extract::FromRequestParts;
 use axum::http::StatusCode;
@@ -13,6 +17,7 @@ use axum_extra::headers::Authorization;
 use axum_extra::headers::authorization::Bearer;
 use based::auth::Sessions;
 use based::auth::User;
+use owl::get;
 use owl::prelude::Model;
 use owl::query;
 use owl::save;
@@ -20,25 +25,85 @@ use serde::Deserialize;
 use serde_json::json;
 use sheepd::DeviceEntry;
 use sheepd::DeviceList;
+use ureq::http::StatusCode;
 
 use super::mqtt::is_within_80_seconds;
+use super::mqtt::send_msg;
+
+pub async fn device_shell_cmd(
+    Path((device_id)): Path<(String)>,
+    Json(payload): Json<api::ShellParam>,
+    session: TypedHeader<Authorization<Bearer>>,
+) -> (StatusCode, Json<api::Result<ShellResponse>>) {
+    // TODO : check auth
+
+    let machine: Option<Model<Machine>> = get!(device_id);
+
+    if let Some(machine) = machine {
+        let resp = send_msg(
+            crate::MQTT.get(),
+            &machine,
+            ClientAction::new(ClientActions::Shell(payload.cmd, payload.cwd)),
+        )
+        .await;
+        if let Some(resp) = resp.wait_for(std::time::Duration::from_secs(60)).await {
+            let r = match resp.response {
+                api::ServerResponses::Shell(shell_response) => shell_response,
+                _ => unreachable!(),
+            };
+            (StatusCode::OK, Json(api::Result::OkVal(r)))
+        } else {
+            (
+                StatusCode::BAD_GATEWAY,
+                Json(api::Result::Err("Did not receive response from device")),
+            )
+        }
+    } else {
+        (StatusCode::NOT_FOUND, Json(api::Result::Err("Not Found")))
+    }
+}
+
+pub async fn device_get_api(
+    Path((device_id)): Path<(String)>,
+    session: TypedHeader<Authorization<Bearer>>,
+) -> (StatusCode, Json<api::Result<DeviceEntry>>) {
+    // TODO : check auth
+
+    let machine: Option<Model<Machine>> = get!(device_id);
+
+    if let Some(machine) = machine {
+        let api = machine.read();
+        let api = DeviceEntry {
+            id: device_id,
+            hostname: api.hostname,
+            online: device_online(&device_id),
+        };
+        (StatusCode::OK, Json(api::Result::OkVal(api)))
+    } else {
+        (StatusCode::NOT_FOUND, Json(api::Result::Err("Not Found")))
+    }
+}
+
+pub fn device_online(id: &str) -> bool {
+    crate::ONLINE
+        .get()
+        .unwrap()
+        .get(&id)
+        .map(|x| is_within_80_seconds(*x.deref()))
+        .unwrap_or(false)
+}
 
 pub async fn devices_list(
     session: TypedHeader<Authorization<Bearer>>,
 ) -> (StatusCode, Json<api::Result<DeviceList>>) {
+    // TODO : auth?
     let machines: Vec<Model<Machine>> = query!(|_| true);
 
     let mut ret = vec![];
 
     for mac in machines {
         let id = mac.read().id.to_string().replace("-", "");
-        let online_state = crate::ONLINE
-            .get()
-            .unwrap()
-            .get(&id)
-            .map(|x| is_within_80_seconds(*x.deref()))
-            .unwrap_or(false);
-
+        let online_state = device_online(&id);
         ret.push(DeviceEntry {
             id: id,
             hostname: mac.read().hostname.clone(),
@@ -86,7 +151,7 @@ pub async fn join_device(
     let machine = Machine::from_join_param(payload);
     let new_token = machine.token.clone();
 
-    // TODO : add device listener instantly
+    listen_to_device(crate::MQTT.get(), &payload.machine_id).await;
 
     save!(machine);
 

src/sheepctl.rs

@@ -1,6 +1,6 @@
 use sheepctl_core::{
     args::{DeviceCommands, SheepctlArgs, SheepctlCommand},
-    cmd::{list_devices, login},
+    cmd::{interactive_shell, list_devices, login},
 };
 
 mod api;
@@ -14,5 +14,6 @@ fn main() {
             DeviceCommands::List(list_devices_command) => list_devices(list_devices_command),
         },
         SheepctlCommand::Login(login_command) => login(login_command),
+        SheepctlCommand::Shell(shell_command) => interactive_shell(shell_command),
     }
 }

src/sheepctl_core/args.rs

@@ -12,6 +12,7 @@ pub struct SheepctlArgs {
 pub enum SheepctlCommand {
     Login(LoginCommand),
     Device(DeviceCommand),
+    Shell(ShellCommand),
 }
 
 #[derive(FromArgs, PartialEq, Debug)]
@@ -49,3 +50,12 @@ pub struct ListDevicesCommand {
     /// only show online devices
     pub online: bool,
 }
+
+#[derive(FromArgs, PartialEq, Debug)]
+#[argh(subcommand, name = "shell")]
+/// Enter interactive shell
+pub struct ShellCommand {
+    #[argh(positional)]
+    /// device ID
+    pub device: String,
+}

src/sheepctl_core/cmd.rs

@@ -1,11 +1,12 @@
 use std::path::PathBuf;
 
 use owl::{Deserialize, Serialize};
-use sheepd::{DeviceList, LoginParam};
+use sheepd::{DeviceList, LoginParam, ShellResponse};
 
-use super::args::{ListDevicesCommand, LoginCommand};
-use crate::api::domain;
+use super::args::{ListDevicesCommand, LoginCommand, ShellCommand};
+use crate::api::{DeviceEntry, ShellParam, domain};
 
+/// Make an POST API call to `path` with `data` returning `Result<T>`
 pub fn api_call<T: Serialize + for<'a> Deserialize<'a>, I: Serialize>(
     server: &str,
     path: &str,
@@ -57,6 +58,11 @@ impl CtlConfig {
     }
 }
 
+pub fn get_machine_api(home: &str, token: &str, id: &str) -> Option<DeviceEntry> {
+    let res = api_call_get::<DeviceEntry>(home, &format!("device/{id}"), token);
+    res.as_result().ok()
+}
+
 pub fn list_devices(arg: ListDevicesCommand) {
     let conf = CtlConfig::load().unwrap();
 
@@ -74,6 +80,56 @@ pub fn list_devices(arg: ListDevicesCommand) {
     }
 }
 
+pub fn interactive_shell(arg: ShellCommand) {
+    let conf = CtlConfig::load().unwrap();
+    let machine = arg.device;
+
+    if let Some(machine) = get_machine_api(&conf.home, &conf.token, &machine) {
+        if !machine.online {
+            println!("Device not online.");
+            std::process::exit(1);
+        }
+
+        let mut cwd = "/".to_string();
+
+        loop {
+            print!("{} [{}]: {cwd} $ ", machine.hostname, machine.id);
+            let mut read = String::new();
+            std::io::stdin().read_line(&mut read).unwrap();
+            if read == "exit" {
+                break;
+            }
+
+            if read.starts_with("cd") {
+                let dir = read.trim_start_matches("cd ").trim_end_matches(";");
+                cwd = dir.to_string();
+                continue;
+            }
+
+            let res = api_call::<ShellResponse, _>(
+                &conf.home,
+                &format!("device/{}/shell", machine.id),
+                ShellParam {
+                    cmd: read.clone(),
+                    cwd: cwd.clone(),
+                },
+            );
+
+            if let Ok(resp) = res.as_result() {
+                println!("{} #{}\n{}", read, resp.status, resp.stdout);
+
+                if !resp.stderr.is_empty() {
+                    println!("Stderr: {}", resp.stderr);
+                }
+            } else {
+                println!("Command execution failed");
+            }
+        }
+    } else {
+        println!("No device with ID {machine}");
+    }
+}
+
 pub fn login(arg: LoginCommand) {
     if let Some(conf) = CtlConfig::load() {
         println!("You are already logged in to {}", conf.home);

src/sheepd_core/mqtt.rs

@@ -1,10 +1,10 @@
-use std::{os, process::Stdio};
+use std::process::Stdio;
 
 use owl::Serialize;
 use rumqttc::AsyncClient;
 use sage::PersonaIdentity;
 
-use crate::api::{ClientAction, ServerResponse, ServerResponses};
+use crate::api::{ClientAction, ServerResponse, ServerResponses, ShellResponse};
 
 // Client MQTT
 pub async fn handle_mqtt(topic: String, data: Vec<u8>) {
@@ -31,6 +31,28 @@ pub async fn handle_mqtt(topic: String, data: Vec<u8>) {
             )
             .await;
         }
+        crate::api::ClientActions::Shell(cmd, cwd) => {
+            log::info!("Received shell command: {cmd} in {cwd}");
+            let res = std::process::Command::new("sh")
+                .arg("-c")
+                .arg(cmd)
+                .current_dir(cwd)
+                .output()
+                .unwrap();
+            send_back(
+                crate::MQTT.get().unwrap(),
+                "respond",
+                ServerResponse::of(
+                    &action,
+                    ServerResponses::Shell(ShellResponse {
+                        stdout: String::from_utf8_lossy(&res.stdout).to_string(),
+                        stderr: String::from_utf8_lossy(&res.stderr).to_string(),
+                        status: res.status.code().unwrap(),
+                    }),
+                ),
+            )
+            .await;
+        }
     }
 }