diff --git a/src/lib.rs b/src/lib.rs index 2760aee..8958c99 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -237,10 +237,10 @@ impl AppState { self.help .push("test_server\ntest the server connection\n".to_string()); self.help.push( - "add_scope\nadd a host, cidr range, or ip address to the current project's scope" + "add_scope\nadd a host, cidr range, or ip address to the current project's scope\n" .to_string(), ); - self.help.push("parse_scope\nparse the projects general.md notes file for the scope copied from the workbook".to_string()); + self.help.push("parse_scope\nparse the projects general.md notes file for the scope copied from the workbook\n".to_string()); self.help.push("exit\nquit the tool\n".to_string()); return Ok(()); } @@ -394,6 +394,18 @@ impl AppState { let _ = self.main_tx.send(ToolMessage::Output((rid, line))); } } + "rebuild_distro_box" | "rdb" => { + if let Some(db) = self.projects[self.selected_project].db.clone() { + if db.created { + let _ = self.main_tx.send(ToolMessage::DestroyDB(db.clone())); + } + let _ = self.main_tx.send(ToolMessage::RebuildDB); + let _ = self.main_tx.send(ToolMessage::Output(( + rid, + "[success] Distorbox made!".to_string(), + ))); + } + } "promote_project" | "pp" => { self.main_tx.send(ToolMessage::Output(( rid, @@ -1230,7 +1242,7 @@ impl Server { self.cert_text = cert; } let mut stream = self.tls_connect()?; - stream.write("HELLO".as_bytes())?; + stream.write("HELLO|attack".as_bytes())?; let mut buf = [0; 8192]; let bytes_read = stream.read(&mut buf)?; let response = String::from_utf8_lossy(&buf[..bytes_read]); diff --git a/src/server.rs b/src/server.rs index e9eb48e..e1a0293 100644 --- a/src/server.rs +++ b/src/server.rs @@ -37,6 +37,7 @@ pub struct Client { last_check: Instant, remove: bool, name: String, + victim: bool, } pub struct Server { @@ -184,6 +185,12 @@ where } let msg = String::from_utf8_lossy(&buf[..n]); if msg.contains("HELLO") { + let mut victim = false; + if let Some((_, ctype)) = msg.split_once("|") { + if ctype.contains("victim") { + victim = true; + } + } let mut connected = false; let mut id = 0; if let Ok(mut lock) = server.lock() { @@ -200,6 +207,7 @@ where last_check: Instant::now(), remove: false, name: String::new(), + victim, }; connected = true; id = new_client.id.clone(); diff --git a/victim_templates/rust_victim b/victim_templates/rust_victim index 12bc4dc..3f7b1bf 100644 --- a/victim_templates/rust_victim +++ b/victim_templates/rust_victim @@ -36,7 +36,7 @@ pub struct Server { impl Server { pub fn connect(&mut self) -> Result> { let mut stream = self.tls_connect()?; - stream.write("HELLO".as_bytes())?; + stream.write("HELLO|victim".as_bytes())?; let mut buf = [0; 8192]; let bytes_read = stream.read(&mut buf)?; let response = String::from_utf8_lossy(&buf[..bytes_read]);