diff --git a/Cargo.toml b/Cargo.toml index 323e3dd..85731fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,9 +9,11 @@ clap = { version = "4.5.39", features = ["derive"] } crossbeam-channel = "0.5.15" lazy_static = "1.5.0" log = "0.4.27" +portable-pty = "0.9.0" pretty_env_logger = "0.5.0" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.140" +term_size = "0.3.2" # slint = "1.11.0" unshell-rs-lib = { path = "./unshell-rs-lib" } uuid = { version = "1.17.0", features = ["v4"] } diff --git a/src/client/cli.rs b/src/client/cli.rs index 3293dfd..2e0679e 100644 --- a/src/client/cli.rs +++ b/src/client/cli.rs @@ -1,101 +1,133 @@ -use std::{io::Write, net::SocketAddr}; - -use unshell_rs_lib::{ - Error, - nodes::{ConnectionConfig, Node}, +use std::{ + io::{Stdin, Stdout, Write}, + net::SocketAddr, }; -use crate::C2Packet; +use clap::{Parser, Subcommand, command}; +use unshell_rs_lib::{ + Error, + nodes::{ConnectionConfig, NodeContainer}, +}; -pub struct Cli; +use crate::client::node_cli::NodeCli; -impl Cli { - pub fn connect(socket: SocketAddr) -> Result<(), Error> { - // let mut client = build_client(TCPClient::connect(&addr)?, vec![])?; +pub trait Cli { + fn name(&self) -> String; + fn parse(&mut self, input: Vec) -> Result<(), Error>; +} - let stdin = std::io::stdin(); - let mut stdout = std::io::stdout(); +#[derive(Debug, Parser)] +pub struct CommandHolder

+where + P: Subcommand, +{ + #[command(subcommand)] + pub command: P, +} - let node = Node::::run_node( - "Client".to_string(), - vec![ConnectionConfig { - socket, - layers: vec![], - }], - vec![], - )?; +pub fn connect_cli(socket: SocketAddr) -> Result<(), Error> { + // let mut client = build_client(TCPClient::connect(&addr)?, vec![])?; - // let mut client_clone = client.try_clone()?; - // thread::spawn(move || { - // // let data = client.read()?; + let node = NodeContainer::connect( + "Client".to_string(), + vec![ConnectionConfig { + socket, + layers: vec![], + }], + vec![], + )?; - // let packet = Packets::decode(client_clone.read().unwrap().as_str()).unwrap(); + let mut current_parser = Box::new(NodeCli::new(node)) as Box; - // match packet { - // Packets::UpdateConnections(items) => { - // for item in items { - // println!("{}", item); - // } - // } - // Packets::UpdateRoutes(items) => { - // for item in items { - // println!("{}", item); - // } - // } - // _ => { - // client_clone - // .write( - // Packets::Error(PacketError::UnsupportedType) - // .encode() - // .unwrap() - // .as_str(), - // ) - // .unwrap(); - // warn!("Invalid packet: {:?}", packet) - // } - // } - // }); + let parse = |current_parser: &mut Box, + stdin: &Stdin, + stdout: &mut Stdout| + -> Result<(), Error> { + let name = current_parser.name(); + print!("Unshell | {}> ", name); + stdout.flush()?; - let selected_node: Option = None; + let mut input = String::new(); + stdin.read_line(&mut input)?; - loop { - print!("> "); - stdout.flush()?; + let input = input.trim(); + if input.is_empty() { + return Ok(()); + } - let mut input = String::new(); - stdin.read_line(&mut input)?; - let input = input.trim(); + let mut input = split_escape(input); + // Clap expects the first arg to be the program name + input.insert(0, name); - let mut node_state = node.state.lock().unwrap(); + current_parser.parse(input)?; - let mut split = input.split(" "); + Ok(()) + }; - match split.next().unwrap() { - "nodes" => { - for (i, node) in node_state.get_all_nodes().iter().enumerate() { - println!("{} -> {}", i, node); - } - } - "ping" => { - // if split.count().clone() <= 1 { - // warn!("You must specify an option"); - // continue; - // } + let stdin = std::io::stdin(); + let mut stdout = std::io::stdout(); - if let Ok(i) = str::parse::(split.next().unwrap()) { - let nodes = node_state.get_all_nodes(); - let node = nodes.get(i).unwrap().clone(); - node_state.send_unrouted(node, &C2Packet::Aa).unwrap(); - } else { - println!(""); - } - } - _ => { - warn!("Invalid command!") - } - } - - // client.write(input)?; + loop { + if let Err(e) = parse(&mut current_parser, &stdin, &mut stdout) { + error!("Failed to parse: {}", e); } } } + +fn split_escape(input: &str) -> Vec { + let mut result = Vec::new(); + let mut current = String::new(); + let mut chars = input.chars().peekable(); + let mut in_single_quote = false; + let mut in_double_quote = false; + + while let Some(ch) = chars.next() { + match ch { + '\\' => { + // Handle escape sequences + if let Some(&next_ch) = chars.peek() { + match next_ch { + '\'' | '"' | '\\' | ' ' => { + // Escape recognized characters + current.push(chars.next().unwrap()); + } + _ => { + // For other characters, keep the backslash + current.push(ch); + } + } + } else { + // Backslash at end of string + current.push(ch); + } + } + '\'' if !in_double_quote => { + in_single_quote = !in_single_quote; + } + '"' if !in_single_quote => { + in_double_quote = !in_double_quote; + } + ' ' if !in_single_quote && !in_double_quote => { + // Split on unquoted spaces + if !current.is_empty() { + result.push(current.clone()); + current.clear(); + } + // Skip consecutive spaces + while chars.peek() == Some(&' ') { + chars.next(); + } + } + _ => { + current.push(ch); + } + } + } + + // Add the last token if it exists + if !current.is_empty() { + result.push(current); + } + + result +} diff --git a/src/client/client_node.rs b/src/client/client_node.rs deleted file mode 100644 index e69de29..0000000 diff --git a/src/client/mod.rs b/src/client/mod.rs index 24ae707..dd81511 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,4 +1,4 @@ mod cli; -mod client_node; +mod node_cli; -pub use cli::Cli; +pub use cli::connect_cli; diff --git a/src/client/node_cli.rs b/src/client/node_cli.rs new file mode 100644 index 0000000..be7feed --- /dev/null +++ b/src/client/node_cli.rs @@ -0,0 +1,112 @@ +use std::time::Instant; + +use clap::{Parser, Subcommand}; +use portable_pty::{PtySize, native_pty_system}; +use unshell_rs_lib::{C2Packet, Error, nodes::NodeContainer}; + +use crate::client::cli::{Cli, CommandHolder}; + +pub struct NodeCli { + node: NodeContainer, + subcommand: Option>, +} + +#[derive(Debug, Subcommand)] +pub enum NodeCliCommands { + /// List out connected nodes + Nodes, + /// Send a ping to a remote node + Ping { n: usize }, + /// Attempt to create a shell at a remote node + Sh { n: usize }, +} + +impl Cli for NodeCli { + fn name(&self) -> String { + "Local".to_string() + } + fn parse(&mut self, input: Vec) -> Result<(), Error> { + if let Some(subcommand) = &mut self.subcommand { + return subcommand.parse(input); + } + let parsed_command = CommandHolder::::try_parse_from(input)?; + + let node_ids = self.node.get_nodes(); + + match parsed_command.command { + NodeCliCommands::Nodes => { + info!("N | Name"); + for (i, node) in node_ids.iter().enumerate() { + info!("[{}] {}", i + 1, node); + } + } + NodeCliCommands::Ping { n } => { + // if split.count().clone() <= 1 { + // warn!("You must specify an option"); + // continue; + // } + + if n <= 0 { + warn!("Node id must be greater than zero"); + } else if n > node_ids.len() { + warn!("Node id {} is out of maximum range {}", n, node_ids.len()); + } else { + let start = Instant::now(); + let node = node_ids.get(n - 1).unwrap().clone(); + self.node.send_unrouted(&node, &C2Packet::Ping).unwrap(); + info!("Sent ping..."); + + let (_, packet) = self.node.read_packet()?; + match packet { + C2Packet::Pong => { + // if src != nod + info!( + "Pong! Latency: {}ms", + (start.elapsed().as_micros() as f32) / 1000. + ); + } + _ => { + error!("Got incorrect packet: {:?}", packet); + } + } + + // node_state = self.node.state.lock().unwrap(); + } + } + NodeCliCommands::Sh { n } => { + if n <= 0 { + warn!("Node id must be greater than zero"); + } else if n > node_ids.len() { + warn!("Node id {} is out of maximum range {}", n, node_ids.len()); + } else { + let node_id = node_ids.get(n - 1).unwrap().clone(); + } + } + } + + Ok(()) + } +} + +impl NodeCli { + pub fn new(node: NodeContainer) -> Self { + Self { + node, + subcommand: None, + } + } + + pub fn run_pty(&mut self) -> Result<(), Error> { + let pty_system = native_pty_system(); + let pty_pair = pty_system.openpty(PtySize { + rows: 24, + cols: 80, + pixel_width: 0, + pixel_height: 0, + })?; + + Ok(()) + + // pty_pair.Ok(()) + } +} diff --git a/src/endpoint/endpoint.rs b/src/endpoint/endpoint.rs index 94a6ac6..8ef1e50 100644 --- a/src/endpoint/endpoint.rs +++ b/src/endpoint/endpoint.rs @@ -1,14 +1,12 @@ use std::net::SocketAddr; use unshell_rs_lib::{ - Error, - nodes::{ConnectionConfig, Node}, + C2Packet, Error, + nodes::{ConnectionConfig, NodeContainer}, }; -use crate::C2Packet; - pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> { - let node = Node::::run_node( + let node = NodeContainer::connect( "Server".to_string(), vec![], vec![ConnectionConfig { @@ -18,16 +16,17 @@ pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> { )?; loop { - match node.rx.recv()? { - C2Packet::Aa => { - info!("1"); + let (src, packet) = node.read_packet()?; + match packet { + C2Packet::Ping => { + info!("Ping from {}!", src); + node.send_unrouted(&src, &C2Packet::Pong)?; + // (&mut node.state.lock().unwrap()).send_unrouted(src, &C2Packet::Pong)?; } - C2Packet::Bb => { - info!("2"); - } - C2Packet::Cc => { - info!("3"); + C2Packet::Pong => { + info!("Pong!"); } + _ => {} } } } diff --git a/src/lib.rs b/src/lib.rs index 2f080b2..0138e06 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,12 +3,10 @@ extern crate log; mod client; mod endpoint; -mod packets; -pub use client::Cli; +pub use client::connect_cli; pub use endpoint::run_endpoint; -pub use packets::C2Packet; // pub use client::UnshellClient; // pub use client::UnshellGui; diff --git a/src/main.rs b/src/main.rs index fd24f83..6d771fa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,11 +7,11 @@ use std::{ use clap::{Parser, Subcommand}; use log::error; -use unshell_rs::{Cli, run_endpoint}; -use unshell_rs_lib::nodes::ConnectionConfig; +use unshell_rs::{connect_cli, run_endpoint}; pub static DEFAULT_CONFIG_FILEPATH: &'static str = "server_config.json"; +pub static DEFAULT_RELAY_HOST: &'static str = "0.0.0.0"; // The default port that this program looks for pub static DEFAULT_SERVICE_PORT: u16 = 13370; // The default website port that this program looks for @@ -33,7 +33,7 @@ enum Commands { // Run as a service, and potentially hosting a website Relay { /// IPv4 to listen for clients on. - #[arg(short, long, default_value_t = ("0.0.0.0".to_string()))] + #[arg(short, long, default_value_t = DEFAULT_RELAY_HOST.to_string())] host: String, /// Port listen to for command clients @@ -143,7 +143,7 @@ fn main() -> Result<(), Box> { // ), Commands::Connect { host, port } => { let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str()); - Cli::connect(if let Ok(addr) = addr { + connect_cli(if let Ok(addr) = addr { addr } else { error!("Could not parse address!"); @@ -153,7 +153,7 @@ fn main() -> Result<(), Box> { Commands::Relay { host, port, - config_filepath, + config_filepath: _, } => { let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str()); run_endpoint(if let Ok(addr) = addr { diff --git a/src/packets.rs b/src/packets.rs deleted file mode 100644 index 05d8635..0000000 --- a/src/packets.rs +++ /dev/null @@ -1,8 +0,0 @@ -use bincode::{Decode, Encode}; - -#[derive(Debug, Encode, Decode, Clone)] -pub enum C2Packet { - Aa, - Bb, - Cc, -} diff --git a/unshell-rs-lib/src/lib.rs b/unshell-rs-lib/src/lib.rs index 89b232e..35e058c 100644 --- a/unshell-rs-lib/src/lib.rs +++ b/unshell-rs-lib/src/lib.rs @@ -8,3 +8,6 @@ static BINCODE_CONFIG: bincode::config::Configuration = bincode::config::standar pub mod layers; pub mod networkers; pub mod nodes; +mod packets; + +pub use packets::C2Packet; diff --git a/unshell-rs-lib/src/nodes/mod.rs b/unshell-rs-lib/src/nodes/mod.rs index 0358d70..baa93c4 100644 --- a/unshell-rs-lib/src/nodes/mod.rs +++ b/unshell-rs-lib/src/nodes/mod.rs @@ -1,8 +1,10 @@ mod listener; mod node; +mod node_container; mod packets; +mod stream; pub use listener::ConnectionConfig; pub use node::Node; -// pub use packets::PacketError; -pub use packets::Packets; +pub use node_container::NodeContainer; +pub use stream::Stream; diff --git a/unshell-rs-lib/src/nodes/node.rs b/unshell-rs-lib/src/nodes/node.rs index 859878f..f795c3c 100644 --- a/unshell-rs-lib/src/nodes/node.rs +++ b/unshell-rs-lib/src/nodes/node.rs @@ -1,13 +1,17 @@ use std::{ collections::HashMap, fmt::Debug, - sync::{Arc, Mutex}, + sync::{ + Arc, Mutex, + mpsc::{self, Receiver, Sender}, + }, thread, time::Duration, }; use bincode::{Decode, Encode}; -use crossbeam_channel::{Receiver, Sender}; +// use std:::{Receiver, Sender}; +#[allow(deprecated)] use rand::{seq::IndexedRandom, thread_rng}; use crate::{ @@ -20,16 +24,6 @@ use crate::{ }, }; -pub struct NodeState

-where - P: Encode + Decode<()> + Debug + Clone + 'static, -{ - id: String, - connections: HashMap>, - map: HashMap>, - packet_listener: Sender

, -} - fn read(c: &mut Box) -> Result { Packets::decode(c.read()?.as_slice()) } @@ -43,7 +37,7 @@ where P: Encode + Decode<()> + Debug + Clone + 'static, { pub state: Arc>>, - pub rx: Receiver

, + pub rx: Receiver<(String, P)>, } impl

Node

@@ -60,7 +54,7 @@ where { // let mut parent = build_client(TCPClient::connect(&parent.socket)?, parent.layers)?; - let (tx, rx) = crossbeam_channel::unbounded(); + let (tx, rx) = mpsc::channel(); let state = Arc::new(Mutex::new(NodeState::

{ id: id, //Uuid::new_v4().to_string(), //TODO: Calling an OS RNG can pose a problem for security; @@ -128,12 +122,22 @@ where write(&mut connection, Packets::SyncUUID(this_uuid.clone()))?; // Recieve UUID - let other_uuid = if let Packets::SyncUUID(source) = read(&mut connection)? { + let uuid_result = read(&mut connection)?; + let other_uuid = if let Packets::SyncUUID(source) = uuid_result { source } else { - return Err("Could not get UUID!".into()); + return Err(format!("Could not get UUID! Got {:?}", uuid_result).into()); }; + if (&mut state.lock().unwrap()).knows_client(&other_uuid) { + write(&mut connection, Packets::ErrorNameExists)?; + return Err(format!( + "Attempted to accept connection from node {} which already exists!", + other_uuid + ) + .into()); + } + info!("New Node! {} (direct)", other_uuid); // Add connection @@ -189,6 +193,16 @@ where } } +pub struct NodeState

+where + P: Encode + Decode<()> + Debug + Clone + 'static, +{ + id: String, + connections: HashMap>, + map: HashMap>, + packet_listener: Sender<(String, P)>, +} + impl

NodeState

where P: Encode + Decode<()> + Debug + Clone + Send + 'static, @@ -207,7 +221,7 @@ where } fn knows_client(&self, id: &String) -> bool { - self.get_known_nodes().contains(id) + self.get_all_nodes().contains(id) } // Remove all nodes where the routes are empty @@ -381,7 +395,7 @@ where fn route_packet(&mut self, src: String, dest: String, data: Vec) -> Result<(), Error> { if dest == self.id { - self.packet_listener.send(decode_vec::

(&data)?)?; + self.packet_listener.send((src, decode_vec::

(&data)?))?; } else { if self.connections.contains_key(&dest) { write( @@ -389,6 +403,7 @@ where Packets::DataUnrouted { src, dest, data }, )?; } else if self.map.contains_key(&dest) { + #[allow(deprecated)] let next_uuid = self .map .get(&dest) diff --git a/unshell-rs-lib/src/nodes/node_container.rs b/unshell-rs-lib/src/nodes/node_container.rs new file mode 100644 index 0000000..bd02003 --- /dev/null +++ b/unshell-rs-lib/src/nodes/node_container.rs @@ -0,0 +1,151 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + thread, +}; + +use crossbeam_channel::{Receiver, Sender}; + +use crate::{ + C2Packet, Error, + nodes::{ + ConnectionConfig, Node, Stream, + node::NodeState, + packets::{decode_vec, encode_vec}, + }, + packets::TransportLayerPacket, +}; + +type Streams = Arc>)>>>>; + +pub struct NodeContainer { + streams: Streams, + state: Arc>>, + spontanious_rx: Receiver<(String, C2Packet)>, +} + +impl NodeContainer { + pub fn connect( + id: String, + clients: Vec, + listeners: Vec, + ) -> Result { + let node = Node::run_node(id, clients, listeners)?; + let streams = Arc::new(Mutex::new(HashMap::new())); + let (spontanious_tx, spontanious_rx) = crossbeam_channel::unbounded(); + + let s = Self { + streams: Arc::clone(&streams), + state: Arc::clone(&node.state), + spontanious_rx, + }; + + // Start node listening thread + thread::spawn(move || { + loop { + if let Err(e) = Self::node_listening_thread(&node, &streams, &spontanious_tx) { + error!("Got error: {}", e); + } + } + }); + + Ok(s) + } + + fn node_listening_thread( + node: &Node, + streams: &Streams, + spontanious_tx: &Sender<(String, C2Packet)>, + ) -> Result<(), Error> { + let (src, packet) = node.rx.recv()?; + + match packet { + TransportLayerPacket::RequestStreamUnrouted { stream_id } => { + let local_stream_id = streams.lock().unwrap().keys().len(); + streams + .lock() + .unwrap() + .insert((local_stream_id, src.clone()), None); + (&mut node.state.lock().unwrap()).send_unrouted( + src, + &TransportLayerPacket::AckStreamUnrouted { + local_stream_id, + remote_stream_id: stream_id, + }, + )?; + + Ok(()) + } + TransportLayerPacket::AckStreamUnrouted { + local_stream_id, + remote_stream_id, + } => { + let key = &(remote_stream_id, src); + if let Some(stream_mut) = streams.lock().unwrap().get_mut(&key) { + if stream_mut.is_none() { + let stream = Self::create_stream(local_stream_id, node, src, stream_mut)?; + Ok(()) + } else { + Err(format!("Stream {:?} already exists!", key).into()) + } + } else { + Err(format!("Could not find stream id by {:?}", key).into()) + } + } + TransportLayerPacket::StreamDataUnrouted { stream_id, data } => todo!(), + TransportLayerPacket::SpontaniousDataUnrouted { data } => { + spontanious_tx.send((src, decode_vec::(&data)?))?; + Ok(()) + } + } + } + + fn create_stream( + remote_stream_id: usize, + dest: String, + node: &Node, + stream_mut: &mut Option<(Stream, Sender>)>, + ) -> Result<(), Error> { + let (recv_tx, recv_rx) = crossbeam_channel::unbounded(); + let (send_tx, send_rx) = crossbeam_channel::unbounded(); + + let stream = Stream::new(send_tx, recv_rx); + + let _ = stream_mut.insert((stream, recv_tx)); + + thread::spawn(move || { + loop { + let packet = send_rx.recv().unwrap(); + (&mut node.state.lock().unwrap()) + .send_unrouted( + dest, + &TransportLayerPacket::StreamDataUnrouted { + stream_id: remote_stream_id, + data: packet, + }, + ) + .unwrap(); + } + }); + + Ok(()) + } + + pub fn get_nodes(&self) -> Vec { + self.state.lock().unwrap().get_all_nodes() + } + + pub fn send_unrouted(&self, dest: &String, data: &C2Packet) -> Result<(), Error> { + (&mut self.state.lock().unwrap()).send_unrouted( + dest.clone(), + &TransportLayerPacket::SpontaniousDataUnrouted { + data: encode_vec(data)?, + }, + )?; + Ok(()) + } + + pub fn read_packet(&self) -> Result<(String, C2Packet), Error> { + Ok(self.spontanious_rx.recv()?) + } +} diff --git a/unshell-rs-lib/src/nodes/packets.rs b/unshell-rs-lib/src/nodes/packets.rs index 9bf7e62..2b0449a 100644 --- a/unshell-rs-lib/src/nodes/packets.rs +++ b/unshell-rs-lib/src/nodes/packets.rs @@ -13,15 +13,24 @@ pub enum Packets { Disconnect { routes: Vec, }, + + // Send single data packet without routing details DataUnrouted { src: String, dest: String, data: Vec, }, + // Send single data packet with routing details DataRouted { path: Vec, data: Vec, }, + + // DataStreamRouted { + // path: Vec, + // data: Vec, + // }, + ErrorNameExists, } impl Packets { diff --git a/unshell-rs-lib/src/nodes/stream.rs b/unshell-rs-lib/src/nodes/stream.rs new file mode 100644 index 0000000..8664b1b --- /dev/null +++ b/unshell-rs-lib/src/nodes/stream.rs @@ -0,0 +1,36 @@ +use crossbeam_channel::{Receiver, Sender}; + +use crate::networkers::Connection; + +pub struct Stream { + tx: Sender>, + rx: Receiver>, +} + +impl Connection for Stream { + fn get_info(&self) -> String { + "unrouted".to_string() + } + + fn is_alive(&self) -> bool { + true + } + + fn read(&mut self) -> Result, crate::Error> { + Ok(self.rx.recv()?) + } + + fn write(&mut self, data: &[u8]) -> Result<(), crate::Error> { + Ok(self.tx.send(data.to_vec())?) + } + + fn try_clone(&self) -> Result, crate::Error> { + todo!() + } +} + +impl Stream { + pub fn new(tx: Sender>, rx: Receiver>) -> Self { + Self { tx, rx } + } +} diff --git a/unshell-rs-lib/src/packets.rs b/unshell-rs-lib/src/packets.rs new file mode 100644 index 0000000..0098587 --- /dev/null +++ b/unshell-rs-lib/src/packets.rs @@ -0,0 +1,30 @@ +use bincode::{Decode, Encode}; +use std::fmt::Debug; + +#[derive(Debug, Encode, Decode, Clone)] +pub enum TransportLayerPacket { + RequestStreamUnrouted { + stream_id: usize, + }, + AckStreamUnrouted { + local_stream_id: usize, + remote_stream_id: usize, + }, + StreamDataUnrouted { + stream_id: usize, + data: Vec, + }, + + SpontaniousDataUnrouted { + data: Vec, + }, +} + +#[derive(Debug, Encode, Decode, Clone)] +pub enum C2Packet { + Ping, + Pong, + + CreatePTY { width: usize, height: usize }, + PTYData, +}