Store CLI

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 <PATH>              Start server with the specified configuration file
   -e, --export <PATH>              Export all store data to a specific path
   -i, --import <PATH>              Import store data from a specific path
+  -o, --console                    Open the store console
   -I, --init <PATH>                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);
+            }
         }
     }
 }

crates/common/src/manager/console.rs

@@ -0,0 +1,264 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
+ *
+ * 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 <from_key> <to_key>");
+                } 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::<Vec<_>>(),
+                                },
+                                AnyKey {
+                                    subspace: to_key.next().unwrap(),
+                                    key: to_key.collect::<Vec<_>>(),
+                                },
+                            )
+                            .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 <key>");
+                } 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 <key>");
+                } else if let Some(key) = parse_key(parts[1]) {
+                    let mut key = key.into_iter();
+                    match store
+                        .get_value::<RawValue>(AnyKey {
+                            subspace: key.next().unwrap(),
+                            key: key.collect::<Vec<_>>(),
+                        })
+                        .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 <key> [<value>]");
+                } 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<Vec<u8>> {
+    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<u8> {
+    if let Some(key) = input.strip_prefix("base64:") {
+        base64_decode(key)
+    } else {
+        parse_binary(input)
+    }
+}
+
+fn base64_decode(input: &str) -> Vec<u8> {
+    general_purpose::STANDARD
+        .decode(input)
+        .expect("Failed to decode base64")
+}
+
+fn parse_binary(input: &str) -> Vec<u8> {
+    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 <from_key> <to_key>");
+    println!("  delete <key>");
+    println!("  get <key>");
+    println!("  put <key> [<value>]");
+    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<u8>);
+
+impl Deserialize for RawValue {
+    fn deserialize(bytes: &[u8]) -> trc::Result<Self> {
+        Ok(RawValue(bytes.to_vec()))
+    }
+}

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;