From 5c53c57d74c24bed1786ef804b59e7a04e51aaad Mon Sep 17 00:00:00 2001 From: JMARyA Date: Tue, 6 May 2025 22:48:24 +0200 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20osquery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/api.rs | 5 +++++ src/herd.rs | 3 ++- src/herd_core/route.rs | 33 +++++++++++++++++++++++++++++++++ src/sheepctl.rs | 3 ++- src/sheepctl_core/args.rs | 13 +++++++++++++ src/sheepctl_core/cmd.rs | 33 +++++++++++++++++++++++++++++++-- 6 files changed, 86 insertions(+), 4 deletions(-) diff --git a/src/api.rs b/src/api.rs index ab26d36..a285fac 100644 --- a/src/api.rs +++ b/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, diff --git a/src/herd.rs b/src/herd.rs index f1c9305..c45f866 100644 --- a/src/herd.rs +++ b/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); diff --git a/src/herd_core/route.rs b/src/herd_core/route.rs index 605e2a1..48e8b54 100644 --- a/src/herd_core/route.rs +++ b/src/herd_core/route.rs @@ -42,6 +42,39 @@ macro_rules! check_admin { }; } +pub async fn device_osquery( + Path(device_id): Path, + APIUser(user): APIUser, + Json(payload): Json, +) -> (StatusCode, Json>) { + check_admin!(user); + + let machine: Option> = 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, APIUser(user): APIUser, diff --git a/src/sheepctl.rs b/src/sheepctl.rs index c5fc437..2c9bf28 100644 --- a/src/sheepctl.rs +++ b/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), } } diff --git a/src/sheepctl_core/args.rs b/src/sheepctl_core/args.rs index 7357db0..09ea27f 100644 --- a/src/sheepctl_core/args.rs +++ b/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, +} diff --git a/src/sheepctl_core/cmd.rs b/src/sheepctl_core/cmd.rs index 216a363..5812607 100644 --- a/src/sheepctl_core/cmd.rs +++ b/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` pub fn api_call 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::( + &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);