Work on adding cli, and transport layer

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;

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;

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<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> {
     Packets::decode(c.read()?.as_slice())
 }
@@ -43,7 +37,7 @@ where
     P: Encode + Decode<()> + Debug + Clone + 'static,
 {
     pub state: Arc<Mutex<NodeState<P>>>,
-    pub rx: Receiver<P>,
+    pub rx: Receiver<(String, P)>,
 }
 
 impl<P> Node<P>
@@ -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::<P> {
             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<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>
 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<u8>) -> Result<(), Error> {
         if dest == self.id {
-            self.packet_listener.send(decode_vec::<P>(&data)?)?;
+            self.packet_listener.send((src, decode_vec::<P>(&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)

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

unshell-rs-lib/src/nodes/packets.rs

@@ -13,15 +13,24 @@ pub enum Packets {
     Disconnect {
         routes: Vec<String>,
     },
+
+    // Send single data packet without routing details
     DataUnrouted {
         src: String,
         dest: String,
         data: Vec<u8>,
     },
+    // Send single data packet with routing details
     DataRouted {
         path: Vec<String>,
         data: Vec<u8>,
     },
+
+    // DataStreamRouted {
+    //     path: Vec<String>,
+    //     data: Vec<u8>,
+    // },
+    ErrorNameExists,
 }
 
 impl Packets {

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

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<u8>,
+    },
+
+    SpontaniousDataUnrouted {
+        data: Vec<u8>,
+    },
+}
+
+#[derive(Debug, Encode, Decode, Clone)]
+pub enum C2Packet {
+    Ping,
+    Pong,
+
+    CreatePTY { width: usize, height: usize },
+    PTYData,
+}