From 4ce461025cecf30861ad80b490d562c6133feac2 Mon Sep 17 00:00:00 2001 From: pyro Date: Tue, 4 Aug 2026 13:35:52 -0500 Subject: [PATCH] made the server and the client/server communication functionality work. --- Cargo.lock | 336 ++++++++++++++++++++++++++++++++++++++- Cargo.toml | 5 + src/funcs.rs | 419 +++++++++++++++++++++--------------------------- src/install.rs | 3 + src/lib.rs | 424 +++++++++++++++++++++++++++++++------------------ src/main.rs | 96 ++++++----- src/server.rs | 317 ++++++++++++++++++++++++++++++++++++ 7 files changed, 1162 insertions(+), 438 deletions(-) create mode 100644 src/server.rs diff --git a/Cargo.lock b/Cargo.lock index c3f6d4a..3b3d0fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -142,9 +142,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "arrayref" @@ -173,6 +173,45 @@ dependencies = [ "libloading", ] +[[package]] +name = "asn1-rs" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f43a50ac4fdca5df8e885c21b835997f0a1cdee65494a6847694a98652d9d8" +dependencies = [ + "asn1-rs-derive", + "asn1-rs-impl", + "displaydoc", + "nom", + "num-traits", + "rusticata-macros", + "thiserror 2.0.18", + "time", +] + +[[package]] +name = "asn1-rs-derive" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", + "synstructure", +] + +[[package]] +name = "asn1-rs-impl" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "async-broadcast" version = "0.7.2" @@ -325,6 +364,29 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "aws-lc-rs" +version = "1.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.43.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43103168cc76fe62678a375e722fc9cb3a0146159ac5828bc4f0dfd755c2224c" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + [[package]] name = "base64" version = "0.22.1" @@ -361,6 +423,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" +dependencies = [ + "serde", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -614,6 +685,15 @@ dependencies = [ "x11rb", ] +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "codespan-reporting" version = "0.12.0" @@ -936,12 +1016,32 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "data-encoding" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" + [[package]] name = "deltae" version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5729f5117e208430e437df2f4843f5e5952997175992d1414f94c57d61e270b4" +[[package]] +name = "der-parser" +version = "10.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6" +dependencies = [ + "asn1-rs", + "displaydoc", + "nom", + "num-bigint", + "num-traits", + "rusticata-macros", +] + [[package]] name = "deranged" version = "0.5.8" @@ -999,6 +1099,17 @@ dependencies = [ "objc2 0.6.4", ] +[[package]] +name = "displaydoc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "dlib" version = "0.5.3" @@ -1035,6 +1146,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "edf234dd1594d6dd434a8fb8cada51ddbbc593e40e4a01556a0b31c62da2775b" +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "either" version = "1.16.0" @@ -2287,6 +2404,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -2304,6 +2431,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -2704,6 +2840,15 @@ dependencies = [ "objc2-foundation 0.2.2", ] +[[package]] +name = "oid-registry" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7" +dependencies = [ + "asn1-rs", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2801,6 +2946,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -3200,6 +3355,20 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rcgen" +version = "0.14.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57f6d249aad744e274e682777a50283a225a32705394ee6d5fcc01efa25e4055" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "x509-parser", + "yasna", +] + [[package]] name = "read-fonts" version = "0.35.0" @@ -3314,6 +3483,20 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + [[package]] name = "roxmltree" version = "0.20.0" @@ -3341,6 +3524,15 @@ dependencies = [ "semver", ] +[[package]] +name = "rusticata-macros" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632" +dependencies = [ + "nom", +] + [[package]] name = "rustix" version = "0.38.44" @@ -3367,6 +3559,51 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.23.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" +dependencies = [ + "aws-lc-rs", + "log", + "once_cell", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "aws-lc-rs", + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -3772,6 +4009,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "svg_fmt" version = "0.4.5" @@ -3811,6 +4054,28 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "sys-locale" version = "0.3.2" @@ -3909,6 +4174,7 @@ dependencies = [ name = "tetanus" version = "0.1.0" dependencies = [ + "anyhow", "clap", "crossterm", "fs_extra", @@ -3916,9 +4182,13 @@ dependencies = [ "ipnet", "ratatui", "rayon", + "rcgen", "rhai", "rustc-hash 2.1.2", + "rustls", + "rustls-pemfile", "tokio", + "tokio-rustls", "walkdir", ] @@ -3978,12 +4248,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" dependencies = [ "deranged", + "itoa", "libc", "num-conv", "num_threads", "powerfmt", "serde_core", "time-core", + "time-macros", ] [[package]] @@ -3992,6 +4264,16 @@ version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" +[[package]] +name = "time-macros" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tiny-keccak" version = "2.0.2" @@ -4082,6 +4364,16 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -4229,6 +4521,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "utf8parse" version = "0.2.2" @@ -5275,6 +5573,24 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd" +[[package]] +name = "x509-parser" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202" +dependencies = [ + "asn1-rs", + "data-encoding", + "der-parser", + "lazy_static", + "nom", + "oid-registry", + "ring", + "rusticata-macros", + "thiserror 2.0.18", + "time", +] + [[package]] name = "xcursor" version = "0.3.10" @@ -5306,6 +5622,16 @@ version = "0.8.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ae8337f8a065cfc972643663ea4279e04e7256de865aa66fe25cec5fb912d3f" +[[package]] +name = "yasna" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282" +dependencies = [ + "bit-vec 0.9.1", + "time", +] + [[package]] name = "yazi" version = "0.2.1" @@ -5399,6 +5725,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index 16d318d..4f1ed90 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] +anyhow = "1.0.103" clap = { version = "4.6.1", features = ["derive"] } crossterm = "0.29.0" fs_extra = "1.3.0" @@ -11,7 +12,11 @@ iced = { version = "0.14.0", features = ["advanced", "tokio"] } ipnet = "2.12.0" ratatui = "0.30.0" rayon = "1.12.0" +rcgen = "0.14.8" rhai = { version = "1.24.0", features = ["metadata", "sync"] } rustc-hash = "2.1.2" +rustls = "0.23.43" +rustls-pemfile = "2.2.0" tokio = { version = "1.52.3", features = ["full"] } +tokio-rustls = "0.26.4" walkdir = "2.5.0" diff --git a/src/funcs.rs b/src/funcs.rs index 8c953fc..513c792 100644 --- a/src/funcs.rs +++ b/src/funcs.rs @@ -1,31 +1,20 @@ -use crate::{ToolMessage::Input, *}; +use crate::*; use crossterm::{ cursor, event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind}, execute, terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode}, }; +use iced::keyboard::key::Code::Sleep; use ratatui::{ prelude::*, widgets::{Block, Borders, List, ListItem, ListState, Paragraph}, }; -use std::sync::Arc; -use std::sync::Mutex; use std::time::Duration; use std::{error::Error, time::Instant}; -use std::{ - io::{Read, Write}, - usize, -}; -use std::{net::TcpStream, thread}; - -pub enum ServerBrokerCmd { - ConnectServer(usize), - RegisterActionOutput((Server, String)), - Disconnect(Server), - DisconnectAll, - Exit, -} +use std::{io::Write, usize}; +use std::{sync::Arc, thread::sleep}; +use std::{sync::Mutex, thread::spawn}; pub fn get_user_input(prompt: &str) -> Result> { println!("{}", prompt); @@ -34,158 +23,6 @@ pub fn get_user_input(prompt: &str) -> Result> { return Ok(response.trim().to_string()); } -pub fn server_broker( - rx: Receiver, - tx: Sender, - servers: Arc>>, -) { - loop { - if let Ok(msg) = rx.try_recv() { - match msg { - ServerBrokerCmd::ConnectServer(id) => { - let _ = tx.send(ToolMessage::Output(( - 0, - format!("attempting to connect to server {}", id), - ))); - if let Ok(mut lock) = servers.lock() { - lock.iter_mut().for_each(|s| { - let _ = - tx.send(ToolMessage::Output((0, format!("got server list lock!")))); - if s.id == id { - let _ = tx.send(ToolMessage::Output(( - 0, - format!("Found the server for {}: {}", s.id, s.address), - ))); - if let Ok(mut stream) = TcpStream::connect(s.address.clone()) { - let _ = tx.send(ToolMessage::Output(( - 0, - String::from("TCP successful to server!"), - ))); - let _ = stream.write(format!("HELLO|{}\n", s.id).as_bytes()); - let _ = tx.send(ToolMessage::Output(( - 0, - String::from("Hello message sent!"), - ))); - let mut buf = [0; 1024]; - if let Ok(bytes_read) = stream.read(&mut buf) { - let response = String::from_utf8_lossy(&buf[..bytes_read]); - if let Some((_, id)) = response.split_once("|") { - if let Ok(id) = id.trim().parse::() { - s.client_id = id; - s.connected = true; - let _ = tx.send(ToolMessage::Output(( - 0, - String::from("got client ID from server!"), - ))); - } else { - let _ = tx.send(ToolMessage::Output(( - 0, - String::from( - "failed to get response from server!", - ), - ))); - } - } - } - } - } - }); - } - } - ServerBrokerCmd::Disconnect(server) => { - if let Ok(mut lock) = servers.lock() { - let mut rid = 0; - let mut found = false; - lock.iter().enumerate().for_each(|(id, s)| { - if s.id == server.id { - rid = id; - found = true; - } - }); - if found { - lock.remove(rid); - } - } - } - ServerBrokerCmd::DisconnectAll => { - println!("todo"); - } - ServerBrokerCmd::Exit => { - println!("todo"); - break; - } - ServerBrokerCmd::RegisterActionOutput(text) => { - if let Ok(mut lock) = servers.lock() { - lock.iter_mut().for_each(|s| { - if s.id == text.0.id { - s.message_que.push(text.1.clone()); - } - }); - } - } - } - } - if let Ok(mut lock) = servers.try_lock() { - lock.iter_mut().for_each(|s| { - if s.connected { - if Instant::now().duration_since(s.last_check) >= s.timer { - if let Ok(mut stream) = TcpStream::connect(s.address.clone()) { - s.last_check = Instant::now(); - let mut success_ids = Vec::new(); - if s.message_que.len() > 0 { - s.message_que.iter().enumerate().for_each(|(id, msg)| { - if let Ok(_) = stream.write( - format!("{}|{}\n", s.client_id.clone(), msg).as_bytes(), - ) { - success_ids.push(id); - } - }); - } else { - let _ = stream.write(format!("{}|PING", s.client_id).as_bytes()); - } - for id in success_ids { - s.message_que.remove(id); - } - let mut buffer = [0; 1024]; - if let Ok(bytes_read) = stream.read(&mut buffer) { - let response = - String::from_utf8_lossy(&buffer[..bytes_read]).to_string(); - let (_, action_string) = - response.split_once("|").unwrap_or(("ACTIONS", "NONE")); - let actions: Vec<&str> = action_string.split(",").collect(); - for action in actions { - let (action, data) = - action.split_once("|").unwrap_or(("PING", "NONE")); - match action { - "OUTPUT" => { - let _ = tx.send(ToolMessage::Output(( - 0, - format!("FROM SERVER: {}", data), - ))); - } - "CMD" => { - let (rid_str, cmd) = - data.split_once("|").unwrap_or(("0", "NONE")); - if let Ok(rid) = rid_str.parse::() { - let _ = tx.send(ToolMessage::Input(( - rid, - cmd.to_string(), - ))); - } - } - _ => {} - } - } - } - } - } - } - }); - } - thread::sleep(Duration::from_millis(10)); - } -} - pub fn run_tui( mut state: AppState, main_rx: Receiver, @@ -198,12 +35,6 @@ pub fn run_tui( cursor::Hide, EnableMouseCapture )?; - let (server_tx, server_rx) = channel(); - let main_tx = state.main_tx.clone(); - let server_clone = state.servers.clone(); - std::thread::spawn(move || { - server_broker(server_rx, main_tx, server_clone); - }); let backend = CrosstermBackend::new(&stdout); let mut terminal = Terminal::new(backend)?; let (event_tx, event_rx) = channel::(); @@ -250,6 +81,7 @@ pub fn run_tui( project_list_state.select(Some(0)); } let mut history_index = state.history.len(); + let mut handles = Vec::new(); loop { terminal.draw(|f| { let main_chunks = Layout::default() @@ -402,17 +234,22 @@ pub fn run_tui( } ToolMessage::Input(cmd) => { state.log.push(cmd.1.clone()); + if cmd.1.contains(" ") { + let (gcmd, args) = cmd.1.split_once(" ").unwrap(); + state.execute_command(gcmd, Some(args.to_string()), cmd.0)?; + } else { + state.execute_command(cmd.1.as_str(), None, cmd.0)?; + } } ToolMessage::Output((rid, txt)) => { state.log.push(txt.clone()); state.output.push(txt.clone()); state.output_scroll = u16::MAX; if rid != 0 { - if let Some(server) = state.selected_server.clone() { - let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(( - server, - format!("OUTPUT|{}|{}", rid, txt), - ))); + if let Some(server) = state.servers.get(state.selected_server) { + if let Ok(mut server) = server.lock() { + server.message_que.push(format!("OUTPUT|{}", txt)); + } } } } @@ -482,12 +319,121 @@ pub fn run_tui( terminal.clear()?; } ToolMessage::ConnectServer => { - if let Some(server) = state.selected_server.clone() { - let _ = server_tx.send(ServerBrokerCmd::ConnectServer(server.id)); + if let Some(server) = state.servers.get_mut(state.selected_server) { + let mut address = String::new(); + let lock_res = server.try_lock(); + let mut connected = false; + match lock_res { + Ok(mut locked_server) => match locked_server.connect() { + Ok(id) => { + state.output.push(format!( + "Server Connected! ID:{} ADDRESS:{}", + id, + locked_server.address.clone() + )); + address = locked_server.address.clone(); + connected = true; + } + Err(e) => { + state + .output + .push(format!("Error connecting to server {e}")); + } + }, + Err(e) => { + state + .output + .push(format!("Error locking server state! {e}")); + } + } + if connected { + let server_clone = server.clone(); + let tx_clone = state.main_tx.clone(); + let mut timer = Duration::from_secs(5); + let mut last_check = Instant::now(); + let server_handle = spawn(move || { + let server = server_clone; + loop { + let mut stop = false; + if Instant::now().duration_since(last_check) >= timer { + let lock_res = server.try_lock(); + match lock_res { + Ok(mut locked_server) => { + locked_server.checkin(); + last_check = locked_server.last_check.clone(); + for action in locked_server.action_que.clone() { + if action.contains("|") { + let (action, data) = + action.split_once("|").unwrap(); + match action.trim() { + "OUTPUT" => { + let _ = tx_clone.send( + ToolMessage::Output(( + 0, + data.trim().to_string(), + )), + ); + } + "CMD" => { + let _ = tx_clone.send( + ToolMessage::Input(( + locked_server + .client_id + .clone(), + data.trim().to_string(), + )), + ); + } + _ => {} + } + } else { + match action.as_str() { + "disconnect" => { + locked_server.connected = false; + stop = true; + } + _ => { + let _ = tx_clone.send( + ToolMessage::Output(( + 0, + action.clone(), + )), + ); + } + } + } + } + timer = locked_server.timer.clone(); + locked_server.action_que.clear(); + if locked_server.message_que.is_empty() { + let id = locked_server.client_id.clone(); + locked_server + .message_que + .push(format!("{}|||PING\n", id)); + } + } + Err(_) => {} + } + } + if stop { + println!("{} disconnected!", address); + break; + } + sleep(Duration::from_secs(1)); + } + }); + handles.push(server_handle); + } } } ToolMessage::SendServer(data) => { - let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(data)); + if let Some(server) = state.servers.get(state.selected_server) { + if let Ok(mut server) = server.lock() { + server + .message_que + .push(format!("{}|||{}\n", data.0, data.1)); + } + } } ToolMessage::DisconnectAllServers => { println!("todo"); @@ -530,73 +476,55 @@ pub fn run_tui( } let (command, args) = trimmed.split_once(' ').unwrap_or((&trimmed, "")); - if state.selected_client != 0 { - if let Some(server) = state.selected_server.clone() { - let _ = server_tx.send( - ServerBrokerCmd::RegisterActionOutput(( - server.clone(), - format!( - "CMD|{}|{}|{} {}", - server.client_id, - state.selected_client, - command, - args - ), - )), - ); + match command { + "exit" | "quit" => break, + "reload-modules" => { + state.initialize_modules(); + state.output.push("Reloading module paths...".into()); } - } else { - match command { - "exit" | "quit" => break, - "reload-modules" => { - state.initialize_modules(); - state.output.push("Reloading module paths...".into()); + "help" => { + let help_text = state.help.clone().join("\n"); + for line in help_text.lines() { + state.output.push(line.to_string()); } - "help" => { - let help_text = state.help.clone().join("\n"); - for line in help_text.lines() { - state.output.push(line.to_string()); + } + "new_project" | "np" => { + if args.split_once(' ').is_some() { + let _ = state.execute_command( + command, + Some(args.to_string()), + 0, + ); + } else { + state + .output + .push("Error: USAGE -> np ".into()); + } + } + command_name => { + if state.module_loader.commands.contains_key(command) { + state.output.push(format!( + "[Worker] Executing script '{}'...", + command_name + )); + if let Err(e) = state.execute_command( + command, + Some(args.to_string()), + 0, + ) { + state + .output + .push(format!("[Error] Pipeline fail: {}", e)); } - } - "new_project" | "np" => { - if args.split_once(' ').is_some() { + } else { + if args == "" { + let _ = state.execute_command(command, None, 0); + } else { let _ = state.execute_command( command, Some(args.to_string()), 0, ); - } else { - state - .output - .push("Error: USAGE -> np ".into()); - } - } - command_name => { - if state.module_loader.commands.contains_key(command) { - state.output.push(format!( - "[Worker] Executing script '{}'...", - command_name - )); - if let Err(e) = state.execute_command( - command, - Some(args.to_string()), - 0, - ) { - state.output.push(format!( - "[Error] Pipeline fail: {}", - e - )); - } - } else { - if args == "" { - let _ = state.execute_command(command, None, 0); - } else { - let _ = state.execute_command( - command, - Some(args.to_string()), - 0, - ); - } } } } @@ -700,5 +628,14 @@ pub fn run_tui( cursor::Show, DisableMouseCapture )?; + println!("disconnecting from servers..."); + state.servers.iter().for_each(|server| { + if let Ok(mut server_lock) = server.lock() { + server_lock.action_que.push("disconnect".to_string()); + } + }); + for handle in handles { + handle.join().unwrap(); + } Ok(()) } diff --git a/src/install.rs b/src/install.rs index 7ed8a93..f11fc8b 100644 --- a/src/install.rs +++ b/src/install.rs @@ -69,6 +69,8 @@ pub fn client_install() -> Result<(), Box> { let term_cmd = get_user_input( "for example: konsole -e |||COMMAND||| or konsole --profile profile_name ENV_NAME=CURRENT_PROJECT_BOX", )?; + let tools_folder = + get_user_input("path where your tools are installed ex /opt or /home/user/tools")?; config_path.push("client.conf"); projects_path.push("default"); create_dir_all(&projects_path)?; @@ -100,6 +102,7 @@ pub fn client_install() -> Result<(), Box> { client_config_file.write(format!("module_path: {}\n", module_path.display()).as_bytes())?; client_config_file.write(format!("template_box: {}\n", template_box).as_bytes())?; client_config_file.write(format!("term_cmd: {}\n", term_cmd).as_bytes())?; + client_config_file.write(format!("tools: {}\n", tools_folder).as_bytes())?; println!("\ndefault config files written!"); println!("downloading default notes and modules..."); create_dir_all("./temp")?; diff --git a/src/lib.rs b/src/lib.rs index 4aff421..ef8e210 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,27 +2,72 @@ use fs_extra::dir::{CopyOptions, copy}; use ipnet::IpNet; use ratatui::crossterm::event; use rhai::{AST, Dynamic, Engine, Scope}; +use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName}; use std::collections::HashMap; use std::error::Error; use std::fs::{self, File, create_dir_all, read_dir, read_to_string, remove_dir, remove_dir_all}; -use std::io::{self, BufRead, BufReader, Write}; +use std::io::{self, BufRead, BufReader, Read, Write}; use std::net::IpAddr; +use std::net::TcpStream; use std::path::PathBuf; use std::process::{Command, Stdio}; use std::sync::mpsc::Receiver; use std::sync::{Arc, Mutex, mpsc::Sender, mpsc::channel}; - -use std::time; use std::time::Duration; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpStream; +use std::time::{self, Instant}; pub mod funcs; pub mod server; +enum AppEvent { + Key(event::KeyEvent), + Worker(ToolMessage), + Mouse(event::MouseEvent), +} + +#[derive(Clone, Debug)] +pub enum Destination { + Server, + Victim, + Attacker, + Control, +} + +#[derive(Clone, Debug)] +pub enum ToolMessage { + Input((usize, String)), + Output((usize, String)), + UpdateHost(Host), + RebuildDB, + DestroyDB(DistroBox), + StopDB(DistroBox), + StopTemplate(DistroBox), + UpdateProject(usize, Project), + RemoveProject, + AddServer, + EndPrompt, + ConnectServer, + SelectServer, + SendServer((usize, String)), + ServerBrokerExit, + AppStateExit, + DisconnectServer, + DisconnectAllServers, +} + +#[derive(Clone, Debug)] +pub enum ToolArg { + Project(Project), + Projects(Vec), + Host(Host), + Hosts(Vec), + Config(HashMap), + Path(PathBuf), +} + pub struct AppState { pub projects: Vec, - pub servers: Arc>>, + pub servers: Vec>>, pub config: HashMap, pub config_file: PathBuf, pub workers: rayon::ThreadPool, @@ -31,9 +76,8 @@ pub struct AppState { pub history: Vec, pub log: Vec, pub output: Vec, - pub selected_server: Option, + pub selected_server: usize, pub selected_project: usize, - pub selected_client: usize, pub curent_intput: String, pub module_loader: ModuleLoader, pub output_scroll: u16, @@ -51,7 +95,7 @@ impl AppState { ( Self { projects: Vec::new(), - servers: Arc::new(Mutex::new(Vec::new())), + servers: Vec::new(), config: HashMap::new(), config_file: PathBuf::new(), workers: rayon::ThreadPoolBuilder::new() @@ -63,9 +107,8 @@ impl AppState { history: Vec::new(), log: Vec::new(), output: Vec::new(), - selected_server: None, + selected_server: 0, selected_project: 0, - selected_client: 0, curent_intput: String::new(), module_loader: ModuleLoader::new(), output_scroll: 0, @@ -100,14 +143,14 @@ impl AppState { address: address.trim().to_string(), connected: false, timer: Duration::from_secs(60), - id: 0, last_check: time::Instant::now(), message_que: Vec::new(), + action_que: Vec::new(), client_id: 0, + selected_client: 0, + config: None, }; - if let Ok(mut lock) = self.servers.lock() { - lock.push(new_server); - } + self.servers.push(Arc::new(Mutex::new(new_server))); } self.config .insert("servers".to_string(), line.trim().to_string()); @@ -223,6 +266,27 @@ impl AppState { ) -> Result<(), Box> { let tx = self.main_tx.clone(); match command_name { + "remote" => { + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + if server.selected_client != 0 { + if let Some(args) = command_args.clone() { + let args = args.trim(); + if args.contains(" ") { + let (cmd, iargs) = args.split_once(" ").unwrap(); + server.message_que.push(format!("CMD|{} {}", cmd, iargs)); + } else { + server.message_que.push(format!("CMD|{}", args)); + } + let _ = self.main_tx.send(ToolMessage::Output(( + 0, + format!("Tasked {} to run {}", server.selected_client, args), + ))); + } + } + } + } + } "new_project" | "np" => { let args = command_args.unwrap(); let (org, name) = args.split_once(" ").unwrap(); @@ -407,17 +471,14 @@ impl AppState { address: format!("{}:{}", self.prompt.responses[0], self.prompt.responses[1]), connected: false, timer: Duration::from_secs(5), - id: 0, last_check: time::Instant::now(), message_que: Vec::new(), + action_que: Vec::new(), client_id: 0, + selected_client: 0, + config: None, }; - if let Ok(mut lock) = self.servers.lock() { - if lock.len() == 0 { - self.selected_server = Some(new_server.clone()); - } - lock.push(new_server); - } + self.servers.push(Arc::new(Mutex::new(new_server))); self.prompt.reset(); } "connect_server" => { @@ -440,71 +501,64 @@ impl AppState { } } "list_servers" => { - let mut selected_addr = String::new(); - if let Some(selected) = self.selected_server.clone() { - selected_addr = selected.address.clone(); - } - if let Ok(lock) = self.servers.lock() { - lock.iter().for_each(|s| { - if selected_addr == s.address { + self.servers.iter().enumerate().for_each(|(id, m)| { + if let Ok(s) = m.lock() { + if id == self.selected_server { let _ = self.main_tx.send(ToolMessage::Output(( rid, - format!("{} - Currently Selected", s.address), + format!( + "Server ID:{} Address:{} ClientID:{} Sleep:{} Last:{} - Currently Selected", + id, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs() + ), ))); } else { - let _ = self - .main_tx - .send(ToolMessage::Output((rid, s.address.clone()))); + let _ = self.main_tx.send(ToolMessage::Output(( + rid, + format!( + "Server ID:{} Address:{} ClientID:{} Sleep:{} Last:{}", + id, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs() + ), + ))); } - }); - } + } + }); } "select_server" => { self.prompt.action = Some(ToolMessage::SelectServer); self.prompt.num_responses = 1; self.prompt.prompts.push(String::from("Selection?")); self.prompt.execute_command = String::from("server_selected"); - let mut counter: usize = 0; - let mut selected_addr = String::new(); - if let Some(selected) = self.selected_server.clone() { - selected_addr = selected.address.clone(); - } - if let Ok(lock) = self.servers.lock() { - lock.iter().for_each(|server| { - if server.address == selected_addr { + self.servers.iter().enumerate().for_each(|(id, servermut)| { + if let Ok(server) = servermut.lock() { + if id == self.selected_server { let _ = self.main_tx.send(ToolMessage::Output(( rid, - format!( - "{}:{} - Currently Selected", - counter, - server.address.clone() - ), + format!("{}:{} - Currently Selected", id, server.address.clone()), ))); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, - format!("{}: {}", counter, server.address.clone()), + format!("{}: {}", id, server.address.clone()), ))); } - counter += 1; - }); - } + } + }); } "server_selected" => { if self.prompt.responses.len() == 1 { - if let Ok(lock) = self.servers.lock() { - if let Ok(id) = self.prompt.responses[0].parse::() { - self.selected_server = Some(lock[id].clone()); + if let Ok(id) = self.prompt.responses[0].parse::() { + self.selected_server = id; + if let Ok(server) = self.servers[id].lock() { let _ = self.main_tx.send(ToolMessage::Output(( rid, - format!("{} Selected!", lock[id].address), - ))); - } else { - let _ = self.main_tx.send(ToolMessage::Output(( - rid, - String::from("Invalid selection."), + format!("{} Selected!", server.address), ))); } + } else { + let _ = self.main_tx.send(ToolMessage::Output(( + rid, + String::from("Invalid selection."), + ))); } } else { let _ = self.main_tx.send(ToolMessage::Output(( @@ -514,46 +568,81 @@ impl AppState { } let _ = self.main_tx.send(ToolMessage::EndPrompt); } + "test_server" => { + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + server.message_que.push("TEST|NONE".to_string()); + } else { + self.output.push(String::from("error getting server lock!")); + } + } else { + self.output + .push(String::from("Error getting selected server!")); + } + } + "server_sleep" => { + if let Some(args) = command_args { + if let Ok(seconds) = args.trim().parse::() { + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + server.timer = Duration::from_secs(seconds); + self.output + .push(format!("Server timer set to {} seconds", seconds)); + } else { + self.output + .push(String::from("error locking server object!")); + } + } else { + self.output + .push(String::from("error not a valid server selected!")); + } + } else { + self.output + .push(String::from("invalid second count provided!")); + } + } else { + self.output.push(String::from( + "Error no seconds provided! please use server_sleep number_of_seconds", + )); + } + } "save_all" => { self.save_all()?; } "save_projects" => { self.save_projects()?; } - "test_server" => { - if let Some(server) = self.selected_server.clone() { - let _ = self - .main_tx - .send(ToolMessage::SendServer((server, "TEST".to_string()))); - } else { - let _ = self.main_tx.send(ToolMessage::Output(( - rid, - "You need to select a server first!".to_string(), - ))); - } - } "list_clients" => { - if let Some(server) = self.selected_server.clone() { - let _ = self.main_tx.send(ToolMessage::SendServer(( - server, - "LIST_CLIENTS".to_string(), - ))); + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + server.message_que.push("LIST_CLIENTS|NONE".to_string()); + } else { + self.output.push(String::from("error getting server lock!")); + } } else { - let _ = self.main_tx.send(ToolMessage::Output(( - rid, - "You need to select a server first!".to_string(), - ))); + self.output + .push(String::from("Error getting selected server!")); } } "select_client" => { if let Some(arg) = command_args { - if let Ok(id) = arg.parse::() { - self.selected_client = id; - } else { - let _ = self.main_tx.send(ToolMessage::Output(( - rid, - format!("error parsing client id from arg: {}", arg), - ))); + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + if let Ok(id) = arg.parse::() { + server.message_que.push(format!("STOP_CONTROL")); + server.message_que.push(format!("CONTROL|{}", id.clone())); + server.selected_client = id; + let _ = self.main_tx.send(ToolMessage::Output(( + 0, + format!("client {} selected!", id), + ))); + } else { + let _ = self.main_tx.send(ToolMessage::Output(( + rid, + format!("error parsing client id from arg: {}", arg), + ))); + } + } } } else { let _ = self.main_tx.send(ToolMessage::Output(( @@ -562,6 +651,17 @@ impl AppState { ))); } } + "deselect_client" => { + if let Some(server) = self.servers.get(self.selected_server) { + if let Ok(mut server) = server.lock() { + server.message_que.push(format!("STOP_CONTROL")); + server.selected_client = 0; + let _ = self + .main_tx + .send(ToolMessage::Output((0, format!("client deselected!")))); + } + } + } "add_scope" => { if let Some(args) = command_args { args.split_whitespace().into_iter().for_each(|host| { @@ -902,17 +1002,17 @@ impl AppState { ), connected: false, timer: Duration::from_secs(60), - id: 0, last_check: time::Instant::now(), message_que: Vec::new(), + action_que: Vec::new(), client_id: 0, + selected_client: 0, + config: None, }; - if let Ok(mut lock) = self.servers.lock() { - lock.push(new_server); - self.prompt.action = None; - self.prompt.responses.clear(); - self.prompt.num_responses = 0; - } + self.servers.push(Arc::new(Mutex::new(new_server))); + self.prompt.action = None; + self.prompt.responses.clear(); + self.prompt.num_responses = 0; return Ok(()); } else { match prompt.responses.len() { @@ -940,34 +1040,90 @@ pub struct Server { pub address: String, pub connected: bool, pub timer: Duration, - pub id: usize, + pub config: Option>, pub last_check: time::Instant, pub message_que: Vec, + pub action_que: Vec, pub client_id: usize, + pub selected_client: usize, } impl Server { - pub async fn checkin(&mut self, outputs: Vec) -> (String, String) { - let mut out_string = (String::new(), String::new()); - if let Ok(mut stream) = TcpStream::connect(self.address.clone()).await { - let payload = format!("{}:{}\n", self.id.clone(), outputs.join(":")); - if let Err(e) = stream.write_all(payload.as_bytes()).await { - out_string.0 = format!("Error sending response output to server! {e}:"); - } else { - out_string.0 = String::from("Success"); + pub fn donwload_certificate(&self) -> Result> { + let mut stream = TcpStream::connect(self.address.clone())?; + stream.write_all("CERT_REQ".as_bytes())?; + let mut buf = [0u8; 8192]; + let bytes_read = stream.read(&mut buf)?; + let response = String::from_utf8_lossy(&buf[..bytes_read]); + if let Some((_, cert)) = response.split_once("|") { + return Ok(cert.to_string()); + } + Err("Server did not send certificate".into()) + } + + pub fn connect(&mut self) -> Result> { + if let Ok(cert) = self.donwload_certificate() { + self.config = Some(self.build_tls_config(cert.as_str())?); + } + let mut stream = self.tls_connect()?; + stream.write("HELLO".as_bytes())?; + let mut buf = [0; 8192]; + let bytes_read = stream.read(&mut buf)?; + let response = String::from_utf8_lossy(&buf[..bytes_read]); + if let Some((_, id)) = response.split_once("|") { + self.client_id = id.trim().parse()?; + } + self.connected = true; + self.last_check = Instant::now(); + return Ok(self.client_id); + } + + pub fn checkin(&mut self) { + if let Ok(mut stream) = self.tls_connect() { + self.last_check = Instant::now(); + let payload = format!( + "{}|||{}\n", + self.client_id.clone(), + self.message_que.join("||") + ); + if let Err(e) = stream.write_all(payload.as_bytes()) { + self.action_que + .push(format!("Error|sending reponse output to server! {e}")); } + self.message_que.clear(); let mut buffer = [0; 4096]; - if let Ok(bytes_read) = stream.read(&mut buffer).await { + if let Ok(bytes_read) = stream.read(&mut buffer) { if bytes_read > 0 { let response = String::from_utf8_lossy(&buffer[..bytes_read]); - if response.starts_with("ACTIONS:") { - let actions_str = response.trim_start_matches("ACTIONS:").trim(); - out_string.1.push_str(actions_str); - } + response.split("||").into_iter().for_each(|action| { + self.action_que.push(action.trim().to_string()); + }); } } } - return out_string; + } + + pub fn build_tls_config(&self, pem: &str) -> Result, Box> { + let mut roots = RootCertStore::empty(); + let mut reader = BufReader::new(pem.as_bytes()); + let certs = rustls_pemfile::certs(&mut reader).collect::, _>>()?; + for cert in certs { + roots.add(cert)?; + } + let config = ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(); + Ok(Arc::new(config)) + } + + fn tls_connect(&self) -> Result, Box> { + let config = self.config.clone().ok_or("TLS config not loaded")?; + let tcp = TcpStream::connect(&self.address)?; + let (host, _) = self.address.split_once(':').ok_or("Invalid address")?; + let server_name = ServerName::try_from(host.to_string())?; + let connection = ClientConnection::new(config, server_name)?; + let tls_stream = StreamOwned::new(connection, tcp); + Ok(tls_stream) } } @@ -1314,14 +1470,6 @@ impl Host { } } -#[derive(Clone, Debug)] -pub enum Destination { - Server, - Victim, - Attacker, - Control, -} - #[derive(Clone, Debug)] pub struct User { pub name: String, @@ -1343,38 +1491,6 @@ impl User { } } -#[derive(Clone, Debug)] -pub enum ToolMessage { - Input((usize, String)), - Output((usize, String)), - UpdateHost(Host), - RebuildDB, - DestroyDB(DistroBox), - StopDB(DistroBox), - StopTemplate(DistroBox), - UpdateProject(usize, Project), - RemoveProject, - AddServer, - EndPrompt, - ConnectServer, - SelectServer, - SendServer((Server, String)), - ServerBrokerExit, - AppStateExit, - DisconnectServer, - DisconnectAllServers, -} - -#[derive(Clone, Debug)] -pub enum ToolArg { - Project(Project), - Projects(Vec), - Host(Host), - Hosts(Vec), - Config(HashMap), - Path(PathBuf), -} - #[derive(Clone, Debug)] pub struct ToolCommand { pub name: String, @@ -1549,12 +1665,6 @@ impl ModuleLoader { } } -enum AppEvent { - Key(event::KeyEvent), - Worker(ToolMessage), - Mouse(event::MouseEvent), -} - #[derive(Debug, Clone)] pub struct DistroBox { pub name: String, diff --git a/src/main.rs b/src/main.rs index 7c46270..326ac63 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,10 @@ use clap::Parser; -use std::sync::Arc; +use rcgen::generate_simple_self_signed; +use std::sync::{Arc, Mutex}; use std::{env, path::PathBuf, process::exit}; -use tetanus::AppState; use tetanus::funcs::*; -use tokio::sync::Mutex; +use tetanus::server::start_server; +use tetanus::{AppState, server}; mod install; @@ -29,21 +30,13 @@ struct Args { #[tokio::main] async fn main() { - let args = Args::parse(); - if args.server { - let server = Arc::new(Mutex::new(tetanus::server::ServerState { - id: 0, - running: true, - address: String::from("127.0.0.1:31337"), - clients: Vec::new(), - key: String::new(), - })); - tetanus::server::start_server(server).await.unwrap(); - } else { - println!("checking for server or client config files..."); - let mut config_path = env::home_dir().unwrap(); - config_path.push(".config/tetanus"); - if !config_path.exists() { + println!("checking for server or client config files..."); + let mut config_path = PathBuf::new(); + let mut config_path_opt = env::home_dir(); + if let Some(home_config_path) = config_path_opt.as_mut() { + println!("made it to home_config_path"); + home_config_path.push(".config/tetanus"); + if !home_config_path.exists() { println!("no config directory found in home directory..."); config_path = PathBuf::from("/etc/tetanus"); if !config_path.exists() { @@ -55,28 +48,55 @@ async fn main() { exit(0); } } + } else { + println!("found home config_path."); + config_path = home_config_path.clone(); } - let args = Args::parse(); - let mut client_config_path = config_path.clone(); - client_config_path.push("client.conf"); - let mut server_config_path = config_path.clone(); - server_config_path.push("server.conf"); - if args.client { - if !client_config_path.exists() { - eprintln!( - "error: no client config path found at {}", - client_config_path.display() - ); - exit(1); + } + let args = Args::parse(); + let mut client_config_path = config_path.clone(); + client_config_path.push("client.conf"); + let mut server_config_path = config_path.clone(); + server_config_path.push("server.conf"); + if args.client { + if !client_config_path.exists() { + eprintln!( + "error: no client config path found at {}", + client_config_path.display() + ); + exit(1); + } + let (mut appstate, rx) = AppState::new(); + let config_load = appstate.load_config(client_config_path.clone()); + if config_load.is_err() { + eprintln!("error loading config!"); + exit(1); + } + appstate.initialize_modules(); + println!("entering tui..."); + let _res = run_tui(appstate, rx).unwrap(); + } else if args.server { + let address = "127.0.0.1:31337".to_string(); + let names = vec!["127.0.0.1".to_string(), "localhost".to_string()]; + let mut certificate_path = config_path.clone(); + certificate_path.push("server"); + let mut key_path = certificate_path.clone(); + certificate_path.push("cert.pem"); + key_path.push("key.pem"); + println!("key paht: {}", key_path.display()); + println!("cert path: {}", certificate_path.display()); + let new_server = server::Server { + address, + clients: Vec::new(), + certificate_path, + key_path, + }; + let res = start_server(Arc::new(Mutex::new(new_server))).await; + match res { + Ok(_) => {} + Err(e) => { + println!("error running server {e}"); } - let (mut appstate, rx) = AppState::new(); - let config_load = appstate.load_config(client_config_path.clone()); - if config_load.is_err() { - eprintln!("error loading config!"); - exit(1); - } - appstate.initialize_modules(); - let _res = run_tui(appstate, rx).unwrap(); } } } diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..186fd06 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,317 @@ +use rcgen::generate_simple_self_signed; +use std::error::Error; +use std::fs::File; +use std::fs::write; +use std::io::BufReader; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::usize; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio_rustls::TlsAcceptor; +use tokio_rustls::rustls::ServerConfig; +use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; + +#[derive(Clone)] +enum ClientAction { + Output(String), + Cmd(String), + Ping, +} + +pub struct Client { + address: String, + id: usize, + hostname: Option, + actions: Vec, + controlling: usize, + controlled: usize, + output_que: Vec, +} + +pub struct Server { + pub address: String, + pub clients: Vec, + pub certificate_path: PathBuf, + pub key_path: PathBuf, +} + +pub async fn start_server(server: Arc>) -> Result<(), Box> { + let lock = server.lock().unwrap(); + if !lock.certificate_path.exists() || !lock.key_path.exists() { + let (server_ip, _) = lock.address.split_once(':').unwrap(); + + let cert = generate_simple_self_signed(vec![ + server_ip.to_string(), + "127.0.0.1".to_string(), + "localhost".to_string(), + ])?; + + write(&lock.certificate_path, cert.cert.pem())?; + + write(&lock.key_path, cert.signing_key.serialize_pem())?; + } + + let certs = load_certs(&lock.certificate_path)?; + let key = load_key(&lock.key_path)?; + + let tls_config = ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certs, key)?; + + let acceptor = TlsAcceptor::from(Arc::new(tls_config)); + + let listener = TcpListener::bind(lock.address.clone()).await?; + + println!("Listening on {}", lock.address); + + drop(lock); + + loop { + let (stream, addr) = listener.accept().await?; + + let acceptor = acceptor.clone(); + let server = server.clone(); + + tokio::spawn(async move { + let mut peek = [0u8; 8]; + if let Ok(n) = stream.peek(&mut peek).await { + if n >= 8 && &peek[..8] == b"CERT_REQ" { + handle_bootstrap(server, stream).await; + return; + } + } + match acceptor.accept(stream).await { + Ok(tls_stream) => { + handle_connection_stat9ic(server, tls_stream, addr).await; + } + + Err(e) => { + eprintln!("TLS handshake failed: {}", e); + } + } + }); + } +} + +fn load_certs(path: &std::path::Path) -> Result>, Box> { + let mut reader = BufReader::new(File::open(path)?); + + Ok(rustls_pemfile::certs(&mut reader).collect::, _>>()?) +} + +fn load_key(path: &std::path::Path) -> Result, Box> { + let mut reader = BufReader::new(File::open(path)?); + + let key = rustls_pemfile::private_key(&mut reader)?.ok_or("No private key found")?; + + Ok(key) +} + +async fn handle_bootstrap(server: Arc>, mut stream: TcpStream) { + let mut buf = [0u8; 1024]; + let n = match stream.read(&mut buf).await { + Ok(n) => n, + Err(_) => return, + }; + if n == 0 { + return; + } + let msg = String::from_utf8_lossy(&buf[..n]); + if msg.trim() != "CERT_REQ" { + return; + } + let cert_path = { + let lock = server.lock().unwrap(); + lock.certificate_path.clone() + }; + let cert_pem = match std::fs::read_to_string(cert_path) { + Ok(c) => c, + Err(_) => return, + }; + let response = format!("CERT|{}\n", cert_pem); + let _ = stream.write_all(response.as_bytes()).await; +} + +async fn handle_connection_stat9ic(server: Arc>, mut stream: S, addr: SocketAddr) +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let mut buf = vec![0u8; 4096]; + let n = match stream.read(&mut buf).await { + Ok(n) => n, + Err(e) => { + eprintln!("Read Error: {e}"); + return; + } + }; + if n == 0 { + return; + } + let msg = String::from_utf8_lossy(&buf[..n]); + if msg.contains("HELLO") { + let mut connected = false; + let mut id = 0; + if let Ok(mut lock) = server.lock() { + let new_client = Client { + address: addr.to_string(), + id: lock.clients.len() + 1, + hostname: None, + actions: Vec::new(), + controlled: 0, + controlling: 0, + output_que: Vec::new(), + }; + connected = true; + id = new_client.id.clone(); + lock.clients.push(new_client); + } + if connected { + stream + .write_all(format!("HELLO|{}\n", id).as_bytes()) + .await + .unwrap(); + println!("client connected! ID:{}", id,); + } + } else { + if let Some((source, data)) = msg.split_once("|||") { + if let Ok(source_id) = source.trim().parse::() { + let mut responses = Vec::new(); + if let Ok(mut lock) = server.lock() { + if let Some(source_client) = lock.clients.iter_mut().find(|c| c.id == source_id) + { + responses = source_client.actions.clone(); + } + data.split("||").into_iter().for_each(|action| { + if let Some((cmd, data)) = action.split_once("|") { + match cmd.trim() { + "OUTPUT" => { + if let Some(dest_client) = + lock.clients.iter_mut().find(|c| c.controlling == source_id) + { + dest_client.actions.push(ClientAction::Output(format!( + "from {}: {}", + dest_client.id, + data.trim() + ))); + } else { + if let Some(source_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + source_client.output_que.push(data.trim().to_string()); + } + } + } + "CMD" => { + if let Some(dest_client) = + lock.clients.iter_mut().find(|c| c.controlled == source_id) + { + dest_client + .actions + .push(ClientAction::Cmd(data.to_string())); + responses.push(ClientAction::Output(format!( + "tasked client {} to run {}", + dest_client.id, cmd + ))); + } + } + "CONTROL" => { + let new_control_id = data.trim().parse::().unwrap(); + if let Some(controlling_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + controlling_client.controlling = new_control_id.clone(); + } + if let Some(contlled_client) = + lock.clients.iter_mut().find(|c| c.id == new_control_id) + { + contlled_client.controlled = source_id; + } else { + if let Some(controlling_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + controlling_client.controlling = 0; + } + } + } + "STOP_CONTROL" => { + if let Some(source_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + if source_client.controlling != 0 { + let dest_client_id = source_client.controlling.clone(); + if let Some(dest_client) = lock + .clients + .iter_mut() + .find(|c| c.id == dest_client_id) + { + dest_client.controlled = 0; + } + } + } + } + "BREAK_CONTROL" => { + if let Some(source_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + if source_client.controlled != 0 { + let dest_id = source_client.controlled.clone(); + source_client.controlled = 0; + if let Some(dest_client) = + lock.clients.iter_mut().find(|c| c.id == dest_id) + { + dest_client.controlling = 0; + } + } + } + } + "LIST_CLIENTS" => { + lock.clients.iter().for_each(|c| { + let out = format!("{}: {}", c.id, c.address); + println!("client list requested!"); + println!("adding {} to response...", out); + responses.push(ClientAction::Output(out)); + }); + } + "TEST" => { + responses.push(ClientAction::Output("TEST BACK".to_string())); + } + _ => {} + } + } + }); + } + if responses.len() > 0 { + let mut messages = Vec::new(); + for r in responses { + match r { + ClientAction::Cmd(cmd) => { + messages.push(format!("CMD|{}", cmd)); + } + ClientAction::Output(text) => { + println!("adding {} to output", text); + messages.push(format!("OUTPUT|{}", text)); + } + ClientAction::Ping => { + messages.push(format!("PONG")); + } + } + } + let full_message = messages.join("||"); + println!("attempting to send {}", full_message); + let buf = full_message.as_bytes(); + stream.write(buf).await.unwrap(); + println!("buffer written!"); + if let Ok(mut lock) = server.lock() { + if let Some(source_client) = + lock.clients.iter_mut().find(|c| c.id == source_id) + { + source_client.actions.clear(); + } + } + } + } + } + } +}