Work on server and client connectivity

Cargo.toml

@@ -1,11 +1,13 @@
 [package]
-name = "server"
+name = "unshell-rs"
 version = "0.1.0"
 edition = "2024"
 
 [dependencies]
 clap = { version = "4.5.39", features = ["derive"] }
+lazy_static = "1.5.0"
 log = "0.4.27"
+native-tls = "0.2.14"
 pretty_env_logger = "0.5.0"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"

payload/src/main.rs

@@ -14,53 +14,6 @@ use unshell_rs_lib::{
     packets::Packet,
 };
 
-// Generic client function
-pub fn run_client<C, Cl>(address: &str) -> Result<(), Box<dyn std::error::Error>>
-where
-    Cl: ClientTrait<C>,
-    C: Connection + 'static,
-    Cl::Error: std::error::Error + 'static,
-    C::Error: std::error::Error + 'static,
-{
-    let recv_conn = Arc::new(Mutex::new(Cl::connect(address)?));
-    let transmit_vec: Arc<Mutex<Vec<Packet>>> = Arc::new(Mutex::new(Vec::new()));
-
-    let transmit_conn = Arc::clone(&recv_conn);
-    let transmit_vec_clone = Arc::clone(&transmit_vec);
-
-    thread::spawn(move || {
-        loop {
-            let mut transmit_vec_lock = transmit_vec.lock().unwrap();
-            if transmit_vec_lock.len() > 0 {
-                let mut conn_lock = recv_conn.lock().unwrap();
-                if let Ok(json) = serde_json::to_string(&transmit_vec_lock.pop().unwrap()) {
-                    conn_lock.write(&json).expect("Failed to send packet!");
-                }
-            } else {
-                thread::sleep(Duration::from_millis(10));
-            }
-        }
-    });
-
-    loop {
-        let mut conn_lock = transmit_conn.lock().unwrap();
-        let data = conn_lock.read();
-        drop(conn_lock);
-        match data {
-            Ok(data_json) => {
-                if data_json.is_empty() {
-                    continue;
-                }
-                let packet = serde_json::from_str::<Packet>(data_json.as_str());
-                println!("{:?}", packet);
-            }
-            Err(e) => {
-                eprintln!("Error reading, {}", e);
-            }
-        }
-    }
-}
-
 fn main() -> Result<(), Box<dyn Error>> {
     run_client::<TCPConnection, TCPClient>("127.0.0.1:3000")?;
 

src/client/client.rs

@@ -0,0 +1,120 @@
+use std::{
+    error::Error,
+    mem,
+    net::SocketAddr,
+    sync::{Arc, Mutex},
+    thread,
+    time::Duration,
+};
+
+use unshell_rs_lib::networkers::{ClientTrait, Connection, TCPClient, TCPConnection};
+
+use crate::packets::{GuiPacket, Parameter, Parameters};
+
+pub struct UnshellClient {
+    addr: SocketAddr,
+    client: Arc<Mutex<TCPConnection>>,
+
+    parameters: Arc<Mutex<Parameters>>,
+    outgoing_packets: Arc<Mutex<Vec<GuiPacket>>>,
+}
+
+impl UnshellClient {
+    pub fn new(addr: SocketAddr) -> Result<Self, Box<dyn Error>> {
+        let client = Arc::new(Mutex::new(TCPClient::connect(&addr)?));
+        let outgoing_packets = Arc::new(Mutex::new(Vec::<GuiPacket>::new()));
+        let parameters = Arc::new(Mutex::new(Parameters::new()));
+
+        let tx_client = Arc::clone(&client);
+        let tx_packets = Arc::clone(&outgoing_packets);
+
+        // Recieve thread
+        thread::spawn(move || {
+            loop {
+                info!("Lock 2");
+                let mut packets_lock = tx_packets.lock().unwrap();
+                info!("Lock 2");
+                if !packets_lock.is_empty() {
+                    info!("Lock 3");
+                    if let Ok(packet) = packets_lock.pop().unwrap().encode() {
+                        info!("Lock 3");
+                        let mut client_lock = tx_client.lock().unwrap();
+                        info!("Wrote {}", packet.as_str());
+                        match client_lock.write(packet.as_str()) {
+                            Err(e) => {
+                                error!("Failed to send packet: {:?}", e);
+                            }
+                            _ => {}
+                        };
+                    }
+                }
+                std::mem::drop(packets_lock);
+
+                thread::sleep(Duration::from_millis(10));
+            }
+        });
+
+        let rx_client = Arc::clone(&client);
+        let rx_params = Arc::clone(&parameters);
+        thread::spawn(move || {
+            loop {
+                info!("Lock 4");
+                let mut client = rx_client.lock().unwrap();
+                info!("Lock 4");
+                if !client.is_alive() {
+                    error!("Disconnected from {}!", client.get_info());
+                }
+                if let Ok(data) = client.read() {
+                    info!("Got {}", data);
+                    if let Ok(packet) = GuiPacket::decode(data.as_str()) {
+                        match packet {
+                            GuiPacket::ParameterUpate(name, parameter) => {
+                                rx_params.lock().unwrap().insert(name, parameter);
+                            }
+                            GuiPacket::Error(error_packet) => {
+                                error!("Got error: {}", print_type_of(&error_packet))
+                            }
+                            GuiPacket::SetAllParameters(parameters) => {
+                                let mut params_lock = rx_params.lock().unwrap();
+                                params_lock.clear();
+                                params_lock.extend(parameters);
+                            }
+                            _ => {
+                                error!("Unsupported packet: {}", data)
+                            }
+                        }
+                    }
+                }
+                // std::mem::drop(client);
+
+                thread::sleep(Duration::from_millis(10));
+            }
+        });
+
+        Ok(Self {
+            addr,
+            client,
+            parameters,
+            outgoing_packets,
+        })
+    }
+
+    pub fn set_parameter(&mut self, key: String, param: Parameter) {
+        self.parameters
+            .lock()
+            .unwrap()
+            .insert(key.clone(), param.clone());
+        self.outgoing_packets
+            .lock()
+            .unwrap()
+            .push(GuiPacket::SetParameter(key, param));
+    }
+
+    pub fn get_parameter(&self, key: &str) -> Option<Parameter> {
+        self.parameters.lock().unwrap().get(key).cloned()
+    }
+}
+
+fn print_type_of<T>(_: &T) -> &'static str {
+    std::any::type_name::<T>()
+}

src/client/gui.rs

@@ -0,0 +1,71 @@
+use slint::{ComponentHandle, ModelRc, VecModel};
+use std::{
+    error::Error,
+    sync::{Arc, Mutex},
+};
+use unshell_rs_lib::config::campaign::CampaignConfig;
+
+use crate::{client::UnshellClient, packets::Parameter};
+
+pub struct UnshellGui {
+    client: UnshellClient,
+    ui: AppWindow,
+    campaign: Option<Arc<Mutex<CampaignConfig>>>,
+}
+
+slint::include_modules!();
+impl UnshellGui {
+    pub fn start(client: UnshellClient) -> Result<(), Box<dyn Error>> {
+        let ui = AppWindow::new()?;
+        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);
+            info!("Lock 1 ");
+            let mut client_lock = client_clone.lock().unwrap();
+            info!("Lock 1 ");
+            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()
+        });
+
+        ui.run()?;
+
+        Ok(())
+    }
+
+    fn update(&mut self) {
+        // self.ui.set_listeners(ModelRc::new(VecModel::from(
+        //     self.campaign
+        //         .listeners
+        //         .iter()
+        //         .map(|l| match l {
+        //             ListenerConfig::Tcp {
+        //                 enabled,
+        //                 name,
+        //                 addr,
+        //                 layers,
+        //                 ..
+        //             } => UITcpListener {
+        //                 enabled: *enabled,
+        //                 name: name.clone().into(),
+        //                 remote_host: addr.to_string().into(),
+        //             },
+        //         })
+        //         .collect::<Vec<UITcpListener>>(),
+        // )));
+    }
+}
+
+// trait

src/client/mod.rs

@@ -0,0 +1,6 @@
+mod gui;
+
+mod client;
+
+pub use client::UnshellClient;
+pub use gui::UnshellGui;

src/gui.rs

@@ -1,62 +0,0 @@
-use slint::{ModelRc, VecModel};
-use std::error::Error;
-use unshell_rs_lib::config::listeners::ListenerConfig;
-
-pub struct Unshell_Gui;
-
-slint::include_modules!();
-impl Unshell_Gui {
-    pub fn start() -> Result<Self, Box<dyn Error>> {
-        let ui = AppWindow::new()?;
-
-        // ui.
-
-        let ui_handle = ui.as_weak();
-        ui.on_tab_clicked(move |index| {
-            let ui = ui_handle.unwrap();
-            ui.set_current_tab(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 listeners: Vec<ListenerConfig> = vec![ListenerConfig::Tcp {
-            enabled: true,
-            name: "test".to_string(),
-            remote_host: "127.0.0.1".to_string(),
-            port: 25565,
-            layers: Vec::new(),
-        }];
-
-        ui.set_listeners(ModelRc::new(VecModel::from(
-            listeners
-                .iter()
-                .map(|l| match l {
-                    ListenerConfig::Tcp {
-                        enabled,
-                        name,
-                        remote_host,
-                        port,
-                        layers,
-                    } => UITcpListener {
-                        enabled: *enabled,
-                        name: name.clone().into(),
-                        remote_host: remote_host.clone().into(),
-                        port: *port as i32,
-                    },
-                })
-                .collect::<Vec<UITcpListener>>(),
-        )));
-
-        ui.run()?;
-
-        Ok(Self {})
-    }
-}

src/lib.rs

@@ -1,5 +1,10 @@
 #[macro_use]
 extern crate log;
 
-mod gui;
+mod client;
+mod packets;
 mod server;
+
+pub use client::UnshellClient;
+pub use client::UnshellGui;
+pub use server::UnshellServer;

src/main.rs

@@ -1,19 +1,25 @@
-use std::error::Error;
+use std::{
-
+    env,
-use clap::{Parser, Subcommand};
+    error::Error,
-use log::trace;
+    net::{IpAddr, Ipv4Addr, SocketAddr},
-use slint::{ModelRc, VecModel};
+    str::FromStr,
-use unshell_rs_lib::{
-    config::listeners::ListenerConfig,
-    listeners::Listener,
-    networkers::{ServerTrait, TCPServer},
 };
 
-/// The default port that this program looks for
+use clap::{Parser, Subcommand};
+use log::error;
+use unshell_rs::{UnshellClient, UnshellGui, UnshellServer};
+// use unshell_rs
+
+pub static DEFAULT_CONFIG_FILEPATH: &'static str = "server_config.json";
+
+// The default port that this program looks for
 pub static DEFAULT_SERVICE_PORT: u16 = 13370;
-/// The default website port that this program looks for
+// The default website port that this program looks for
 pub static DEFAULT_WEB_PORT: u16 = 8082;
 
+pub static LOCAL_SOCKET: SocketAddr =
+    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 12, 34, 56)), 13370);
+
 #[derive(Debug, Parser)]
 #[command(name = "unshell-rs")]
 #[command(about = "Slick reverse shell tool in rust", long_about = None)]
@@ -26,48 +32,84 @@ struct Args {
 enum Commands {
     /// Run as a service, and potentially hosting a website
     #[command(arg_required_else_help = true)]
-    Serve {
+    Server {
-        /// Only listen for command clients locally
+        /// IPv4 to listen for clients on.
-        #[arg(short, long, default_value_t = false)]
+        host: String,
-        local: bool,
 
         /// Port listen to for command clients
         #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
-        service_port: u16,
+        port: u16,
+
+        /// Json file to store config
+        #[arg(short, long, default_value_t = DEFAULT_CONFIG_FILEPATH.to_string())]
+        config_filepath: String,
         // /// Port to listen for website traffic (0 is disabled)
         // #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
         // web_port: u16,
     },
-    Gui {
+    /// Run GUI and connect to remote server
-        /// Listen for command clients remotely aswell
+    Remote {
-        #[arg(short, long, default_value_t = true)]
+        /// Remote server to connect to
-        remote: bool,
+        host: String,
 
         /// Port listen to for command clients
         #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
-        service_port: u16,
+        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>> {
+    if env::var("RUST_LOG").is_err() {
+        unsafe { env::set_var("RUST_LOG", "info") }
+    }
+
     pretty_env_logger::init();
     let args = Args::parse();
 
     match args.command {
-        Commands::Gui {
+        Commands::Local { config_filepath } => {
-            remote,
+            let mut server = UnshellServer::from_filepath(config_filepath.as_str());
-            service_port,
+            server.run(LOCAL_SOCKET)?;
-        } => {}
-        Commands::Serve {
-            local,
-            service_port,
-            // web_port,
-        } => {}
-    }
 
-    // let mut server = Listener::new(TCPServer::bind("0.0.0.0:3000")?);
+            let client = UnshellClient::new(LOCAL_SOCKET)?;
 
-    // server.run_listener()?;
+            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 {
+            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)?;
+            } else {
+                error!("Could not parse address!");
+                return Ok(());
+            }
+
+            loop {}
+        }
+    };
 
     Ok(())
 }

src/packets/mod.rs

@@ -0,0 +1,37 @@
+use serde_json::Result;
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+use unshell_rs_lib::connection::ErrorPacket;
+
+pub type Parameters = HashMap<String, Parameter>;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub enum GuiPacket {
+    GetParameter(String),
+    AckGetParameter(String, Option<Parameter>),
+    ParameterUpate(String, Parameter),
+
+    SetParameter(String, Parameter),
+    AckSetParameter(bool),
+
+    SetAllParameters(Parameters),
+
+    Error(ErrorPacket),
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub enum Parameter {
+    Test1,
+    CurrentTab(i32),
+}
+
+impl GuiPacket {
+    pub fn encode(&self) -> Result<String> {
+        serde_json::to_string(self)
+    }
+
+    pub fn decode(string: &str) -> Result<Self> {
+        serde_json::from_str::<Self>(string)
+    }
+}

src/server/config.rs

@@ -0,0 +1,29 @@
+use lazy_static::lazy_static;
+use serde::{Deserialize, Serialize};
+use unshell_rs_lib::config::campaign::CampaignConfig;
+
+use std::collections::HashMap;
+
+use crate::packets::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

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

src/server/server.rs

@@ -0,0 +1,153 @@
+use std::{
+    collections::HashMap,
+    error::Error,
+    fs::File,
+    io::Read,
+    net::SocketAddr,
+    sync::{Arc, Mutex},
+    thread,
+    time::Duration,
+};
+
+use serde::{Deserialize, Serialize};
+
+use unshell_rs_lib::{
+    config::campaign::CampaignConfig,
+    connection::ErrorPacket,
+    networkers::{Connection, ServerTrait, TCPConnection, TCPServer, run_listener_state},
+};
+
+use crate::{
+    packets::{GuiPacket, Parameters},
+    server::{DEFAULT_CAMPAIGN, DEFAULT_USERS, User, config::DEFAULT_PARAMETERS},
+};
+
+#[derive(Serialize, Deserialize)]
+pub struct UnshellServerConfig {
+    campaign: CampaignConfig,
+    parameters: Parameters,
+    users: Vec<User>,
+
+    #[serde(skip)]
+    client_count: usize,
+    #[serde(skip)]
+    broadcast_flag: HashMap<usize, Option<String>>,
+}
+
+pub struct UnshellServer {
+    config: Arc<Mutex<UnshellServerConfig>>,
+}
+
+impl UnshellServer {
+    pub fn from_filepath(filepath: &str) -> Self {
+        (|| -> 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(),
+
+                    client_count: 0,
+                    broadcast_flag: HashMap::new(),
+                })),
+            }
+        })
+    }
+
+    pub fn run(&mut self, addr: SocketAddr) -> Result<(), Box<dyn Error>> {
+        let on_connect = |connection: TCPConnection,
+                          config_clone: Arc<Mutex<UnshellServerConfig>>| {
+            // Recieve loop
+            thread::spawn(move || {
+                let config = Arc::clone(&config_clone);
+
+                let mut connection = connection;
+
+                let send = |c: &mut TCPConnection, packet: GuiPacket| {
+                    if let Ok(packet) = packet.encode() {
+                        info!("Send {}", packet);
+                        c.write(packet.as_str()).unwrap();
+                    }
+                };
+
+                let mut config_lock = config.lock().unwrap();
+                let client_id = config_lock.client_count.clone();
+                config_lock.client_count += 1;
+                send(
+                    &mut connection,
+                    GuiPacket::SetAllParameters(config_lock.parameters.clone()),
+                );
+                std::mem::drop(config_lock);
+
+                loop {
+                    if !connection.is_alive() {
+                        warn!("Client {} disconnected!", connection.get_info());
+                        break;
+                    }
+                    if let Ok(data) = connection.read() {
+                        if let Ok(packet) = GuiPacket::decode(data.as_str()) {
+                            match packet {
+                                GuiPacket::GetParameter(param) => send(
+                                    &mut connection,
+                                    GuiPacket::AckGetParameter(param.clone(), {
+                                        let config_lock = config.lock().unwrap();
+                                        let result = config_lock.parameters.get(&param);
+                                        result.cloned()
+                                    }),
+                                ),
+                                GuiPacket::SetParameter(name, param) => send(
+                                    &mut connection,
+                                    GuiPacket::AckSetParameter({
+                                        let mut config_lock = config.lock().unwrap();
+                                        config_lock.parameters.insert(name.clone(), param);
+                                        config_lock.broadcast_flag.insert(client_id, Some(name));
+
+                                        true
+                                    }),
+                                ),
+                                _ => send(
+                                    &mut connection,
+                                    GuiPacket::Error(ErrorPacket::UnsupportedRequestError),
+                                ),
+                            }
+                        }
+                    }
+
+                    let mut config_lock = config.lock().unwrap();
+                    if let Some(Some(key)) = config_lock.broadcast_flag.get(&client_id) {
+                        send(
+                            &mut connection,
+                            GuiPacket::ParameterUpate(
+                                key.clone(),
+                                config_lock.parameters.get(key).unwrap().clone(),
+                            ),
+                        );
+                        config_lock.broadcast_flag.insert(client_id, None);
+                    }
+
+                    thread::sleep(Duration::from_millis(10));
+                }
+            });
+        };
+
+        let config_clone = Arc::clone(&self.config);
+        run_listener_state(TCPServer::bind(&addr)?, on_connect, config_clone);
+
+        Ok(())
+    }
+}

ui/pages/listeners.slint

@@ -33,16 +33,6 @@ component ListenerCard inherits BorderedRectangle {
                 font-size: 14px;
             }
         }
-
-        TitleText {
-            text: "Port: ";
-
-            Text {
-                text: listener.port;
-                font-weight: 500;
-                font-size: 14px;
-            }
-        }
     }
 }
 

ui/structs.slint

@@ -2,5 +2,4 @@ export struct UITcpListener {
     enabled: bool,
     name: string,
     remote_host: string,
-    port: int,
 }

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

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

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

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

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

@@ -1,14 +1,48 @@
+use std::{
+    error::Error,
+    net::SocketAddr,
+    sync::{Arc, Mutex},
+    thread,
+};
+
 use serde::{Deserialize, Serialize};
 
-use crate::config::layers::LayerConfig;
+use crate::{
+    config::layers::LayerConfig,
+    networkers::{ServerTrait, TCPConnection, TCPServer},
+};
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub enum ListenerConfig {
     Tcp {
         enabled: bool,
         name: String,
-        remote_host: String,
+        addr: SocketAddr,
-        port: u16,
         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/connection/mod.rs

@@ -0,0 +1,4 @@
+mod packets;
+
+pub use packets::C2Packet;
+pub use packets::ErrorPacket;

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

@@ -0,0 +1,43 @@
+use std::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,
+
+    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")
+    }
+}
+
+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/lib.rs

@@ -2,7 +2,6 @@
 extern crate log;
 
 pub mod config;
+pub mod connection;
 pub mod layers;
-pub mod listeners;
 pub mod networkers;
-pub mod packets;

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

@@ -1,9 +0,0 @@
-pub struct Client<C> {
-    pub stream: C,
-}
-
-impl<C> Client<C> {
-    pub fn new(stream: C) -> Self {
-        Self { stream }
-    }
-}

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

@@ -1,6 +0,0 @@
-use crate::layers::Layer;
-
-mod client;
-mod server;
-
-pub use server::Listener;

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

@@ -1,42 +0,0 @@
-use std::sync::{Arc, Mutex};
-
-use crate::{
-    listeners::client::Client,
-    networkers::{Connection, ServerTrait},
-};
-
-pub struct Listener<S, C> {
-    pub server: Arc<Mutex<S>>,
-    pub clients: Arc<Mutex<Vec<Client<C>>>>,
-}
-
-impl<S, C> Listener<S, C> {
-    pub fn new(server: S) -> Self {
-        Self {
-            server: Arc::new(Mutex::new(server)),
-            clients: Arc::new(Mutex::new(Vec::new())),
-        }
-    }
-
-    pub fn run_listener(&mut self) -> Result<(), Box<dyn std::error::Error>>
-    where
-        S: ServerTrait<C>,
-        C: Connection + 'static,
-        S::Error: std::error::Error + 'static,
-        C::Error: std::error::Error + 'static,
-    {
-        loop {
-            let mut conn_lock = self.server.lock().unwrap();
-
-            match conn_lock.accept() {
-                Ok(conn) => {
-                    let mut clients_lock = self.clients.lock().unwrap();
-                    clients_lock.push(Client::new(conn));
-                }
-                Err(e) => {
-                    error!("Failed to accept connection: {:?}", e);
-                }
-            }
-        }
-    }
-}

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

@@ -3,6 +3,9 @@
 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>;
 }
@@ -10,20 +13,88 @@ pub trait Connection: Send + Sync {
 pub trait ServerTrait<C: Connection> {
     type Error: std::fmt::Debug;
 
-    fn accept(&mut self) -> Result<C, Self::Error>;
+    fn get_info(&self) -> String;
-    fn bind(address: &str) -> Result<Self, Self::Error>
+    fn accept(&self) -> Result<C, Self::Error>;
+    fn bind(address: &SocketAddr) -> Result<Self, Self::Error>
     where
         Self: Sized;
 }
 
-pub trait ClientTrait<C: Connection> {
+pub trait ClientTrait<C: Connection + Sized> {
     type Error: std::fmt::Debug;
 
-    fn connect(address: &str) -> Result<C, Self::Error>;
+    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 tcp;
 
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::thread;
+
 pub use tcp::TCPClient;
 pub use tcp::TCPConnection;
 pub use tcp::TCPServer;

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

@@ -1,6 +1,6 @@
 use std::{
     io::{self, BufRead, BufReader, Write},
-    net::{TcpListener, TcpStream},
+    net::{SocketAddr, TcpListener, TcpStream},
 };
 
 use crate::networkers::{ClientTrait, Connection, ServerTrait};
@@ -8,20 +8,43 @@ use crate::networkers::{ClientTrait, Connection, ServerTrait};
 pub struct TCPConnection {
     stream: TcpStream,
     reader: BufReader<TcpStream>,
+    is_alive: bool,
 }
 
 impl Connection for TCPConnection {
     type Error = io::Error;
 
+    fn get_info(&self) -> String {
+        format!(
+            "tcp://{}",
+            if let Ok(addr) = &self.stream.peer_addr() {
+                addr.to_string()
+            } else {
+                "ERROR".to_string()
+            }
+        )
+    }
+
+    fn is_alive(&self) -> bool {
+        self.is_alive
+    }
+
     fn read(&mut self) -> Result<String, Self::Error> {
         let mut line = String::new();
-        self.reader.read_line(&mut line)?;
+        let n = self.reader.read_line(&mut line)?;
+
+        // Stream sends a null buffer if it is disconnected
+        if n == 0 {
+            self.is_alive = false;
+        }
+
         Ok(line.trim_end().to_string())
     }
 
     fn write(&mut self, data: &str) -> Result<(), Self::Error> {
         writeln!(self.stream, "{}", data)?;
-        self.stream.flush()
+        self.stream.flush()?;
+        Ok(())
     }
 }
 
@@ -32,13 +55,28 @@ pub struct TCPServer {
 impl ServerTrait<TCPConnection> for TCPServer {
     type Error = io::Error;
 
-    fn accept(&mut self) -> Result<TCPConnection, Self::Error> {
+    fn get_info(&self) -> String {
-        let (stream, _) = self.listener.accept()?;
+        format!(
-        let reader = BufReader::new(stream.try_clone()?);
+            "tcp://{}",
-        Ok(TCPConnection { stream, reader })
+            if let Ok(addr) = &self.listener.local_addr() {
+                addr.to_string()
+            } else {
+                "ERROR".to_string()
+            }
+        )
     }
 
-    fn bind(address: &str) -> Result<Self, Self::Error> {
+    fn accept(&self) -> Result<TCPConnection, Self::Error> {
+        let (stream, _) = self.listener.accept()?;
+        let reader = BufReader::new(stream.try_clone()?);
+        Ok(TCPConnection {
+            stream,
+            reader,
+            is_alive: true,
+        })
+    }
+
+    fn bind(address: &SocketAddr) -> Result<Self, Self::Error> {
         let listener = TcpListener::bind(address)?;
         Ok(Self { listener })
     }
@@ -49,9 +87,15 @@ pub struct TCPClient;
 impl ClientTrait<TCPConnection> for TCPClient {
     type Error = io::Error;
 
-    fn connect(address: &str) -> Result<TCPConnection, Self::Error> {
+    fn connect(address: &SocketAddr) -> Result<TCPConnection, Self::Error> {
         let stream = TcpStream::connect(address)?;
         let reader = BufReader::new(stream.try_clone()?);
-        Ok(TCPConnection { stream, reader })
+        let conn = TCPConnection {
+            stream,
+            reader,
+            is_alive: true,
+        };
+        info!("Connected to {}", conn.get_info());
+        Ok(conn)
     }
 }

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

@@ -1,9 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-mod sysinfo;
-
-#[derive(Serialize, Deserialize, Debug)]
-pub enum Packet {
-    Heartbeat,
-    Sysinfo(sysinfo::Sysinfo),
-}

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

@@ -1,6 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct Sysinfo {
-    hostname: String,
-}