mirror of
https://github.com/EFForg/rayhunter.git
synced 2026-04-29 08:59:59 -07:00
get config and set config
This commit is contained in:
@@ -1,10 +1,10 @@
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use rayhunter::analysis::analyzer::AnalyzerConfig;
|
||||
|
||||
use crate::error::RayhunterError;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct Config {
|
||||
pub qmdl_store_path: String,
|
||||
@@ -32,11 +32,11 @@ impl Default for Config {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_config<P>(path: P) -> Result<Config, RayhunterError>
|
||||
pub async fn parse_config<P>(path: P) -> Result<Config, RayhunterError>
|
||||
where
|
||||
P: AsRef<std::path::Path>,
|
||||
{
|
||||
if let Ok(config_file) = std::fs::read_to_string(&path) {
|
||||
if let Ok(config_file) = tokio::fs::read_to_string(&path).await {
|
||||
Ok(toml::from_str(&config_file).map_err(RayhunterError::ConfigFileParsingError)?)
|
||||
} else {
|
||||
Ok(Config::default())
|
||||
|
||||
@@ -20,7 +20,7 @@ use crate::diag::run_diag_read_thread;
|
||||
use crate::error::RayhunterError;
|
||||
use crate::pcap::get_pcap;
|
||||
use crate::qmdl_store::RecordingStore;
|
||||
use crate::server::{get_qmdl, restart_daemon, serve_static, ServerState};
|
||||
use crate::server::{get_config, get_qmdl, restart_daemon, serve_static, set_config, ServerState};
|
||||
use crate::stats::{get_qmdl_manifest, get_system_stats};
|
||||
|
||||
use analysis::{
|
||||
@@ -60,6 +60,8 @@ fn get_router() -> AppRouter {
|
||||
.route("/api/analysis", get(get_analysis_status))
|
||||
.route("/api/analysis/{name}", post(start_analysis))
|
||||
.route("/api/restart-daemon", post(restart_daemon))
|
||||
.route("/api/config", get(get_config))
|
||||
.route("/api/config", post(set_config))
|
||||
.route("/", get(|| async { Redirect::permanent("/index.html") }))
|
||||
.route("/{*path}", get(serve_static))
|
||||
}
|
||||
@@ -187,8 +189,8 @@ async fn main() -> Result<(), RayhunterError> {
|
||||
let args = parse_args();
|
||||
|
||||
loop {
|
||||
let config = parse_config(&args.config_path)?;
|
||||
if !run_with_config(&config).await? {
|
||||
let config = parse_config(&args.config_path).await?;
|
||||
if !run_with_config(&args, &config).await? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -199,7 +201,10 @@ async fn main() -> Result<(), RayhunterError> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_with_config(config: &config::Config) -> Result<bool, RayhunterError> {
|
||||
async fn run_with_config(
|
||||
args: &config::Args,
|
||||
config: &config::Config,
|
||||
) -> Result<bool, RayhunterError> {
|
||||
// TaskTrackers give us an interface to spawn tokio threads, and then
|
||||
// eventually await all of them ending
|
||||
let task_tracker = TaskTracker::new();
|
||||
@@ -264,6 +269,7 @@ async fn run_with_config(config: &config::Config) -> Result<bool, RayhunterError
|
||||
analysis_tx.clone(),
|
||||
);
|
||||
let state = Arc::new(ServerState {
|
||||
config_path: args.config_path.clone(),
|
||||
qmdl_store_lock: qmdl_store_lock.clone(),
|
||||
diag_device_ctrl_sender: diag_tx,
|
||||
ui_update_sender: ui_update_tx,
|
||||
|
||||
@@ -4,18 +4,23 @@ use axum::extract::State;
|
||||
use axum::http::header::{self, CONTENT_LENGTH, CONTENT_TYPE};
|
||||
use axum::http::{HeaderValue, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::Json;
|
||||
use include_dir::{include_dir, Dir};
|
||||
use std::sync::Arc;
|
||||
use tokio::fs::write;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::sync::mpsc::Sender;
|
||||
use tokio::sync::{oneshot, RwLock};
|
||||
use tokio_util::io::ReaderStream;
|
||||
|
||||
use crate::analysis::{AnalysisCtrlMessage, AnalysisStatus};
|
||||
use crate::config::parse_config;
|
||||
use crate::config::Config;
|
||||
use crate::qmdl_store::RecordingStore;
|
||||
use crate::{display, DiagDeviceCtrlMessage};
|
||||
|
||||
pub struct ServerState {
|
||||
pub config_path: String,
|
||||
pub qmdl_store_lock: Arc<RwLock<RecordingStore>>,
|
||||
pub diag_device_ctrl_sender: Sender<DiagDeviceCtrlMessage>,
|
||||
pub ui_update_sender: Sender<display::DisplayState>,
|
||||
@@ -35,12 +40,15 @@ pub async fn get_qmdl(
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("couldn't find qmdl file with name {}", qmdl_idx),
|
||||
))?;
|
||||
let qmdl_file = qmdl_store.open_entry_qmdl(entry_index).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("error opening QMDL file: {}", e),
|
||||
)
|
||||
})?;
|
||||
let qmdl_file = qmdl_store
|
||||
.open_entry_qmdl(entry_index)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("error opening QMDL file: {}", err),
|
||||
)
|
||||
})?;
|
||||
let limited_qmdl_file = qmdl_file.take(entry.qmdl_size_bytes as u64);
|
||||
let qmdl_stream = ReaderStream::new(limited_qmdl_file);
|
||||
|
||||
@@ -84,13 +92,12 @@ pub async fn restart_daemon(
|
||||
let mut restart_tx = state.daemon_restart_tx.write().await;
|
||||
|
||||
if let Some(sender) = restart_tx.take() {
|
||||
sender.send(()).map_err(|()| {
|
||||
sender.send(()).map_err(|_| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"couldn't send restart signal".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok((StatusCode::ACCEPTED, "restart signal sent".to_string()))
|
||||
} else {
|
||||
Ok((
|
||||
@@ -99,3 +106,37 @@ pub async fn restart_daemon(
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_config(
|
||||
State(state): State<Arc<ServerState>>,
|
||||
) -> Result<Json<Config>, (StatusCode, String)> {
|
||||
let config = parse_config(&state.config_path).await.map_err(|err| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("failed to read config file: {}", err),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(Json(config))
|
||||
}
|
||||
|
||||
pub async fn set_config(
|
||||
State(state): State<Arc<ServerState>>,
|
||||
Json(config): Json<Config>,
|
||||
) -> Result<(StatusCode, String), (StatusCode, String)> {
|
||||
let config_str = toml::to_string_pretty(&config).map_err(|err| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("failed to serialize config as TOML: {}", err),
|
||||
)
|
||||
})?;
|
||||
|
||||
write(&state.config_path, config_str).await.map_err(|err| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("failed to write config file: {}", err),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok((StatusCode::ACCEPTED, "wrote config".to_string()))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user