From fda5e9ea0284989ec9ce1e7220400bf06b5abc89 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Fri, 6 Jun 2025 19:20:49 -0600 Subject: [PATCH] Work on server and client connectivity --- Cargo.toml | 4 +- payload/src/main.rs | 47 ------- src/client/client.rs | 120 ++++++++++++++++++ src/client/gui.rs | 71 +++++++++++ src/client/mod.rs | 6 + src/gui.rs | 62 --------- src/lib.rs | 7 +- src/main.rs | 108 +++++++++++----- src/packets/mod.rs | 37 ++++++ src/server.rs | 0 src/server/config.rs | 29 +++++ src/server/mod.rs | 6 + src/server/server.rs | 153 +++++++++++++++++++++++ ui/pages/listeners.slint | 10 -- ui/structs.slint | 1 - unshell-rs-lib/src/config/campaign.rs | 9 +- unshell-rs-lib/src/config/layers.rs | 2 +- unshell-rs-lib/src/config/listeners.rs | 42 ++++++- unshell-rs-lib/src/connection/mod.rs | 4 + unshell-rs-lib/src/connection/packets.rs | 43 +++++++ unshell-rs-lib/src/lib.rs | 3 +- unshell-rs-lib/src/listeners/client.rs | 9 -- unshell-rs-lib/src/listeners/mod.rs | 6 - unshell-rs-lib/src/listeners/server.rs | 42 ------- unshell-rs-lib/src/networkers/mod.rs | 79 +++++++++++- unshell-rs-lib/src/networkers/tcp.rs | 64 ++++++++-- unshell-rs-lib/src/packets/mod.rs | 9 -- unshell-rs-lib/src/packets/sysinfo.rs | 6 - 28 files changed, 728 insertions(+), 251 deletions(-) create mode 100644 src/client/client.rs create mode 100644 src/client/gui.rs create mode 100644 src/client/mod.rs delete mode 100644 src/gui.rs create mode 100644 src/packets/mod.rs delete mode 100644 src/server.rs create mode 100644 src/server/config.rs create mode 100644 src/server/mod.rs create mode 100644 src/server/server.rs create mode 100644 unshell-rs-lib/src/connection/mod.rs create mode 100644 unshell-rs-lib/src/connection/packets.rs delete mode 100644 unshell-rs-lib/src/listeners/client.rs delete mode 100644 unshell-rs-lib/src/listeners/mod.rs delete mode 100644 unshell-rs-lib/src/listeners/server.rs delete mode 100644 unshell-rs-lib/src/packets/mod.rs delete mode 100644 unshell-rs-lib/src/packets/sysinfo.rs diff --git a/Cargo.toml b/Cargo.toml index 893260e..36b0d26 100644 --- a/Cargo.toml +++ b/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" diff --git a/payload/src/main.rs b/payload/src/main.rs index 81739e8..f62bbe4 100644 --- a/payload/src/main.rs +++ b/payload/src/main.rs @@ -14,53 +14,6 @@ use unshell_rs_lib::{ packets::Packet, }; -// Generic client function -pub fn run_client(address: &str) -> Result<(), Box> -where - Cl: ClientTrait, - 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>> = 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::(data_json.as_str()); - println!("{:?}", packet); - } - Err(e) => { - eprintln!("Error reading, {}", e); - } - } - } -} - fn main() -> Result<(), Box> { run_client::("127.0.0.1:3000")?; diff --git a/src/client/client.rs b/src/client/client.rs new file mode 100644 index 0000000..dec2f73 --- /dev/null +++ b/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>, + + parameters: Arc>, + outgoing_packets: Arc>>, +} + +impl UnshellClient { + pub fn new(addr: SocketAddr) -> Result> { + let client = Arc::new(Mutex::new(TCPClient::connect(&addr)?)); + let outgoing_packets = Arc::new(Mutex::new(Vec::::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(¶meters); + 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 { + self.parameters.lock().unwrap().get(key).cloned() + } +} + +fn print_type_of(_: &T) -> &'static str { + std::any::type_name::() +} diff --git a/src/client/gui.rs b/src/client/gui.rs new file mode 100644 index 0000000..0a67e71 --- /dev/null +++ b/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>>, +} + +slint::include_modules!(); +impl UnshellGui { + pub fn start(client: UnshellClient) -> Result<(), Box> { + 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::>(), + // ))); + } +} + +// trait diff --git a/src/client/mod.rs b/src/client/mod.rs new file mode 100644 index 0000000..e17ee56 --- /dev/null +++ b/src/client/mod.rs @@ -0,0 +1,6 @@ +mod gui; + +mod client; + +pub use client::UnshellClient; +pub use gui::UnshellGui; diff --git a/src/gui.rs b/src/gui.rs deleted file mode 100644 index 01d020d..0000000 --- a/src/gui.rs +++ /dev/null @@ -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> { - 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 = 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::>(), - ))); - - ui.run()?; - - Ok(Self {}) - } -} diff --git a/src/lib.rs b/src/lib.rs index 095dab7..885f676 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/main.rs b/src/main.rs index 19411d2..a415bd8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,19 +1,25 @@ -use std::error::Error; - -use clap::{Parser, Subcommand}; -use log::trace; -use slint::{ModelRc, VecModel}; -use unshell_rs_lib::{ - config::listeners::ListenerConfig, - listeners::Listener, - networkers::{ServerTrait, TCPServer}, +use std::{ + env, + error::Error, + net::{IpAddr, Ipv4Addr, SocketAddr}, + str::FromStr, }; -/// 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 { - /// Only listen for command clients locally - #[arg(short, long, default_value_t = false)] - local: bool, + Server { + /// IPv4 to listen for clients on. + host: String, /// 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 { - /// Listen for command clients remotely aswell - #[arg(short, long, default_value_t = true)] - remote: bool, + /// Run GUI and connect to remote server + Remote { + /// Remote server to connect to + 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> { + 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 { - remote, - service_port, - } => {} - Commands::Serve { - local, - service_port, - // web_port, - } => {} - } + Commands::Local { config_filepath } => { + let mut server = UnshellServer::from_filepath(config_filepath.as_str()); + server.run(LOCAL_SOCKET)?; - // 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(()) } diff --git a/src/packets/mod.rs b/src/packets/mod.rs new file mode 100644 index 0000000..64ba9e7 --- /dev/null +++ b/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; + +#[derive(Debug, Serialize, Deserialize)] +pub enum GuiPacket { + GetParameter(String), + AckGetParameter(String, Option), + 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 { + serde_json::to_string(self) + } + + pub fn decode(string: &str) -> Result { + serde_json::from_str::(string) + } +} diff --git a/src/server.rs b/src/server.rs deleted file mode 100644 index e69de29..0000000 diff --git a/src/server/config.rs b/src/server/config.rs new file mode 100644 index 0000000..5524fbc --- /dev/null +++ b/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 = 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, +} diff --git a/src/server/mod.rs b/src/server/mod.rs new file mode 100644 index 0000000..fc2fc81 --- /dev/null +++ b/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; diff --git a/src/server/server.rs b/src/server/server.rs new file mode 100644 index 0000000..46fb03a --- /dev/null +++ b/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, + + #[serde(skip)] + client_count: usize, + #[serde(skip)] + broadcast_flag: HashMap>, +} + +pub struct UnshellServer { + config: Arc>, +} + +impl UnshellServer { + pub fn from_filepath(filepath: &str) -> Self { + (|| -> Result> { + 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::(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> { + let on_connect = |connection: TCPConnection, + config_clone: Arc>| { + // 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(¶m); + 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(()) + } +} diff --git a/ui/pages/listeners.slint b/ui/pages/listeners.slint index b411010..e985265 100644 --- a/ui/pages/listeners.slint +++ b/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; - } - } } } diff --git a/ui/structs.slint b/ui/structs.slint index 6e03039..58bdbab 100644 --- a/ui/structs.slint +++ b/ui/structs.slint @@ -2,5 +2,4 @@ export struct UITcpListener { enabled: bool, name: string, remote_host: string, - port: int, } diff --git a/unshell-rs-lib/src/config/campaign.rs b/unshell-rs-lib/src/config/campaign.rs index e9141be..1b27d66 100644 --- a/unshell-rs-lib/src/config/campaign.rs +++ b/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, } diff --git a/unshell-rs-lib/src/config/layers.rs b/unshell-rs-lib/src/config/layers.rs index 32290af..a72cb66 100644 --- a/unshell-rs-lib/src/config/layers.rs +++ b/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 {} diff --git a/unshell-rs-lib/src/config/listeners.rs b/unshell-rs-lib/src/config/listeners.rs index 6211a6d..9d2e33c 100644 --- a/unshell-rs-lib/src/config/listeners.rs +++ b/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, + + #[serde(skip)] + connections: Option>>>, }, } + +impl ListenerConfig { + pub fn start(self) -> Result<(), Box> { + match self { + ListenerConfig::Tcp { + mut enabled, + addr, + layers, + mut connections, + .. + } => { + let server = TCPServer::bind(&addr)?; + + enabled = true; + + // connections = Some(run_listener(server)); + } + } + + Ok(()) + } +} diff --git a/unshell-rs-lib/src/connection/mod.rs b/unshell-rs-lib/src/connection/mod.rs new file mode 100644 index 0000000..ded8de8 --- /dev/null +++ b/unshell-rs-lib/src/connection/mod.rs @@ -0,0 +1,4 @@ +mod packets; + +pub use packets::C2Packet; +pub use packets::ErrorPacket; diff --git a/unshell-rs-lib/src/connection/packets.rs b/unshell-rs-lib/src/connection/packets.rs new file mode 100644 index 0000000..6fb6c51 --- /dev/null +++ b/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 { + serde_json::to_string(self) + } + + pub fn decode(string: &str) -> Result { + serde_json::from_str::(string) + } +} diff --git a/unshell-rs-lib/src/lib.rs b/unshell-rs-lib/src/lib.rs index 14ecc9e..69eaa11 100644 --- a/unshell-rs-lib/src/lib.rs +++ b/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; diff --git a/unshell-rs-lib/src/listeners/client.rs b/unshell-rs-lib/src/listeners/client.rs deleted file mode 100644 index 2397f36..0000000 --- a/unshell-rs-lib/src/listeners/client.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub struct Client { - pub stream: C, -} - -impl Client { - pub fn new(stream: C) -> Self { - Self { stream } - } -} diff --git a/unshell-rs-lib/src/listeners/mod.rs b/unshell-rs-lib/src/listeners/mod.rs deleted file mode 100644 index 12b6dd8..0000000 --- a/unshell-rs-lib/src/listeners/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -use crate::layers::Layer; - -mod client; -mod server; - -pub use server::Listener; diff --git a/unshell-rs-lib/src/listeners/server.rs b/unshell-rs-lib/src/listeners/server.rs deleted file mode 100644 index 29fad41..0000000 --- a/unshell-rs-lib/src/listeners/server.rs +++ /dev/null @@ -1,42 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use crate::{ - listeners::client::Client, - networkers::{Connection, ServerTrait}, -}; - -pub struct Listener { - pub server: Arc>, - pub clients: Arc>>>, -} - -impl Listener { - 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> - where - S: ServerTrait, - 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); - } - } - } - } -} diff --git a/unshell-rs-lib/src/networkers/mod.rs b/unshell-rs-lib/src/networkers/mod.rs index 8861fdb..7eda544 100644 --- a/unshell-rs-lib/src/networkers/mod.rs +++ b/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; fn write(&mut self, data: &str) -> Result<(), Self::Error>; } @@ -10,20 +13,88 @@ pub trait Connection: Send + Sync { pub trait ServerTrait { type Error: std::fmt::Debug; - fn accept(&mut self) -> Result; - fn bind(address: &str) -> Result + fn get_info(&self) -> String; + fn accept(&self) -> Result; + fn bind(address: &SocketAddr) -> Result where Self: Sized; } -pub trait ClientTrait { +pub trait ClientTrait { type Error: std::fmt::Debug; - fn connect(address: &str) -> Result; + fn connect(address: &SocketAddr) -> Result; +} + +pub fn run_listener_state(server: S, on_connect_callback: R, state: Arc) +/*-> Arc>>*/ +where + S: ServerTrait + Sync + Send + 'static, + C: Connection + 'static, + R: Fn(C, Arc) + Sync + Send + 'static, + A: Sync + Send + 'static, +{ + info!("Started listener {}", server.get_info()); + // let clients: Arc>> = 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(server: S, on_connect_callback: R) +/*-> Arc>>*/ +where + S: ServerTrait + Sync + Send + 'static, + C: Connection + 'static, + R: Fn(C) + Sync + Send + 'static, +{ + info!("Started listener {}", server.get_info()); + // let clients: Arc>> = 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; diff --git a/unshell-rs-lib/src/networkers/tcp.rs b/unshell-rs-lib/src/networkers/tcp.rs index 794133b..72cf393 100644 --- a/unshell-rs-lib/src/networkers/tcp.rs +++ b/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, + 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 { 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 for TCPServer { type Error = io::Error; - fn accept(&mut self) -> Result { - 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 { + fn accept(&self) -> Result { + let (stream, _) = self.listener.accept()?; + let reader = BufReader::new(stream.try_clone()?); + Ok(TCPConnection { + stream, + reader, + is_alive: true, + }) + } + + fn bind(address: &SocketAddr) -> Result { let listener = TcpListener::bind(address)?; Ok(Self { listener }) } @@ -49,9 +87,15 @@ pub struct TCPClient; impl ClientTrait for TCPClient { type Error = io::Error; - fn connect(address: &str) -> Result { + fn connect(address: &SocketAddr) -> Result { 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) } } diff --git a/unshell-rs-lib/src/packets/mod.rs b/unshell-rs-lib/src/packets/mod.rs deleted file mode 100644 index c727c6d..0000000 --- a/unshell-rs-lib/src/packets/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -use serde::{Deserialize, Serialize}; - -mod sysinfo; - -#[derive(Serialize, Deserialize, Debug)] -pub enum Packet { - Heartbeat, - Sysinfo(sysinfo::Sysinfo), -} diff --git a/unshell-rs-lib/src/packets/sysinfo.rs b/unshell-rs-lib/src/packets/sysinfo.rs deleted file mode 100644 index 5b3e892..0000000 --- a/unshell-rs-lib/src/packets/sysinfo.rs +++ /dev/null @@ -1,6 +0,0 @@ -use serde::{Deserialize, Serialize}; - -#[derive(Serialize, Deserialize, Debug)] -pub struct Sysinfo { - hostname: String, -}