diff --git a/crates/common/src/manager/boot.rs b/crates/common/src/manager/boot.rs index bf3dd8b5..79079ff9 100644 --- a/crates/common/src/manager/boot.rs +++ b/crates/common/src/manager/boot.rs @@ -27,6 +27,7 @@ use crate::{ use super::{ backup::BackupParams, config::{ConfigManager, Patterns}, + console::store_console, WEBADMIN_KEY, }; @@ -56,6 +57,7 @@ Options: -c, --config Start server with the specified configuration file -e, --export Export all store data to a specific path -i, --import Import store data from a specific path + -o, --console Open the store console -I, --init Initialize a new server at a specific path -h, --help Print help -V, --version Print version @@ -63,16 +65,17 @@ Options: ); #[derive(PartialEq, Eq)] -enum ImportExport { +enum StoreOp { Export(BackupParams), Import(PathBuf), + Console, None, } impl BootManager { pub async fn init() -> Self { let mut config_path = std::env::var("CONFIG_PATH").ok(); - let mut import_export = ImportExport::None; + let mut import_export = StoreOp::None; if config_path.is_none() { let mut args = std::env::args().skip(1); @@ -105,10 +108,13 @@ impl BootManager { std::process::exit(0); } ("export" | "e", Some(value)) => { - import_export = ImportExport::Export(BackupParams::new(value.into())); + import_export = StoreOp::Export(BackupParams::new(value.into())); } ("import" | "i", Some(value)) => { - import_export = ImportExport::Import(value.into()); + import_export = StoreOp::Import(value.into()); + } + ("console" | "o", None) => { + import_export = StoreOp::Console; } (_, None) => { failed(&format!("Unrecognized command '{key}', try '--help'.")); @@ -120,7 +126,7 @@ impl BootManager { } if config_path.is_none() { - if import_export == ImportExport::None { + if import_export == StoreOp::None { eprintln!("{HELP}"); } else { eprintln!("Missing '--config' argument for import/export.") @@ -181,7 +187,7 @@ impl BootManager { let telemetry = Telemetry::parse(&mut config, &stores); match import_export { - ImportExport::None => { + StoreOp::None => { // Add hostname lookup if missing let mut insert_keys = Vec::new(); if config @@ -357,7 +363,7 @@ impl BootManager { ipc_rxs, } } - ImportExport::Export(path) => { + StoreOp::Export(path) => { // Enable telemetry telemetry.enable(false); @@ -368,7 +374,7 @@ impl BootManager { .await; std::process::exit(0); } - ImportExport::Import(path) => { + StoreOp::Import(path) => { // Enable telemetry telemetry.enable(false); @@ -379,6 +385,11 @@ impl BootManager { .await; std::process::exit(0); } + StoreOp::Console => { + // Store console + store_console(Core::parse(&mut config, stores, manager).await.storage.data).await; + std::process::exit(0); + } } } } diff --git a/crates/common/src/manager/console.rs b/crates/common/src/manager/console.rs new file mode 100644 index 00000000..438d7137 --- /dev/null +++ b/crates/common/src/manager/console.rs @@ -0,0 +1,264 @@ +/* + * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd + * + * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + */ + +use std::env; +use std::io::{self, Write}; + +use base64::engine::general_purpose; +use base64::Engine; +use store::write::{AnyClass, AnyKey, BatchBuilder, ValueClass}; +use store::{ + Deserialize, IterateParams, Store, SUBSPACE_BITMAP_ID, SUBSPACE_BITMAP_TAG, + SUBSPACE_BITMAP_TEXT, SUBSPACE_INDEXES, +}; + +const HELP: &str = concat!( + "Stalwart Mail Server v", + env!("CARGO_PKG_VERSION"), + r#" Data Store CLI + +Enter commands (type 'help' for available commands). +"# +); + +pub async fn store_console(store: Store) { + print!("{HELP}"); + + if matches!(store, Store::None) { + println!("No store available. Verify your configuration."); + return; + } + + loop { + print!("> "); + io::stdout().flush().unwrap(); + + let mut input = String::new(); + io::stdin().read_line(&mut input).unwrap(); + let input = input.trim(); + + let parts: Vec<&str> = input.split_whitespace().collect(); + + if parts.is_empty() { + continue; + } + + match parts[0] { + "scan" => { + if parts.len() != 3 { + println!("Usage: scan "); + } else if let (Some(from_key), Some(to_key)) = + (parse_key(parts[1]), parse_key(parts[2])) + { + println!("Scanning from {:?} to {:?}", from_key, to_key); + let mut from_key = from_key.into_iter(); + let mut to_key = to_key.into_iter(); + let subspace = from_key.next().unwrap(); + + store + .iterate( + IterateParams::new( + AnyKey { + subspace, + key: from_key.collect::>(), + }, + AnyKey { + subspace: to_key.next().unwrap(), + key: to_key.collect::>(), + }, + ) + .set_values( + ![ + SUBSPACE_INDEXES, + SUBSPACE_BITMAP_ID, + SUBSPACE_BITMAP_TAG, + SUBSPACE_BITMAP_TEXT, + ] + .contains(&subspace), + ), + |key, value| { + print_escaped(key); + print!(" : "); + print_escaped(value); + println!(); + Ok(true) + }, + ) + .await + .expect("Failed to scan keys"); + } + } + "delete" => { + if parts.len() != 2 { + println!("Usage: delete "); + } else if let Some(key) = parse_key(parts[1]) { + println!("Deleting key: {:?}", key); + let mut key = key.into_iter(); + let mut batch = BatchBuilder::new(); + batch.clear(ValueClass::Any(AnyClass { + subspace: key.next().unwrap(), + key: key.collect(), + })); + if let Err(err) = store.write(batch.build()).await { + println!("Failed to delete key: {}", err); + } + } + } + "get" => { + if parts.len() != 2 { + println!("Usage: get "); + } else if let Some(key) = parse_key(parts[1]) { + let mut key = key.into_iter(); + match store + .get_value::(AnyKey { + subspace: key.next().unwrap(), + key: key.collect::>(), + }) + .await + { + Ok(Some(data)) => { + print_escaped(&data.0); + println!(); + } + Ok(None) => { + println!("Key not found."); + } + Err(err) => { + println!("Failed to retrieve key: {}", err); + } + } + } + } + "put" => { + if parts.len() != 2 { + println!("Usage: put []"); + } else if let Some(key) = parse_key(parts[1]) { + let value = parts.get(2).map(|v| parse_value(v)).unwrap_or_default(); + println!("Putting key: {key:?}"); + + let mut key = key.into_iter(); + let mut batch = BatchBuilder::new(); + batch.set( + ValueClass::Any(AnyClass { + subspace: key.next().unwrap(), + key: key.collect(), + }), + value, + ); + if let Err(err) = store.write(batch.build()).await { + println!("Failed to insert key: {}", err); + } + } + } + "help" => { + print_help(); + } + "exit" | "quit" => { + println!("Exiting..."); + break; + } + _ => { + println!("Unknown command. Type 'help' for available commands."); + } + } + } +} + +fn parse_key(input: &str) -> Option> { + let result = if let Some(key) = input.strip_prefix("base64:") { + base64_decode(key) + } else { + parse_binary(input) + }; + if matches!(result.first(), Some(ch) if ch.is_ascii_alphabetic() && ch.is_ascii_lowercase()) { + Some(result) + } else { + println!("Invalid key: {result:?}"); + None + } +} + +fn parse_value(input: &str) -> Vec { + if let Some(key) = input.strip_prefix("base64:") { + base64_decode(key) + } else { + parse_binary(input) + } +} + +fn base64_decode(input: &str) -> Vec { + general_purpose::STANDARD + .decode(input) + .expect("Failed to decode base64") +} + +fn parse_binary(input: &str) -> Vec { + let mut result = Vec::new(); + let mut chars = input.chars().peekable(); + + while let Some(c) = chars.next() { + if c == '\\' { + match chars.next() { + Some('x') => { + let hex: String = chars.by_ref().take(2).collect(); + if hex.len() == 2 { + if let Ok(byte) = u8::from_str_radix(&hex, 16) { + result.push(byte); + } else { + result.extend_from_slice(b"\\x"); + result.extend_from_slice(hex.as_bytes()); + } + } else { + result.push(b'\\'); + result.push(b'x'); + result.extend_from_slice(hex.as_bytes()); + } + } + Some(other) => { + result.push(b'\\'); + result.push(other as u8); + } + None => { + result.push(b'\\'); + } + } + } else { + result.push(c as u8); + } + } + + result +} + +fn print_escaped(bytes: &[u8]) { + for ch in bytes { + if ch.is_ascii() && !ch.is_ascii_control() && *ch != b'\\' { + print!("{}", *ch as char); + } else { + print!("\\x{:02x}", ch); + } + } +} + +fn print_help() { + println!("Available commands:"); + println!(" scan "); + println!(" delete "); + println!(" get "); + println!(" put []"); + println!(" help"); + println!(" exit/quit"); + println!("Note: Keys and values can be prefixed with 'base64:' for base64 encoding"); + println!(" or use escaped hex values (e.g., \\x41 for 'A')"); +} + +struct RawValue(Vec); + +impl Deserialize for RawValue { + fn deserialize(bytes: &[u8]) -> trc::Result { + Ok(RawValue(bytes.to_vec())) + } +} diff --git a/crates/common/src/manager/mod.rs b/crates/common/src/manager/mod.rs index a5cc7e42..a58ab779 100644 --- a/crates/common/src/manager/mod.rs +++ b/crates/common/src/manager/mod.rs @@ -13,6 +13,7 @@ use self::config::ConfigManager; pub mod backup; pub mod boot; pub mod config; +pub mod console; pub mod reload; pub mod restore; pub mod webadmin;