Work on adding cli, and transport layer

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<String>) -> Result<(), Error>;
+}
 
-        let stdin = std::io::stdin();
-        let mut stdout = std::io::stdout();
+#[derive(Debug, Parser)]
+pub struct CommandHolder<P>
+where
+    P: Subcommand,
+{
+    #[command(subcommand)]
+    pub command: P,
+}
 
-        let node = Node::<C2Packet>::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<dyn Cli>;
 
-        //     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<dyn Cli + 'static>,
+                 stdin: &Stdin,
+                 stdout: &mut Stdout|
+     -> Result<(), Error> {
+        let name = current_parser.name();
+        print!("Unshell | {}> ", name);
+        stdout.flush()?;
 
-        let selected_node: Option<usize> = 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::<usize>(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<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
+}

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;

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<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(())
+    }
+}

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::<C2Packet>::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!");
             }
+            _ => {}
         }
     }
 }

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;

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<dyn Error>> {
         // ),
         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<dyn Error>> {
         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 {

src/packets.rs

@@ -1,8 +0,0 @@
-use bincode::{Decode, Encode};
-
-#[derive(Debug, Encode, Decode, Clone)]
-pub enum C2Packet {
-    Aa,
-    Bb,
-    Cc,
-}