Work on server and client connectivity

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

@@ -1,6 +1,9 @@
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Serialize, Deserialize)]
-pub struct CampignConfig {
-    name: String,
+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 +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,
-        port: u16,
+        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/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 bind(address: &str) -> Result<Self, Self::Error>
+    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> {
+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> {
-        let (stream, _) = self.listener.accept()?;
-        let reader = BufReader::new(stream.try_clone()?);
-        Ok(TCPConnection { stream, reader })
+    fn get_info(&self) -> String {
+        format!(
+            "tcp://{}",
+            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,
-}