✨ osquery

src/api.rs

@@ -46,6 +46,11 @@ pub struct ShellParam {
     pub cwd: String,
 }
 
+#[derive(Deserialize, Serialize)]
+pub struct QueryParam {
+    pub query: String,
+}
+
 #[derive(Deserialize, Serialize, Debug)]
 pub struct ShellResponse {
     pub stdout: String,

src/herd.rs

@@ -14,7 +14,7 @@ mod herd_core;
 use crate::herd_core::mqtt::{handle_mqtt, listen_to_devices};
 use herd_core::{
     config::Config,
-    route::{device_get_api, device_shell_cmd, join_device, login_user},
+    route::{device_get_api, device_osquery, device_shell_cmd, join_device, login_user},
 };
 use herd_core::{model::Machine, route::devices_list};
 use rumqttc::AsyncClient;
@@ -69,6 +69,7 @@ async fn main() {
         .route("/login", post(login_user))
         .route("/device/{device_id}", get(device_get_api))
         .route("/device/{device_id}/shell", post(device_shell_cmd))
+        .route("/device/{device_id}/osquery", post(device_osquery))
         .route("/devices", get(devices_list));
 
     let app = Router::new().merge(device).merge(user);

src/herd_core/route.rs

@@ -42,6 +42,39 @@ macro_rules! check_admin {
     };
 }
 
+pub async fn device_osquery(
+    Path(device_id): Path<String>,
+    APIUser(user): APIUser,
+    Json(payload): Json<api::QueryParam>,
+) -> (StatusCode, Json<api::Result<String>>) {
+    check_admin!(user);
+
+    let machine: Option<Model<Machine>> = get!(device_id);
+
+    if let Some(machine) = machine {
+        let resp = send_msg(
+            crate::MQTT.get().unwrap(),
+            &machine,
+            ClientAction::new(ClientActions::OSQuery(payload.query)),
+        )
+        .await;
+        if let Some(resp) = resp.wait_for(std::time::Duration::from_secs(60)).await {
+            let r = match resp.response {
+                api::ServerResponses::OSQuery(res) => res,
+                _ => 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_shell_cmd(
     Path(device_id): Path<String>,
     APIUser(user): APIUser,

src/sheepctl.rs

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

src/sheepctl_core/args.rs

@@ -13,6 +13,7 @@ pub enum SheepctlCommand {
     Login(LoginCommand),
     Device(DeviceCommand),
     Shell(ShellCommand),
+    Query(QueryCommand),
 }
 
 #[derive(FromArgs, PartialEq, Debug)]
@@ -59,3 +60,15 @@ pub struct ShellCommand {
     /// device ID
     pub device: String,
 }
+
+#[derive(FromArgs, PartialEq, Debug)]
+#[argh(subcommand, name = "query")]
+/// Run an osquery
+pub struct QueryCommand {
+    #[argh(positional)]
+    /// device ID
+    pub device: String,
+    #[argh(positional)]
+    /// query
+    pub query: String,
+}

src/sheepctl_core/cmd.rs

@@ -3,8 +3,8 @@ use std::{io::Write, path::PathBuf};
 use owl::{Deserialize, Serialize};
 use sheepd::{DeviceList, LoginParam, ShellResponse};
 
-use super::args::{ListDevicesCommand, LoginCommand, ShellCommand};
-use crate::api::{DeviceEntry, ShellParam, domain};
+use super::args::{ListDevicesCommand, LoginCommand, QueryCommand, ShellCommand};
+use crate::api::{DeviceEntry, QueryParam, 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>(
@@ -147,6 +147,35 @@ pub fn interactive_shell(arg: ShellCommand) {
     }
 }
 
+pub fn run_osquery(args: QueryCommand) {
+    // TODO : sanity checks
+
+    let conf = CtlConfig::load().unwrap();
+    let machine = args.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 res = api_call_post_auth::<String, _>(
+            &conf.home,
+            &format!("device/{}/osquery", machine.id),
+            &conf.token,
+            QueryParam { query: args.query },
+        );
+
+        if let Ok(res) = res.as_result() {
+            println!("{res}");
+        } else {
+            println!("Error doing query");
+        }
+    } 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);