diff --git a/src/cli.rs b/src/cli.rs index 9d1ac3a..e44b354 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,7 +1,8 @@ use clap::Parser; -#[derive(Parser, Debug)] -#[command(author = "perp")] +#[derive(Parser, Debug, Clone)] +#[command(author = "perp and sad")] + // CLI flags/arguments pub struct Cli { #[arg( diff --git a/src/handler.rs b/src/handler.rs index ee83edb..afadb4e 100644 --- a/src/handler.rs +++ b/src/handler.rs @@ -44,27 +44,23 @@ pub fn parse_signatures(file_path: &str) -> Result> { fn unescape_string(s: &str) -> Result> { let mut result = Vec::new(); - let mut chars = s.chars(); + let mut chars = s.chars().peekable(); while let Some(c) = chars.next() { if c == '\\' { match chars.next() { Some('x') => { - let hex = chars.next().and_then(|c1| { - chars.next().map(|c2| format!("{}{}", c1, c2)) - }).unwrap_or_else(|| { - result.push(b'\\'); - result.push(b'x'); - String::new() - }); - if !hex.is_empty() { + let hex: String = chars.by_ref().take(2).collect(); + if hex.len() == 2 { if let Ok(byte) = u8::from_str_radix(&hex, 16) { result.push(byte); } else { - result.push(b'\\'); - result.push(b'x'); + result.extend(b"\\x"); result.extend(hex.bytes()); } + } else { + result.extend(b"\\x"); + result.extend(hex.bytes()); } }, Some('0') => result.push(0), diff --git a/src/main.rs b/src/main.rs index f45138d..e6f5f77 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use clap::Parser; use rand::seq::SliceRandom; use tokio::net::TcpListener; +use tokio::io::AsyncWriteExt; use tracing::{debug, error, info, Level}; mod cli; @@ -29,23 +30,22 @@ async fn main() -> anyhow::Result<()> { return Err(e); } }; - debug!("Read {} signatures", signatures.len()); + //debug!("Read {} signatures", signatures.len()); - // Bind listener let listener = TcpListener::bind(&cli.listen).await?; info!("Started listener on {}", cli.listen); loop { // Accept connection - let (stream, address) = listener.accept().await?; + let (mut stream, address) = listener.accept().await?; if cli.debug { debug!("Accepted connection from {}", address); } else if cli.verbose { info!("Accepted connection from {}", address); } - // Clone signatures let sigs = signatures.clone(); + let cli_clone = cli.clone(); // Spawn async thread tokio::spawn(async move { @@ -53,23 +53,29 @@ async fn main() -> anyhow::Result<()> { let signature = sigs.choose(&mut rand::thread_rng()); if let Some(sig) = signature { - // Generate payload let payload = generate_payload(sig); // Write payload - if let Err(e) = stream.try_write(&payload) { - error!("Failed to write payload to {}: {}", address, e); - return; - } - - if cli.debug { - debug!( - "Sent payload to {}: {:?}", - address, - String::from_utf8_lossy(&payload) - ); - } else if cli.verbose { - info!("Sent payload to {}", address); + match stream.write_all(&payload).await { + Ok(()) => { + if cli_clone.debug { + debug!( + "Sent payload to {}: {:?} ({} bytes)", + address, + String::from_utf8_lossy(&payload), + payload.len() + ); + } else if cli_clone.verbose { + info!("Sent payload ({} bytes) to {}", payload.len(), address); + } + } + Err(e) => { + if e.kind() == std::io::ErrorKind::ConnectionReset { + debug!("Connection reset by peer: {}", address); + } else { + error!("Failed to write payload to {}: {}", address, e); + } + } } } else { debug!("No signature available");