Start rewrite, get layers working

Cargo.toml

@@ -8,14 +8,12 @@ clap = { version = "4.5.39", features = ["derive"] }
 crossbeam-channel = "0.5.15"
 lazy_static = "1.5.0"
 log = "0.4.27"
-mio = { version = "1.0.4", features = ["os-poll"] }
-native-tls = "0.2.14"
 pretty_env_logger = "0.5.0"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"
-slint = "1.11.0"
+# slint = "1.11.0"
 unshell-rs-lib = { path = "./unshell-rs-lib" }
 
 
-[build-dependencies]
-slint-build = "1.11.0"
+# [build-dependencies]
+# slint-build = "1.11.0"

payload/Cargo.toml

@@ -1,9 +0,0 @@
-[package]
-name = "payload"
-version = "0.1.0"
-edition = "2024"
-
-[dependencies]
-serde_json = "1.0.140"
-# libc = "0.2.172"
-unshell-rs-lib = { path = "../unshell-rs-lib" }

payload/src/execute.rs

@@ -1,79 +0,0 @@
-use std::error::Error;
-
-#[allow(dead_code)]
-#[cfg(unix)]
-unsafe fn execute_in_memory(binary_data: &[u8]) -> Result<(), Box<dyn Error>> {
-    use std::mem;
-
-    // Allocate executable memory
-    let size = binary_data.len();
-    let page_size = 4096; // Typical page size
-    let aligned_size = (size + page_size - 1) & !(page_size - 1);
-
-    let ptr = libc::mmap(
-        std::ptr::null_mut(),
-        aligned_size,
-        libc::PROT_READ | libc::PROT_WRITE,
-        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
-        -1,
-        0,
-    );
-
-    if ptr == libc::MAP_FAILED {
-        return Err(Box::new(std::io::Error::last_os_error()));
-    }
-
-    // Copy binary data to allocated memory
-    std::ptr::copy_nonoverlapping(binary_data.as_ptr(), ptr as *mut u8, size);
-
-    // Make memory executable
-    if libc::mprotect(ptr, aligned_size, libc::PROT_READ | libc::PROT_EXEC) != 0 {
-        libc::munmap(ptr, aligned_size);
-        return Err(Box::new(std::io::Error::last_os_error()));
-    }
-
-    // Cast to function pointer and execute
-    // This assumes the binary is a simple executable that can be called as a function
-    // For ELF binaries, you'd need proper ELF parsing and loading
-    let func: extern "C" fn() = mem::transmute(ptr);
-
-    println!("Executing binary...");
-    func();
-
-    // Clean up
-    libc::munmap(ptr, aligned_size);
-
-    Ok(())
-}
-
-#[cfg(windows)]
-unsafe fn execute_in_memory(binary_data: &[u8]) -> Result<(), Box<dyn Error>> {
-    use std::mem;
-    use std::ptr;
-
-    // Allocate executable memory
-    let ptr = winapi::um::memoryapi::VirtualAlloc(
-        ptr::null_mut(),
-        binary_data.len(),
-        winapi::um::winnt::MEM_COMMIT | winapi::um::winnt::MEM_RESERVE,
-        winapi::um::winnt::PAGE_EXECUTE_READWRITE,
-    );
-
-    if ptr.is_null() {
-        return Err(Box::new(std::io::Error::last_os_error()));
-    }
-
-    // Copy binary data to allocated memory
-    ptr::copy_nonoverlapping(binary_data.as_ptr(), ptr as *mut u8, binary_data.len());
-
-    // Cast to function pointer and execute
-    let func: extern "C" fn() = mem::transmute(ptr);
-
-    println!("Executing binary...");
-    func();
-
-    // Clean up
-    winapi::um::memoryapi::VirtualFree(ptr, 0, winapi::um::winnt::MEM_RELEASE);
-
-    Ok(())
-}

payload/src/main.rs

@@ -1,21 +0,0 @@
-// #[allow(unsafe_op_in_unsafe_fn)]
-// mod execute;
-
-use std::error::Error;
-
-use std::{
-    sync::{Arc, Mutex},
-    thread,
-    time::Duration,
-};
-
-use unshell_rs_lib::{
-    networkers::{ClientTrait, Connection, TCPClient, TCPConnection},
-    packets::Packet,
-};
-
-fn main() -> Result<(), Box<dyn Error>> {
-    run_client::<TCPConnection, TCPClient>("127.0.0.1:3000")?;
-
-    Ok(())
-}

src/client/cli.rs

@@ -0,0 +1,33 @@
+use std::{error::Error, io::Write, net::SocketAddr};
+
+use unshell_rs_lib::{
+    layers::{LayerConfig, build_client},
+    networkers::{ClientTrait, Connection, TCPClient},
+};
+
+use crate::client;
+
+pub struct Cli;
+
+impl Cli {
+    pub fn connect(addr: SocketAddr) -> Result<(), Box<dyn Error>> {
+        let mut client = build_client(
+            TCPClient::connect(&addr)?,
+            vec![LayerConfig::Handshake, LayerConfig::Base64],
+        )?;
+
+        let stdin = std::io::stdin();
+        let mut stdout = std::io::stdout();
+
+        loop {
+            print!("> ");
+            stdout.flush()?;
+
+            let mut input = String::new();
+            stdin.read_line(&mut input)?;
+            let input = input.trim();
+
+            client.write(input)?;
+        }
+    }
+}

src/client/client.rs

@@ -1,46 +0,0 @@
-use std::{
-    error::Error,
-    net::SocketAddr,
-    sync::{Arc, Mutex},
-};
-
-use crossbeam_channel::{Receiver, Sender};
-use unshell_rs_lib::{
-    connection::{C2Packet, Parameter, Parameters},
-    networkers::{AsyncConnection, ClientTrait, TCPClient, TCPConnection},
-};
-
-pub struct UnshellClient {
-    #[allow(dead_code)]
-    addr: SocketAddr,
-    tx: Sender<C2Packet>,
-    pub rx: Receiver<C2Packet>,
-    parameters: Arc<Mutex<Parameters>>,
-}
-
-impl UnshellClient {
-    pub fn new(addr: SocketAddr) -> Result<Self, Box<dyn Error>> {
-        let client = TCPClient::connect(&addr)?;
-
-        let (tx, rx) = TCPConnection::as_async(client);
-
-        let parameters = Arc::new(Mutex::new(Parameters::new()));
-
-        Ok(Self {
-            addr,
-            tx,
-            rx,
-            parameters,
-        })
-    }
-
-    pub fn set_parameter(&mut self, key: String, param: Parameter) {
-        let mut params_lock = self.parameters.lock().unwrap();
-        params_lock.insert(key.clone(), param.clone());
-        self.tx.send(C2Packet::SetParameter(key, param)).unwrap();
-    }
-
-    pub fn get_parameter(&self, key: &str) -> Option<Parameter> {
-        self.parameters.lock().unwrap().get(key).cloned()
-    }
-}

src/client/gui.rs

@@ -1,76 +0,0 @@
-use slint::{ComponentHandle, Weak};
-use std::{
-    error::Error,
-    sync::{Arc, Mutex},
-    thread,
-};
-use unshell_rs_lib::connection::{C2Packet, Parameter};
-
-use crate::client::UnshellClient;
-
-pub struct UnshellGui {}
-
-slint::include_modules!();
-impl UnshellGui {
-    pub fn start(client: UnshellClient) -> Result<(), Box<dyn Error>> {
-        let ui = AppWindow::new()?;
-        let rx = client.rx.clone();
-        let client = Arc::new(Mutex::new(client));
-
-        let ui_handle = ui.as_weak();
-        let client_clone = Arc::clone(&client);
-        ui.on_tab_clicked(move |index| {
-            let ui = ui_handle.unwrap();
-            ui.set_current_tab(index);
-            let mut client_lock = client_clone.lock().unwrap();
-            client_lock.set_parameter("Current Tab".to_string(), Parameter::CurrentTab(index));
-            trace!("Tab {} selected", index);
-        });
-
-        ui.set_app_info({
-            (String::new()
-                + "Unshell\n"
-                + "Version "
-                + env!("CARGO_PKG_VERSION")
-                + "\n\n View the source code at:\n https://github.com/astatin3/unshell-rs")
-                .into()
-        });
-
-        let ui_handle = ui.as_weak();
-        thread::spawn(move || {
-            fn on_param_update(ui_handle: Weak<AppWindow>, parameter: &Parameter) {
-                // info!("{}", name);
-                match parameter {
-                    Parameter::Test1 => todo!(),
-                    Parameter::CurrentTab(i) => {
-                        let i = i.clone();
-                        slint::invoke_from_event_loop(move || {
-                            ui_handle.unwrap().set_current_tab(i)
-                        })
-                        .unwrap();
-                    }
-                }
-            }
-
-            loop {
-                if let Ok(data) = rx.recv() {
-                    match data {
-                        C2Packet::SetAllParameters(parameters) => {
-                            for key in parameters.keys() {
-                                on_param_update(ui_handle.clone(), parameters.get(key).unwrap());
-                            }
-                        }
-                        C2Packet::ParameterUpate(name, parameter) => {
-                            on_param_update(ui_handle.clone(), &parameter);
-                        }
-                        _ => {}
-                    }
-                }
-            }
-        });
-
-        ui.run()?;
-
-        Ok(())
-    }
-}

src/client/mod.rs

@@ -1,6 +1,3 @@
-mod gui;
+mod cli;
 
-mod client;
-
-pub use client::UnshellClient;
-pub use gui::UnshellGui;
+pub use cli::Cli;

src/lib.rs

@@ -2,8 +2,10 @@
 extern crate log;
 
 mod client;
-mod server;
+// mod server;
 
-pub use client::UnshellClient;
-pub use client::UnshellGui;
-pub use server::UnshellServer;
+pub use client::Cli;
+
+// pub use client::UnshellClient;
+// pub use client::UnshellGui;
+// pub use server::UnshellServer;

src/main.rs

@@ -7,7 +7,9 @@ use std::{
 
 use clap::{Parser, Subcommand};
 use log::error;
-use unshell_rs::{UnshellClient, UnshellGui, UnshellServer};
+use unshell_rs::Cli;
+use unshell_rs_lib::connection::Node;
+// use unshell_rs::{UnshellClient, UnshellGui, UnshellServer};
 // use unshell_rs
 
 pub static DEFAULT_CONFIG_FILEPATH: &'static str = "server_config.json";
@@ -32,7 +34,7 @@ struct Args {
 enum Commands {
     /// Run as a service, and potentially hosting a website
     #[command(arg_required_else_help = true)]
-    Server {
+    Relay {
         /// IPv4 to listen for clients on.
         host: String,
 
@@ -47,8 +49,8 @@ enum Commands {
         // #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
         // web_port: u16,
     },
-    /// Run GUI and connect to remote server
-    Remote {
+    /// Connect to remote server
+    Connect {
         /// Remote server to connect to
         host: String,
 
@@ -56,12 +58,6 @@ enum Commands {
         #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
         port: u16,
     },
-    /// Run both server and GUI on local machine.
-    Local {
-        /// Json file to store config
-        #[arg(short, long, default_value_t = DEFAULT_CONFIG_FILEPATH.to_string())]
-        config_filepath: String,
-    },
 }
 
 fn main() -> Result<(), Box<dyn Error>> {
@@ -73,41 +69,31 @@ fn main() -> Result<(), Box<dyn Error>> {
     let args = Args::parse();
 
     match args.command {
-        Commands::Local { config_filepath } => {
-            let mut server = UnshellServer::from_filepath(config_filepath.as_str());
-            server.run(LOCAL_SOCKET)?;
-
-            let client = UnshellClient::new(LOCAL_SOCKET)?;
-
-            UnshellGui::start(client)?;
-        }
-        Commands::Remote { host, port } => {
-            let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
-            let client = UnshellClient::new(if let Ok(addr) = addr {
-                addr
-            } else {
-                error!("Could not parse address!");
-                return Ok(());
-            })?;
-
-            UnshellGui::start(client)?;
-        }
-        Commands::Server {
+        Commands::Relay {
             host,
             port,
             config_filepath,
         } => {
-            let mut unshell_server = UnshellServer::from_filepath(config_filepath.as_str());
-
             let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
-            if let Ok(addr) = addr {
-                unshell_server.run(addr)?;
+            if let Err(e) = Node::run(if let Ok(addr) = addr {
+                addr
             } else {
                 error!("Could not parse address!");
                 return Ok(());
+            }) {
+                error!("{}", e);
+            }
+        }
+        Commands::Connect { host, port } => {
+            let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
+            if let Err(e) = Cli::connect(if let Ok(addr) = addr {
+                addr
+            } else {
+                error!("Could not parse address!");
+                return Ok(());
+            }) {
+                error!("{}", e);
             }
-
-            loop {}
         }
     };
 

src/server/config.rs

@@ -1,27 +0,0 @@
-use lazy_static::lazy_static;
-use serde::{Deserialize, Serialize};
-use unshell_rs_lib::config::campaign::CampaignConfig;
-
-use unshell_rs_lib::connection::Parameters;
-
-lazy_static! {
-    pub static ref DEFAULT_CAMPAIGN: CampaignConfig = CampaignConfig {
-        name: "Default Campaign".to_string(),
-        listeners: Vec::new(),
-    };
-    pub static ref DEFAULT_USERS: Vec<User> = vec![User {
-        name: "User".into(),
-        key: "CHANGEME".to_string(),
-    }];
-    pub static ref DEFAULT_PARAMETERS: Parameters = {
-        let p = Parameters::new();
-
-        p
-    };
-}
-
-#[derive(Serialize, Deserialize, Clone)]
-pub struct User {
-    pub name: String,
-    pub key: String,
-}

src/server/mod.rs

@@ -1,6 +0,0 @@
-mod config;
-mod server;
-
-pub use crate::server::config::{DEFAULT_CAMPAIGN, DEFAULT_USERS, User};
-
-pub use server::UnshellServer;

src/server/server.rs

@@ -1,155 +0,0 @@
-use std::{
-    error::Error,
-    fs::File,
-    io::Read,
-    net::SocketAddr,
-    sync::{Arc, Mutex},
-    thread,
-};
-
-use crossbeam_channel::Sender;
-use serde::{Deserialize, Serialize};
-
-use unshell_rs_lib::{
-    config::campaign::CampaignConfig,
-    connection::{C2Packet, ErrorPacket, Parameters},
-    networkers::{AsyncConnection, ServerTrait, TCPConnection, TCPServer, run_listener_state},
-};
-
-use crate::server::{DEFAULT_CAMPAIGN, DEFAULT_USERS, User, config::DEFAULT_PARAMETERS};
-
-#[derive(Serialize, Deserialize)]
-pub struct UnshellServerConfig {
-    campaign: CampaignConfig,
-    parameters: Parameters,
-    users: Vec<User>,
-
-    #[serde(skip)]
-    clients: Vec<Client>,
-}
-
-impl UnshellServerConfig {
-    pub fn broadcast_update_param(&self, name: String) {
-        for i in 0..self.clients.len() {
-            let _ = self.clients.get(i).unwrap().broadcast_tx.send(name.clone());
-        }
-    }
-}
-
-pub struct UnshellServer {
-    config: Arc<Mutex<UnshellServerConfig>>,
-}
-
-impl UnshellServer {
-    pub fn from_filepath(filepath: &str) -> Self {
-        let s = (|| -> Result<Self, Box<dyn Error>> {
-            let mut file = File::open(filepath.to_string())?;
-
-            let mut contents = String::new();
-            file.read_to_string(&mut contents)?;
-
-            let config = serde_json::from_str::<UnshellServerConfig>(contents.as_str())?;
-
-            info!("Loaded server config from {}", filepath);
-
-            Ok(Self {
-                config: Arc::new(Mutex::new(config)),
-            })
-        })()
-        .unwrap_or({
-            warn!("Loaded default server config");
-            Self {
-                config: Arc::new(Mutex::new(UnshellServerConfig {
-                    campaign: DEFAULT_CAMPAIGN.clone(),
-                    users: DEFAULT_USERS.clone(),
-                    parameters: DEFAULT_PARAMETERS.clone(),
-
-                    clients: Vec::new(),
-                })),
-            }
-        });
-
-        s
-    }
-
-    pub fn run(&mut self, addr: SocketAddr) -> Result<(), Box<dyn Error>> {
-        let config_clone = Arc::clone(&self.config);
-        run_listener_state(TCPServer::bind(&addr)?, Client::run, config_clone);
-
-        Ok(())
-    }
-}
-
-/// Remote client type for unshell parameters
-struct Client {
-    pub broadcast_tx: Sender<String>,
-}
-
-impl Client {
-    pub fn run(connection: TCPConnection, config: Arc<Mutex<UnshellServerConfig>>) {
-        let (tx, rx) = TCPConnection::as_async::<C2Packet>(connection);
-
-        let (broadcast_tx, broadcast_rx) = crossbeam_channel::unbounded::<String>();
-
-        let s = Self { broadcast_tx };
-
-        let mut config_lock = config.lock().unwrap();
-        config_lock.clients.push(s);
-        let config_clone = Arc::clone(&config);
-        let tx_clone = tx.clone();
-        thread::spawn(move || {
-            loop {
-                if let Ok(key) = broadcast_rx.recv() {
-                    let config_lock = config_clone.lock().unwrap();
-                    if let Err(e) = tx_clone.send(C2Packet::ParameterUpate(
-                        key.clone(),
-                        config_lock.parameters.get(&key).unwrap().clone(),
-                    )) {
-                        error!("Failed to send packet: {}", e);
-                    };
-                }
-            }
-        });
-
-        if let Err(e) = tx.send(C2Packet::SetAllParameters(config_lock.parameters.clone())) {
-            error!("Failed to send packet: {}", e);
-        };
-        std::mem::drop(config_lock);
-
-        thread::spawn(move || {
-            loop {
-                // if !connection.is_alive() {
-                //     warn!("Client {} disconnected!", connection.get_info());
-                //     break;
-                // }
-                if let Ok(packet) = rx.recv() {
-                    if let Err(e) = match packet {
-                        C2Packet::GetParameter(param) => {
-                            tx.send(C2Packet::AckGetParameter(param.clone(), {
-                                let config_lock = config.lock().unwrap();
-                                let result = config_lock.parameters.get(&param);
-                                result.cloned()
-                            }))
-                        }
-                        C2Packet::SetParameter(name, param) => {
-                            tx.send(C2Packet::AckSetParameter({
-                                let mut config_lock = config.lock().unwrap();
-                                config_lock.parameters.insert(name.clone(), param);
-                                config_lock.broadcast_update_param(name);
-                                true
-                            }))
-                        }
-
-                        C2Packet::Error(error) => {
-                            warn!("Got error: {:?}", error);
-                            Ok(())
-                        }
-                        _ => tx.send(C2Packet::Error(ErrorPacket::UnsupportedRequestError)),
-                    } {
-                        error!("Failed to send packet: {}", e);
-                    }
-                }
-            }
-        });
-    }
-}

unshell-rs-lib/Cargo.toml

@@ -6,6 +6,5 @@ edition = "2024"
 base64 = "0.22.1"
 crossbeam-channel = "0.5.15"
 log = "0.4.27"
-mio = { version = "1.0.4", features = ["os-poll"] }
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"

unshell-rs-lib/src/config/campaign.rs

@@ -1,9 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-use crate::config::listeners::ListenerConfig;
-
-#[derive(Serialize, Deserialize, Clone)]
-pub struct CampaignConfig {
-    pub name: String,
-    pub listeners: Vec<ListenerConfig>,
-}

unshell-rs-lib/src/config/layers.rs

@@ -1,4 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Serialize, Deserialize, Clone)]
-pub enum LayerConfig {}

unshell-rs-lib/src/config/listeners.rs

@@ -1,47 +0,0 @@
-use std::{
-    error::Error,
-    net::SocketAddr,
-    sync::{Arc, Mutex},
-};
-
-use serde::{Deserialize, Serialize};
-
-use crate::{
-    config::layers::LayerConfig,
-    networkers::{ServerTrait, TCPConnection, TCPServer},
-};
-
-#[derive(Serialize, Deserialize, Clone)]
-pub enum ListenerConfig {
-    Tcp {
-        enabled: bool,
-        name: String,
-        addr: SocketAddr,
-        layers: Vec<LayerConfig>,
-
-        #[serde(skip)]
-        connections: Option<Arc<Mutex<Vec<TCPConnection>>>>,
-    },
-}
-
-impl ListenerConfig {
-    pub fn start(self) -> Result<(), Box<dyn Error>> {
-        match self {
-            ListenerConfig::Tcp {
-                mut enabled,
-                addr,
-                layers,
-                mut connections,
-                ..
-            } => {
-                let server = TCPServer::bind(&addr)?;
-
-                enabled = true;
-
-                // connections = Some(run_listener(server));
-            }
-        }
-
-        Ok(())
-    }
-}

unshell-rs-lib/src/config/mod.rs

@@ -1,3 +0,0 @@
-pub mod campaign;
-pub mod layers;
-pub mod listeners;

unshell-rs-lib/src/connection/mod.rs

@@ -1,6 +1,2 @@
-mod packets;
-
-pub use packets::C2Packet;
-pub use packets::ErrorPacket;
-pub use packets::Parameter;
-pub use packets::Parameters;
+mod node;
+pub use node::Node;

unshell-rs-lib/src/connection/node.rs

@@ -0,0 +1,37 @@
+use std::{net::SocketAddr, thread};
+
+use crate::{
+    Error,
+    layers::LayerConfig,
+    networkers::{Connection, ServerTrait, TCPConnection, TCPServer, run_listener},
+};
+
+pub struct Node;
+
+impl Node {
+    pub fn run(addr: SocketAddr) -> Result<(), Error> {
+        let layers = vec![LayerConfig::Handshake, LayerConfig::Base64];
+
+        run_listener(
+            TCPServer::bind(&addr)?,
+            layers,
+            |connection: Box<dyn Connection + Send + 'static>| {
+                thread::spawn(move || {
+                    let mut connection = connection;
+
+                    loop {
+                        if let Ok(data) = connection.read() {
+                            if !connection.is_alive() {
+                                warn!("{} Disconnected!", connection.get_info());
+                                break;
+                            }
+                            println!("Data: {}", data);
+                        }
+                    }
+                });
+            },
+        );
+
+        Ok(())
+    }
+}

unshell-rs-lib/src/connection/packets.rs

@@ -1,60 +0,0 @@
-use std::{collections::HashMap, fmt};
-
-use serde::{Deserialize, Serialize};
-use serde_json::Result;
-
-use crate::config::campaign::CampaignConfig;
-
-#[derive(Serialize, Deserialize, Debug)]
-pub enum C2Packet {
-    GetClients,
-    AckGetClients,
-
-    RequestCampaign,
-    AckRequestCampaign(CampaignConfig),
-
-    SetCampaign(CampaignConfig),
-    AckSetCampaign,
-
-    GetParameter(String),
-    AckGetParameter(String, Option<Parameter>),
-    ParameterUpate(String, Parameter),
-
-    SetParameter(String, Parameter),
-    AckSetParameter(bool),
-
-    SetAllParameters(Parameters),
-
-    Error(ErrorPacket),
-
-    Sysinfo { hostname: String },
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub enum ErrorPacket {
-    UnsupportedRequestError,
-}
-
-impl fmt::Debug for CampaignConfig {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "CampaignConfig")
-    }
-}
-
-pub type Parameters = HashMap<String, Parameter>;
-
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub enum Parameter {
-    Test1,
-    CurrentTab(i32),
-}
-
-impl C2Packet {
-    pub fn encode(&self) -> Result<String> {
-        serde_json::to_string(self)
-    }
-
-    pub fn decode(string: &str) -> Result<Self> {
-        serde_json::from_str::<Self>(string)
-    }
-}

unshell-rs-lib/src/layers/base64.rs

@@ -1,18 +1,44 @@
-use crate::layers::Layer;
-use base64;
+use crate::{
+    Error,
+    networkers::{Connection, ProtocolLayer},
+};
+use base64::{Engine as _, engine::general_purpose};
 use serde::{Deserialize, Serialize};
 
-#[derive(Default, Serialize, Deserialize)]
-pub struct Base64;
-
-impl Layer for Base64 {
-    fn encode(&mut self, data: &[u8]) -> Vec<u8> {
-        #[allow(deprecated)]
-        base64::encode(str::from_utf8(data).unwrap()).into_bytes()
+pub struct Base64Layer<C: Connection> {
+    inner: C,
 }
 
-    fn decode(&mut self, data: &[u8]) -> Vec<u8> {
-        #[allow(deprecated)]
-        base64::decode(str::from_utf8(data).unwrap()).unwrap()
+impl<C: Connection> Connection for Base64Layer<C> {
+    fn get_info(&self) -> String {
+        format!("b64->{}", self.inner.get_info())
+    }
+
+    fn is_alive(&self) -> bool {
+        self.inner.is_alive()
+    }
+
+    fn read(&mut self) -> Result<String, Error> {
+        Ok(str::from_utf8(
+            &general_purpose::STANDARD
+                .decode(&self.inner.read()?)
+                .unwrap(),
+        )
+        .unwrap()
+        .to_string())
+    }
+
+    fn write(&mut self, data: &str) -> Result<(), Error> {
+        info!("Bsae");
+
+        self.inner.write(&general_purpose::STANDARD.encode(data))?;
+
+        Ok(())
+    }
+}
+
+impl<C: Connection> ProtocolLayer<C> for Base64Layer<C> {
+    fn new(inner: C) -> Result<Self, Error> {
+        Ok(Base64Layer { inner })
     }
 }

unshell-rs-lib/src/layers/builder.rs

@@ -0,0 +1,69 @@
+use crate::{
+    Error,
+    layers::{Base64Layer, HandshakeLayer, LayerConfig},
+    networkers::{Connection, ProtocolLayer},
+};
+
+impl Connection for Box<dyn Connection + Send> {
+    fn get_info(&self) -> String {
+        (**self).get_info()
+    }
+
+    fn is_alive(&self) -> bool {
+        (**self).is_alive()
+    }
+
+    fn read(&mut self) -> Result<String, Error> {
+        (**self).read()
+    }
+
+    fn write(&mut self, data: &str) -> Result<(), Error> {
+        (**self).write(data)
+    }
+}
+
+pub fn build_client<C>(base_conn: C, layers: Vec<LayerConfig>) -> Result<Box<dyn Connection>, Error>
+where
+    C: Connection + 'static,
+{
+    let mut current_conn: Box<dyn Connection + Send> = Box::new(base_conn);
+
+    for layer_config in &layers {
+        current_conn = match layer_config {
+            LayerConfig::Base64 => Box::new(Base64Layer::new(current_conn)?),
+            LayerConfig::Handshake => {
+                let mut handshake_layer = HandshakeLayer::new(current_conn)?;
+                handshake_layer.initialize_client()?;
+                Box::new(handshake_layer)
+            }
+        };
+    }
+
+    Ok(current_conn)
+}
+
+pub fn create_server_builder<C>(
+    layers: Vec<LayerConfig>,
+) -> Result<Box<dyn Fn(C) -> Result<Box<dyn Connection + Send>, Error>>, Error>
+where
+    C: Connection + 'static,
+{
+    Ok(Box::new(
+        move |base_conn: C| -> Result<Box<dyn Connection + Send>, Error> {
+            let mut current_conn: Box<dyn Connection + Send> = Box::new(base_conn);
+
+            for layer_config in &layers {
+                current_conn = match layer_config {
+                    LayerConfig::Base64 => Box::new(Base64Layer::new(current_conn)?),
+                    LayerConfig::Handshake => {
+                        let mut handshake_layer = HandshakeLayer::new(current_conn)?;
+                        handshake_layer.initialize_server()?;
+                        Box::new(handshake_layer)
+                    }
+                };
+            }
+
+            Ok(current_conn)
+        },
+    ))
+}

unshell-rs-lib/src/layers/handshake.rs

@@ -0,0 +1,103 @@
+use crate::{
+    layers::Base64Layer,
+    networkers::{Connection, ProtocolLayer},
+};
+
+type Error = Box<dyn std::error::Error>;
+
+// 4-Way Handshake Layer
+pub struct HandshakeLayer<C: Connection> {
+    inner: C,
+    finished_handshake: bool,
+}
+
+impl<C: Connection> Connection for HandshakeLayer<C> {
+    fn get_info(&self) -> String {
+        format!("handshake->{}", self.inner.get_info())
+    }
+
+    fn is_alive(&self) -> bool {
+        self.inner.is_alive()
+    }
+
+    fn read(&mut self) -> Result<String, Error> {
+        if !self.finished_handshake {
+            return Err("NotComplete".into());
+        }
+        self.inner.read()
+    }
+
+    fn write(&mut self, data: &str) -> Result<(), Error> {
+        if !self.finished_handshake {
+            return Err("NotComplete".into());
+        }
+        self.inner.write(data)
+    }
+}
+
+impl<C: Connection + 'static> ProtocolLayer<C> for HandshakeLayer<C> {
+    fn new(inner: C) -> Result<Self, Error> {
+        Ok(HandshakeLayer {
+            inner,
+            finished_handshake: false,
+        })
+    }
+
+    fn initialize_client(&mut self) -> Result<(), Error> {
+        println!("Starting client handshake...");
+
+        // Step 1: Client sends SYN
+        self.inner.write("SYN")?;
+        println!("Client: Sent SYN");
+
+        // Step 2: Client receives SYN-ACK
+        let response = self.inner.read()?;
+        if response != "SYN-ACK" {
+            return Err(format!("Expected SYN-ACK, got: {}", response).into());
+        }
+        println!("Client: Received SYN-ACK");
+
+        // Step 3: Client sends ACK
+        self.inner.write("ACK")?;
+        println!("Client: Sent ACK");
+
+        // Step 4: Client receives FIN (final confirmation)
+        let response = self.inner.read()?;
+        if response != "FIN" {
+            return Err(format!("Expected FIN, got: {}", response).into());
+        }
+        println!("Client: Received FIN - Handshake complete!");
+
+        self.finished_handshake = true;
+        Ok(())
+    }
+
+    fn initialize_server(&mut self) -> Result<(), Error> {
+        println!("Starting server handshake...");
+
+        // Step 1: Server receives SYN
+        let request = self.inner.read()?;
+        if request != "SYN" {
+            return Err(format!("Expected SYN, got: {}", request).into());
+        }
+        println!("Server: Received SYN");
+
+        // Step 2: Server sends SYN-ACK
+        self.inner.write("SYN-ACK")?;
+        println!("Server: Sent SYN-ACK");
+
+        // Step 3: Server receives ACK
+        let response = self.inner.read()?;
+        if response != "ACK" {
+            return Err(format!("Expected ACK, got: {}", response).into());
+        }
+        println!("Server: Received ACK");
+
+        // Step 4: Server sends FIN (final confirmation)
+        self.inner.write("FIN")?;
+        println!("Server: Sent FIN - Handshake complete!");
+
+        self.finished_handshake = true;
+        Ok(())
+    }
+}

unshell-rs-lib/src/layers/mod.rs

@@ -1,9 +1,14 @@
-pub trait Layer: Serialize + Deserialize<'static> + Sized {
-    fn encode(&mut self, data: &[u8]) -> Vec<u8>;
-    fn decode(&mut self, data: &[u8]) -> Vec<u8>;
+pub enum LayerConfig {
+    Base64,
+    Handshake,
 }
 
 pub mod base64;
+mod builder;
+pub mod handshake;
 
-pub use base64::Base64;
-use serde::{Deserialize, Serialize};
+pub use base64::Base64Layer;
+pub use handshake::HandshakeLayer;
+
+pub use builder::build_client;
+pub use builder::create_server_builder;

unshell-rs-lib/src/lib.rs

@@ -1,7 +1,9 @@
 #[macro_use]
 extern crate log;
 
-pub mod config;
+pub type Error = Box<dyn std::error::Error>;
+
+// pub mod config;
 pub mod connection;
 pub mod layers;
 pub mod networkers;

unshell-rs-lib/src/networkers/mod.rs

@@ -1,115 +1,15 @@
-/// This is the lowset-level data transmission type
-
-pub trait Connection: Send + Sync {
-    type Error: std::fmt::Debug;
-
-    fn get_info(&self) -> String;
-    fn is_alive(&self) -> bool;
-
-    fn read(&mut self) -> Result<String, Self::Error>;
-    fn write(&mut self, data: &str) -> Result<(), Self::Error>;
-}
-
-pub trait AsyncConnection<C>
-where
-    C: Connection,
-{
-    type Error: std::fmt::Debug;
-
-    fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
-        connection: C,
-    ) -> (Sender<T>, Receiver<T>);
-}
-
-pub trait ServerTrait<C: Connection> {
-    type Error: std::fmt::Debug;
-
-    fn get_info(&self) -> String;
-    fn accept(&self) -> Result<C, Self::Error>;
-    fn bind(address: &SocketAddr) -> Result<Self, Self::Error>
-    where
-        Self: Sized;
-}
-
-pub trait ClientTrait<C: Connection + Sized> {
-    type Error: std::fmt::Debug;
-
-    fn connect(address: &SocketAddr) -> Result<C, Self::Error>;
-}
-
-pub fn run_listener_state<S, C, R, A>(server: S, on_connect_callback: R, state: Arc<A>)
-/*-> Arc<Mutex<Vec<C>>>*/
-where
-    S: ServerTrait<C> + Sync + Send + 'static,
-    C: Connection + 'static,
-    R: Fn(C, Arc<A>) + Sync + Send + 'static,
-    A: Sync + Send + 'static,
-{
-    info!("Started listener {}", server.get_info());
-    // let clients: Arc<Mutex<Vec<C>>> = Arc::new(Mutex::new(Vec::new()));
-    // let clients_clone = Arc::clone(&clients);
-
-    thread::spawn(move || {
-        loop {
-            match server.accept() {
-                Ok(conn) => {
-                    info!("New connection ({})", conn.get_info());
-
-                    on_connect_callback(conn, Arc::clone(&state));
-
-                    // OnConnectCallback::on_connect(&mut on_connect_callback, conn);
-                    // let mut clients_lock = clients_clone.lock().unwrap();
-                    // clients_lock.push(conn);
-                }
-                Err(e) => {
-                    error!("Failed to accept connection: {:?}", e);
-                }
-            }
-        }
-    });
-}
-
-pub fn run_listener<S, C, R>(server: S, on_connect_callback: R)
-/*-> Arc<Mutex<Vec<C>>>*/
-where
-    S: ServerTrait<C> + Sync + Send + 'static,
-    C: Connection + 'static,
-    R: Fn(C) + Sync + Send + 'static,
-{
-    info!("Started listener {}", server.get_info());
-    // let clients: Arc<Mutex<Vec<C>>> = Arc::new(Mutex::new(Vec::new()));
-    // let clients_clone = Arc::clone(&clients);
-
-    thread::spawn(move || {
-        loop {
-            match server.accept() {
-                Ok(conn) => {
-                    info!("New connection ({})", conn.get_info());
-
-                    on_connect_callback(conn);
-
-                    // OnConnectCallback::on_connect(&mut on_connect_callback, conn);
-                    // let mut clients_lock = clients_clone.lock().unwrap();
-                    // clients_lock.push(conn);
-                }
-                Err(e) => {
-                    error!("Failed to accept connection: {:?}", e);
-                }
-            }
-        }
-    });
-}
-
+mod server;
 mod tcp;
+mod traits;
 
-use std::net::SocketAddr;
-use std::sync::Arc;
-use std::thread;
-
-use crossbeam_channel::Receiver;
-use crossbeam_channel::Sender;
-use serde::Serialize;
-use serde::de::DeserializeOwned;
 pub use tcp::TCPClient;
 pub use tcp::TCPConnection;
 pub use tcp::TCPServer;
+
+// pub use traits::AsyncConnection;
+pub use traits::ClientTrait;
+pub use traits::Connection;
+pub use traits::ProtocolLayer;
+pub use traits::ServerTrait;
+
+pub use server::run_listener;

unshell-rs-lib/src/networkers/server.rs

@@ -0,0 +1,81 @@
+use std::sync::Arc;
+use std::thread;
+
+use crate::{
+    layers::{LayerConfig, create_server_builder},
+    networkers::{Connection, ServerTrait},
+};
+
+// Helper macros for building layered connections
+macro_rules! build_layered_connection {
+    ($base:expr) => {
+        $base
+    };
+    ($base:expr, $layer:ty) => {
+        <$layer>::new($base)?
+    };
+    ($base:expr, $layer:ty, $($layers:ty),+) => {
+        build_layered_connection!(<$layer>::new($base)?, $($layers),+)
+    };
+}
+
+pub fn run_listener_state<S, C, R, A>(server: S, on_connect_callback: R, state: Arc<A>)
+/*-> Arc<Mutex<Vec<C>>>*/
+where
+    S: ServerTrait<C> + Sync + Send + 'static,
+    C: Connection + 'static,
+    R: Fn(C, Arc<A>) + Sync + Send + 'static,
+    A: Sync + Send + 'static,
+{
+    info!("Started listener {}", server.get_info());
+    // let clients: Arc<Mutex<Vec<C>>> = Arc::new(Mutex::new(Vec::new()));
+    // let clients_clone = Arc::clone(&clients);
+
+    loop {
+        match server.accept() {
+            Ok(conn) => {
+                info!("New connection ({})", conn.get_info());
+
+                on_connect_callback(conn, Arc::clone(&state));
+
+                // OnConnectCallback::on_connect(&mut on_connect_callback, conn);
+                // let mut clients_lock = clients_clone.lock().unwrap();
+                // clients_lock.push(conn);
+            }
+            Err(e) => {
+                error!("Failed to accept connection: {:?}", e);
+            }
+        }
+    }
+}
+
+pub fn run_listener<S, C, R>(server: S, layers: Vec<LayerConfig>, on_connect_callback: R)
+/*-> Arc<Mutex<Vec<C>>>*/
+where
+    S: ServerTrait<C> + Sync + Send + 'static,
+    C: Connection + 'static,
+    R: Fn(Box<dyn Connection + Send + 'static>) + Sync + Send + 'static,
+{
+    let layer_builder = create_server_builder::<C>(layers).unwrap();
+
+    info!("Started listener {}", server.get_info());
+    // let clients: Arc<Mutex<Vec<C>>> = Arc::new(Mutex::new(Vec::new()));
+    // let clients_clone = Arc::clone(&clients);
+
+    loop {
+        match server.accept() {
+            Ok(conn) => match layer_builder(conn) {
+                Ok(conn) => {
+                    info!("New connection ({})", conn.get_info());
+                    on_connect_callback(conn);
+                }
+                Err(e) => {
+                    error!("Failed to create layers: {:?}", e);
+                }
+            },
+            Err(e) => {
+                error!("Failed to accept connection: {:?}", e);
+            }
+        }
+    }
+}

unshell-rs-lib/src/networkers/tcp.rs

@@ -1,13 +1,12 @@
 use std::{
     io::{self, BufRead, BufReader, Write},
     net::{SocketAddr, TcpListener, TcpStream},
-    thread,
 };
 
-use crossbeam_channel::{Receiver, Sender};
-use serde::{Serialize, de::DeserializeOwned};
-
-use crate::networkers::{AsyncConnection, ClientTrait, Connection, ServerTrait};
+use crate::{
+    Error,
+    networkers::{ClientTrait, Connection, ServerTrait},
+};
 
 pub struct TCPConnection {
     stream: TcpStream,
@@ -16,8 +15,6 @@ pub struct TCPConnection {
 }
 
 impl Connection for TCPConnection {
-    type Error = io::Error;
-
     fn get_info(&self) -> String {
         format!(
             "tcp://{}",
@@ -33,7 +30,7 @@ impl Connection for TCPConnection {
         self.is_alive
     }
 
-    fn read(&mut self) -> Result<String, Self::Error> {
+    fn read(&mut self) -> Result<String, Error> {
         let mut line = String::new();
         let n = self.reader.read_line(&mut line)?;
 
@@ -45,79 +42,78 @@ impl Connection for TCPConnection {
         Ok(line.trim_end().to_string())
     }
 
-    fn write(&mut self, data: &str) -> Result<(), Self::Error> {
+    fn write(&mut self, data: &str) -> Result<(), Error> {
+        info!("Sent: {}", data);
         writeln!(self.stream, "{}", data)?;
         self.stream.flush()?;
         Ok(())
     }
 }
 
-impl AsyncConnection<TCPConnection> for TCPConnection {
-    type Error = io::Error;
+// impl AsyncConnection<TCPConnection> for TCPConnection {
+//     type Error = io::Error;
 
-    fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
-        connection: TCPConnection,
-    ) -> (Sender<T>, Receiver<T>) {
-        let (send_tx, send_rx) = crossbeam_channel::unbounded::<T>();
-        let (recv_tx, recv_rx) = crossbeam_channel::unbounded::<T>();
+//     fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
+//         connection: TCPConnection,
+//     ) -> (Sender<T>, Receiver<T>) {
+//         let (send_tx, send_rx) = crossbeam_channel::unbounded::<T>();
+//         let (recv_tx, recv_rx) = crossbeam_channel::unbounded::<T>();
 
-        thread::spawn(move || {
-            let mut reader = connection.reader;
+//         thread::spawn(move || {
+//             let mut reader = connection.reader;
 
-            let mut read = || -> Result<String, Self::Error> {
-                let mut line = String::new();
-                let _ = reader.read_line(&mut line)?;
+//             let mut read = || -> Result<String, Self::Error> {
+//                 let mut line = String::new();
+//                 let _ = reader.read_line(&mut line)?;
 
-                Ok(line.trim_end().to_string())
-            };
+//                 Ok(line.trim_end().to_string())
+//             };
 
-            loop {
-                if let Ok(data) = read() {
-                    if data.is_empty() {
-                        break;
-                    }
-                    info!("Got {}", data);
-                    if let Ok(decoded) = serde_json::from_str::<T>(&data) {
-                        if let Err(e) = send_tx.send(decoded) {
-                            error!("Got error: {}", e);
-                        }
-                    }
-                }
-            }
-        });
+//             loop {
+//                 if let Ok(data) = read() {
+//                     if data.is_empty() {
+//                         break;
+//                     }
+//                     info!("Got {}", data);
+//                     if let Ok(decoded) = serde_json::from_str::<T>(&data) {
+//                         if let Err(e) = send_tx.send(decoded) {
+//                             error!("Got error: {}", e);
+//                         }
+//                     }
+//                 }
+//             }
+//         });
 
-        thread::spawn(move || {
-            let mut stream = connection.stream;
+//         thread::spawn(move || {
+//             let mut stream = connection.stream;
 
-            let mut write = |data: String| -> Result<(), Self::Error> {
-                writeln!(stream, "{}", data)?;
-                stream.flush()?;
-                Ok(())
-            };
+//             let mut write = |data: String| -> Result<(), Self::Error> {
+//                 writeln!(stream, "{}", data)?;
+//                 stream.flush()?;
+//                 Ok(())
+//             };
 
-            loop {
-                if let Ok(data) = recv_rx.recv() {
-                    if let Ok(encoded) = serde_json::to_string(&data) {
-                        info!("Write {}", encoded);
-                        if let Err(e) = write(encoded) {
-                            error!("Got error: {}", e);
-                        }
-                    }
-                }
-            }
-        });
+//             loop {
+//                 if let Ok(data) = recv_rx.recv() {
+//                     if let Ok(encoded) = serde_json::to_string(&data) {
+//                         info!("Write {}", encoded);
+//                         if let Err(e) = write(encoded) {
+//                             error!("Got error: {}", e);
+//                         }
+//                     }
+//                 }
+//             }
+//         });
 
-        (recv_tx, send_rx)
-    }
-}
+//         (recv_tx, send_rx)
+//     }
+// }
 
 pub struct TCPServer {
     listener: TcpListener,
 }
 
 impl ServerTrait<TCPConnection> for TCPServer {
-    type Error = io::Error;
-
     fn get_info(&self) -> String {
         format!(
             "tcp://{}",
@@ -129,7 +125,7 @@ impl ServerTrait<TCPConnection> for TCPServer {
         )
     }
 
-    fn accept(&self) -> Result<TCPConnection, Self::Error> {
+    fn accept(&self) -> Result<TCPConnection, Error> {
         let (stream, _) = self.listener.accept()?;
         let reader = BufReader::new(stream.try_clone()?);
         Ok(TCPConnection {
@@ -139,7 +135,7 @@ impl ServerTrait<TCPConnection> for TCPServer {
         })
     }
 
-    fn bind(address: &SocketAddr) -> Result<Self, Self::Error> {
+    fn bind(address: &SocketAddr) -> Result<Self, Error> {
         let listener = TcpListener::bind(address)?;
         Ok(Self { listener })
     }
@@ -148,9 +144,7 @@ impl ServerTrait<TCPConnection> for TCPServer {
 pub struct TCPClient;
 
 impl ClientTrait<TCPConnection> for TCPClient {
-    type Error = io::Error;
-
-    fn connect(address: &SocketAddr) -> Result<TCPConnection, Self::Error> {
+    fn connect(address: &SocketAddr) -> Result<TCPConnection, Error> {
         let stream = TcpStream::connect(address)?;
         let reader = BufReader::new(stream.try_clone()?);
         let conn = TCPConnection {
@@ -158,7 +152,6 @@ impl ClientTrait<TCPConnection> for TCPClient {
             reader,
             is_alive: true,
         };
-        info!("Connected to {}", conn.get_info());
         Ok(conn)
     }
 }

unshell-rs-lib/src/networkers/traits.rs

@@ -0,0 +1,50 @@
+use std::net::SocketAddr;
+use std::ops::Deref;
+use std::ops::DerefMut;
+
+use crate::Error;
+
+// This is the lowset-level data transmission type
+pub trait Connection: Send {
+    fn get_info(&self) -> String;
+    fn is_alive(&self) -> bool;
+
+    fn read(&mut self) -> Result<String, Error>;
+    fn write(&mut self, data: &str) -> Result<(), Error>;
+}
+
+// Trait for protocol layers that can be initialized
+pub trait ProtocolLayer<C: Connection>: Connection {
+    fn new(inner: C) -> Result<Self, Error>
+    where
+        Self: Sized;
+    fn initialize_client(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+    fn initialize_server(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+// impl Sized for dyn Connection {}
+
+// pub trait AsyncConnection<C>
+// where
+//     C: Connection,
+// {
+//     fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
+//         connection: C,
+//     ) -> (Sender<T>, Receiver<T>);
+// }
+
+pub trait ServerTrait<C: Connection> {
+    fn get_info(&self) -> String;
+    fn accept(&self) -> Result<C, Error>;
+    fn bind(address: &SocketAddr) -> Result<Self, Error>
+    where
+        Self: Sized;
+}
+
+pub trait ClientTrait<C: Connection + Sized> {
+    fn connect(address: &SocketAddr) -> Result<C, Error>;
+}