mirror of
https://github.com/Astatin3/unshell-nodes-rs.git
synced 2026-06-08 16:18:08 -06:00
Work on adding cli, and transport layer
This commit is contained in:
@@ -9,9 +9,11 @@ clap = { version = "4.5.39", features = ["derive"] }
|
|||||||
crossbeam-channel = "0.5.15"
|
crossbeam-channel = "0.5.15"
|
||||||
lazy_static = "1.5.0"
|
lazy_static = "1.5.0"
|
||||||
log = "0.4.27"
|
log = "0.4.27"
|
||||||
|
portable-pty = "0.9.0"
|
||||||
pretty_env_logger = "0.5.0"
|
pretty_env_logger = "0.5.0"
|
||||||
serde = { version = "1.0.219", features = ["derive"] }
|
serde = { version = "1.0.219", features = ["derive"] }
|
||||||
serde_json = "1.0.140"
|
serde_json = "1.0.140"
|
||||||
|
term_size = "0.3.2"
|
||||||
# slint = "1.11.0"
|
# slint = "1.11.0"
|
||||||
unshell-rs-lib = { path = "./unshell-rs-lib" }
|
unshell-rs-lib = { path = "./unshell-rs-lib" }
|
||||||
uuid = { version = "1.17.0", features = ["v4"] }
|
uuid = { version = "1.17.0", features = ["v4"] }
|
||||||
|
|||||||
+114
-82
@@ -1,101 +1,133 @@
|
|||||||
use std::{io::Write, net::SocketAddr};
|
use std::{
|
||||||
|
io::{Stdin, Stdout, Write},
|
||||||
use unshell_rs_lib::{
|
net::SocketAddr,
|
||||||
Error,
|
|
||||||
nodes::{ConnectionConfig, Node},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
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 trait Cli {
|
||||||
pub fn connect(socket: SocketAddr) -> Result<(), Error> {
|
fn name(&self) -> String;
|
||||||
// let mut client = build_client(TCPClient::connect(&addr)?, vec![])?;
|
fn parse(&mut self, input: Vec<String>) -> Result<(), Error>;
|
||||||
|
}
|
||||||
|
|
||||||
let stdin = std::io::stdin();
|
#[derive(Debug, Parser)]
|
||||||
let mut stdout = std::io::stdout();
|
pub struct CommandHolder<P>
|
||||||
|
where
|
||||||
|
P: Subcommand,
|
||||||
|
{
|
||||||
|
#[command(subcommand)]
|
||||||
|
pub command: P,
|
||||||
|
}
|
||||||
|
|
||||||
let node = Node::<C2Packet>::run_node(
|
pub fn connect_cli(socket: SocketAddr) -> Result<(), Error> {
|
||||||
"Client".to_string(),
|
// let mut client = build_client(TCPClient::connect(&addr)?, vec![])?;
|
||||||
vec![ConnectionConfig {
|
|
||||||
socket,
|
|
||||||
layers: vec![],
|
|
||||||
}],
|
|
||||||
vec![],
|
|
||||||
)?;
|
|
||||||
|
|
||||||
// let mut client_clone = client.try_clone()?;
|
let node = NodeContainer::connect(
|
||||||
// thread::spawn(move || {
|
"Client".to_string(),
|
||||||
// // let data = client.read()?;
|
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<dyn Cli>;
|
||||||
|
|
||||||
// match packet {
|
let parse = |current_parser: &mut Box<dyn Cli + 'static>,
|
||||||
// Packets::UpdateConnections(items) => {
|
stdin: &Stdin,
|
||||||
// for item in items {
|
stdout: &mut Stdout|
|
||||||
// println!("{}", item);
|
-> Result<(), Error> {
|
||||||
// }
|
let name = current_parser.name();
|
||||||
// }
|
print!("Unshell | {}> ", name);
|
||||||
// Packets::UpdateRoutes(items) => {
|
stdout.flush()?;
|
||||||
// for item in items {
|
|
||||||
// println!("{}", item);
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// _ => {
|
|
||||||
// client_clone
|
|
||||||
// .write(
|
|
||||||
// Packets::Error(PacketError::UnsupportedType)
|
|
||||||
// .encode()
|
|
||||||
// .unwrap()
|
|
||||||
// .as_str(),
|
|
||||||
// )
|
|
||||||
// .unwrap();
|
|
||||||
// warn!("Invalid packet: {:?}", packet)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// });
|
|
||||||
|
|
||||||
let selected_node: Option<usize> = None;
|
let mut input = String::new();
|
||||||
|
stdin.read_line(&mut input)?;
|
||||||
|
|
||||||
loop {
|
let input = input.trim();
|
||||||
print!("> ");
|
if input.is_empty() {
|
||||||
stdout.flush()?;
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
let mut input = String::new();
|
let mut input = split_escape(input);
|
||||||
stdin.read_line(&mut input)?;
|
// Clap expects the first arg to be the program name
|
||||||
let input = input.trim();
|
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() {
|
let stdin = std::io::stdin();
|
||||||
"nodes" => {
|
let mut stdout = std::io::stdout();
|
||||||
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;
|
|
||||||
// }
|
|
||||||
|
|
||||||
if let Ok(i) = str::parse::<usize>(split.next().unwrap()) {
|
loop {
|
||||||
let nodes = node_state.get_all_nodes();
|
if let Err(e) = parse(&mut current_parser, &stdin, &mut stdout) {
|
||||||
let node = nodes.get(i).unwrap().clone();
|
error!("Failed to parse: {}", e);
|
||||||
node_state.send_unrouted(node, &C2Packet::Aa).unwrap();
|
|
||||||
} else {
|
|
||||||
println!("");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
_ => {
|
|
||||||
warn!("Invalid command!")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// client.write(input)?;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn split_escape(input: &str) -> Vec<String> {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|||||||
+2
-2
@@ -1,4 +1,4 @@
|
|||||||
mod cli;
|
mod cli;
|
||||||
mod client_node;
|
mod node_cli;
|
||||||
|
|
||||||
pub use cli::Cli;
|
pub use cli::connect_cli;
|
||||||
|
|||||||
@@ -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<Box<dyn Cli>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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<String>) -> Result<(), Error> {
|
||||||
|
if let Some(subcommand) = &mut self.subcommand {
|
||||||
|
return subcommand.parse(input);
|
||||||
|
}
|
||||||
|
let parsed_command = CommandHolder::<NodeCliCommands>::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(())
|
||||||
|
}
|
||||||
|
}
|
||||||
+12
-13
@@ -1,14 +1,12 @@
|
|||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
|
|
||||||
use unshell_rs_lib::{
|
use unshell_rs_lib::{
|
||||||
Error,
|
C2Packet, Error,
|
||||||
nodes::{ConnectionConfig, Node},
|
nodes::{ConnectionConfig, NodeContainer},
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::C2Packet;
|
|
||||||
|
|
||||||
pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> {
|
pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> {
|
||||||
let node = Node::<C2Packet>::run_node(
|
let node = NodeContainer::connect(
|
||||||
"Server".to_string(),
|
"Server".to_string(),
|
||||||
vec![],
|
vec![],
|
||||||
vec![ConnectionConfig {
|
vec![ConnectionConfig {
|
||||||
@@ -18,16 +16,17 @@ pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> {
|
|||||||
)?;
|
)?;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
match node.rx.recv()? {
|
let (src, packet) = node.read_packet()?;
|
||||||
C2Packet::Aa => {
|
match packet {
|
||||||
info!("1");
|
C2Packet::Ping => {
|
||||||
|
info!("Ping from {}!", src);
|
||||||
|
node.send_unrouted(&src, &C2Packet::Pong)?;
|
||||||
|
// (&mut node.state.lock().unwrap()).send_unrouted(src, &C2Packet::Pong)?;
|
||||||
}
|
}
|
||||||
C2Packet::Bb => {
|
C2Packet::Pong => {
|
||||||
info!("2");
|
info!("Pong!");
|
||||||
}
|
|
||||||
C2Packet::Cc => {
|
|
||||||
info!("3");
|
|
||||||
}
|
}
|
||||||
|
_ => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+1
-3
@@ -3,12 +3,10 @@ extern crate log;
|
|||||||
|
|
||||||
mod client;
|
mod client;
|
||||||
mod endpoint;
|
mod endpoint;
|
||||||
mod packets;
|
|
||||||
|
|
||||||
pub use client::Cli;
|
pub use client::connect_cli;
|
||||||
|
|
||||||
pub use endpoint::run_endpoint;
|
pub use endpoint::run_endpoint;
|
||||||
pub use packets::C2Packet;
|
|
||||||
|
|
||||||
// pub use client::UnshellClient;
|
// pub use client::UnshellClient;
|
||||||
// pub use client::UnshellGui;
|
// pub use client::UnshellGui;
|
||||||
|
|||||||
+5
-5
@@ -7,11 +7,11 @@ use std::{
|
|||||||
|
|
||||||
use clap::{Parser, Subcommand};
|
use clap::{Parser, Subcommand};
|
||||||
use log::error;
|
use log::error;
|
||||||
use unshell_rs::{Cli, run_endpoint};
|
use unshell_rs::{connect_cli, run_endpoint};
|
||||||
use unshell_rs_lib::nodes::ConnectionConfig;
|
|
||||||
|
|
||||||
pub static DEFAULT_CONFIG_FILEPATH: &'static str = "server_config.json";
|
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
|
// The default port that this program looks for
|
||||||
pub static DEFAULT_SERVICE_PORT: u16 = 13370;
|
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
|
||||||
@@ -33,7 +33,7 @@ enum Commands {
|
|||||||
// Run as a service, and potentially hosting a website
|
// Run as a service, and potentially hosting a website
|
||||||
Relay {
|
Relay {
|
||||||
/// IPv4 to listen for clients on.
|
/// 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,
|
host: String,
|
||||||
|
|
||||||
/// Port listen to for command clients
|
/// Port listen to for command clients
|
||||||
@@ -143,7 +143,7 @@ fn main() -> Result<(), Box<dyn Error>> {
|
|||||||
// ),
|
// ),
|
||||||
Commands::Connect { host, port } => {
|
Commands::Connect { host, port } => {
|
||||||
let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
|
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
|
addr
|
||||||
} else {
|
} else {
|
||||||
error!("Could not parse address!");
|
error!("Could not parse address!");
|
||||||
@@ -153,7 +153,7 @@ fn main() -> Result<(), Box<dyn Error>> {
|
|||||||
Commands::Relay {
|
Commands::Relay {
|
||||||
host,
|
host,
|
||||||
port,
|
port,
|
||||||
config_filepath,
|
config_filepath: _,
|
||||||
} => {
|
} => {
|
||||||
let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
|
let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
|
||||||
run_endpoint(if let Ok(addr) = addr {
|
run_endpoint(if let Ok(addr) = addr {
|
||||||
|
|||||||
@@ -1,8 +0,0 @@
|
|||||||
use bincode::{Decode, Encode};
|
|
||||||
|
|
||||||
#[derive(Debug, Encode, Decode, Clone)]
|
|
||||||
pub enum C2Packet {
|
|
||||||
Aa,
|
|
||||||
Bb,
|
|
||||||
Cc,
|
|
||||||
}
|
|
||||||
@@ -8,3 +8,6 @@ static BINCODE_CONFIG: bincode::config::Configuration = bincode::config::standar
|
|||||||
pub mod layers;
|
pub mod layers;
|
||||||
pub mod networkers;
|
pub mod networkers;
|
||||||
pub mod nodes;
|
pub mod nodes;
|
||||||
|
mod packets;
|
||||||
|
|
||||||
|
pub use packets::C2Packet;
|
||||||
|
|||||||
@@ -1,8 +1,10 @@
|
|||||||
mod listener;
|
mod listener;
|
||||||
mod node;
|
mod node;
|
||||||
|
mod node_container;
|
||||||
mod packets;
|
mod packets;
|
||||||
|
mod stream;
|
||||||
|
|
||||||
pub use listener::ConnectionConfig;
|
pub use listener::ConnectionConfig;
|
||||||
pub use node::Node;
|
pub use node::Node;
|
||||||
// pub use packets::PacketError;
|
pub use node_container::NodeContainer;
|
||||||
pub use packets::Packets;
|
pub use stream::Stream;
|
||||||
|
|||||||
@@ -1,13 +1,17 @@
|
|||||||
use std::{
|
use std::{
|
||||||
collections::HashMap,
|
collections::HashMap,
|
||||||
fmt::Debug,
|
fmt::Debug,
|
||||||
sync::{Arc, Mutex},
|
sync::{
|
||||||
|
Arc, Mutex,
|
||||||
|
mpsc::{self, Receiver, Sender},
|
||||||
|
},
|
||||||
thread,
|
thread,
|
||||||
time::Duration,
|
time::Duration,
|
||||||
};
|
};
|
||||||
|
|
||||||
use bincode::{Decode, Encode};
|
use bincode::{Decode, Encode};
|
||||||
use crossbeam_channel::{Receiver, Sender};
|
// use std:::{Receiver, Sender};
|
||||||
|
#[allow(deprecated)]
|
||||||
use rand::{seq::IndexedRandom, thread_rng};
|
use rand::{seq::IndexedRandom, thread_rng};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
@@ -20,16 +24,6 @@ use crate::{
|
|||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
pub struct NodeState<P>
|
|
||||||
where
|
|
||||||
P: Encode + Decode<()> + Debug + Clone + 'static,
|
|
||||||
{
|
|
||||||
id: String,
|
|
||||||
connections: HashMap<String, Box<dyn Connection + Send>>,
|
|
||||||
map: HashMap<String, Vec<String>>,
|
|
||||||
packet_listener: Sender<P>,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn read(c: &mut Box<dyn Connection + Send>) -> Result<Packets, Error> {
|
fn read(c: &mut Box<dyn Connection + Send>) -> Result<Packets, Error> {
|
||||||
Packets::decode(c.read()?.as_slice())
|
Packets::decode(c.read()?.as_slice())
|
||||||
}
|
}
|
||||||
@@ -43,7 +37,7 @@ where
|
|||||||
P: Encode + Decode<()> + Debug + Clone + 'static,
|
P: Encode + Decode<()> + Debug + Clone + 'static,
|
||||||
{
|
{
|
||||||
pub state: Arc<Mutex<NodeState<P>>>,
|
pub state: Arc<Mutex<NodeState<P>>>,
|
||||||
pub rx: Receiver<P>,
|
pub rx: Receiver<(String, P)>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<P> Node<P>
|
impl<P> Node<P>
|
||||||
@@ -60,7 +54,7 @@ where
|
|||||||
{
|
{
|
||||||
// let mut parent = build_client(TCPClient::connect(&parent.socket)?, parent.layers)?;
|
// 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::<P> {
|
let state = Arc::new(Mutex::new(NodeState::<P> {
|
||||||
id: id, //Uuid::new_v4().to_string(), //TODO: Calling an OS RNG can pose a problem for security;
|
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()))?;
|
write(&mut connection, Packets::SyncUUID(this_uuid.clone()))?;
|
||||||
|
|
||||||
// Recieve UUID
|
// 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
|
source
|
||||||
} else {
|
} 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);
|
info!("New Node! {} (direct)", other_uuid);
|
||||||
|
|
||||||
// Add connection
|
// Add connection
|
||||||
@@ -189,6 +193,16 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub struct NodeState<P>
|
||||||
|
where
|
||||||
|
P: Encode + Decode<()> + Debug + Clone + 'static,
|
||||||
|
{
|
||||||
|
id: String,
|
||||||
|
connections: HashMap<String, Box<dyn Connection + Send>>,
|
||||||
|
map: HashMap<String, Vec<String>>,
|
||||||
|
packet_listener: Sender<(String, P)>,
|
||||||
|
}
|
||||||
|
|
||||||
impl<P> NodeState<P>
|
impl<P> NodeState<P>
|
||||||
where
|
where
|
||||||
P: Encode + Decode<()> + Debug + Clone + Send + 'static,
|
P: Encode + Decode<()> + Debug + Clone + Send + 'static,
|
||||||
@@ -207,7 +221,7 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn knows_client(&self, id: &String) -> bool {
|
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
|
// Remove all nodes where the routes are empty
|
||||||
@@ -381,7 +395,7 @@ where
|
|||||||
|
|
||||||
fn route_packet(&mut self, src: String, dest: String, data: Vec<u8>) -> Result<(), Error> {
|
fn route_packet(&mut self, src: String, dest: String, data: Vec<u8>) -> Result<(), Error> {
|
||||||
if dest == self.id {
|
if dest == self.id {
|
||||||
self.packet_listener.send(decode_vec::<P>(&data)?)?;
|
self.packet_listener.send((src, decode_vec::<P>(&data)?))?;
|
||||||
} else {
|
} else {
|
||||||
if self.connections.contains_key(&dest) {
|
if self.connections.contains_key(&dest) {
|
||||||
write(
|
write(
|
||||||
@@ -389,6 +403,7 @@ where
|
|||||||
Packets::DataUnrouted { src, dest, data },
|
Packets::DataUnrouted { src, dest, data },
|
||||||
)?;
|
)?;
|
||||||
} else if self.map.contains_key(&dest) {
|
} else if self.map.contains_key(&dest) {
|
||||||
|
#[allow(deprecated)]
|
||||||
let next_uuid = self
|
let next_uuid = self
|
||||||
.map
|
.map
|
||||||
.get(&dest)
|
.get(&dest)
|
||||||
|
|||||||
@@ -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<Mutex<HashMap<(usize, String), Option<(Stream, Sender<Vec<u8>>)>>>>;
|
||||||
|
|
||||||
|
pub struct NodeContainer {
|
||||||
|
streams: Streams,
|
||||||
|
state: Arc<Mutex<NodeState<TransportLayerPacket>>>,
|
||||||
|
spontanious_rx: Receiver<(String, C2Packet)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl NodeContainer {
|
||||||
|
pub fn connect(
|
||||||
|
id: String,
|
||||||
|
clients: Vec<ConnectionConfig>,
|
||||||
|
listeners: Vec<ConnectionConfig>,
|
||||||
|
) -> Result<Self, Error> {
|
||||||
|
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<TransportLayerPacket>,
|
||||||
|
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::<C2Packet>(&data)?))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn create_stream(
|
||||||
|
remote_stream_id: usize,
|
||||||
|
dest: String,
|
||||||
|
node: &Node<TransportLayerPacket>,
|
||||||
|
stream_mut: &mut Option<(Stream, Sender<Vec<u8>>)>,
|
||||||
|
) -> 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<String> {
|
||||||
|
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()?)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -13,15 +13,24 @@ pub enum Packets {
|
|||||||
Disconnect {
|
Disconnect {
|
||||||
routes: Vec<String>,
|
routes: Vec<String>,
|
||||||
},
|
},
|
||||||
|
|
||||||
|
// Send single data packet without routing details
|
||||||
DataUnrouted {
|
DataUnrouted {
|
||||||
src: String,
|
src: String,
|
||||||
dest: String,
|
dest: String,
|
||||||
data: Vec<u8>,
|
data: Vec<u8>,
|
||||||
},
|
},
|
||||||
|
// Send single data packet with routing details
|
||||||
DataRouted {
|
DataRouted {
|
||||||
path: Vec<String>,
|
path: Vec<String>,
|
||||||
data: Vec<u8>,
|
data: Vec<u8>,
|
||||||
},
|
},
|
||||||
|
|
||||||
|
// DataStreamRouted {
|
||||||
|
// path: Vec<String>,
|
||||||
|
// data: Vec<u8>,
|
||||||
|
// },
|
||||||
|
ErrorNameExists,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Packets {
|
impl Packets {
|
||||||
|
|||||||
@@ -0,0 +1,36 @@
|
|||||||
|
use crossbeam_channel::{Receiver, Sender};
|
||||||
|
|
||||||
|
use crate::networkers::Connection;
|
||||||
|
|
||||||
|
pub struct Stream {
|
||||||
|
tx: Sender<Vec<u8>>,
|
||||||
|
rx: Receiver<Vec<u8>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Connection for Stream {
|
||||||
|
fn get_info(&self) -> String {
|
||||||
|
"unrouted".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_alive(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read(&mut self) -> Result<Vec<u8>, 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<Box<dyn Connection + Send + Sync>, crate::Error> {
|
||||||
|
todo!()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Stream {
|
||||||
|
pub fn new(tx: Sender<Vec<u8>>, rx: Receiver<Vec<u8>>) -> Self {
|
||||||
|
Self { tx, rx }
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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<u8>,
|
||||||
|
},
|
||||||
|
|
||||||
|
SpontaniousDataUnrouted {
|
||||||
|
data: Vec<u8>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Encode, Decode, Clone)]
|
||||||
|
pub enum C2Packet {
|
||||||
|
Ping,
|
||||||
|
Pong,
|
||||||
|
|
||||||
|
CreatePTY { width: usize, height: usize },
|
||||||
|
PTYData,
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user