astatin3/unshell-nodes-rs
1
0
Fork 0
mirror of https://github.com/Astatin3/unshell-nodes-rs.git synced 2026-06-08 16:18:08 -06:00
Code Issues Packages Projects Releases Wiki Activity

Work on server and client connectivity

Browse Source
This commit is contained in:
Michael Mikovsky
2025-06-06 19:20:49 -06:00
parent 92c9f08a5c
commit fda5e9ea02
28 changed files with 728 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+3 -1
View File
@@ -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
-47
View File
@@ -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
+120
View File
@@ -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
+71
View File
@@ -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
+6
View File
@@ -0,0 +1,6 @@
mod gui;
mod client;
pub use client::UnshellClient;
pub use gui::UnshellGui;
src/gui.rs
-62
View File
@@ -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
+6 -1
View File
@@ -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
+75 -33
View File
@@ -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<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 {
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 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 {
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(());
}
// let mut server = Listener::new(TCPServer::bind("0.0.0.0:3000")?);
// server.run_listener()?;
loop {}
}
};
Ok(())
}
src/packets/mod.rs
+37
View File
@@ -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.rs
View File
src/server/config.rs
+29
View File
@@ -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
+6
View File
@@ -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
+153
View File
@@ -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
-10
View File
@@ -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
-1
View File
@@ -2,5 +2,4 @@ export struct UITcpListener {
enabled: bool,
name: string,
remote_host: string,
port: int,
}
unshell-rs-lib/src/config/campaign.rs
+6 -3
View File
@@ -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 -1
View File
@@ -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
+38 -4
View File
@@ -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
+4
View File
@@ -0,0 +1,4 @@
mod packets;
pub use packets::C2Packet;
pub use packets::ErrorPacket;
unshell-rs-lib/src/connection/packets.rs
+43
View File
@@ -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
+1 -2
View File
@@ -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
-9
View File
@@ -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
-6
View File
@@ -1,6 +0,0 @@
use crate::layers::Layer;
mod client;
mod server;
pub use server::Listener;
unshell-rs-lib/src/listeners/server.rs
-42
View File
@@ -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
+75 -4
View File
@@ -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
+54 -10
View File
@@ -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
-9
View File
@@ -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
-6
View File
@@ -1,6 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug)]
pub struct Sysinfo {
hostname: String,
}