From 1f29c0e62c1c5aeb1d3572ff9a8848e6ebda7050 Mon Sep 17 00:00:00 2001
From: sad <sad@im.gay>
Date: Fri, 4 Oct 2024 17:27:08 +0000
Subject: [PATCH] cleanup, fix fmt and lint, fix tests, fix signature checking

---
 src/cli.rs     |  5 +++--
 src/handler.rs | 18 +++++++-----------
 src/main.rs    | 42 ++++++++++++++++++++++++------------------
 3 files changed, 34 insertions(+), 31 deletions(-)

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<Vec<Signature>> {
 
 fn unescape_string(s: &str) -> Result<Vec<u8>> {
     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");