implement regex parsing

Cargo.lock

@@ -17,6 +17,15 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "aho-corasick"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "anstream"
 version = "0.3.2"
@@ -209,6 +218,7 @@ dependencies = [
  "anyhow",
  "clap",
  "rand",
+ "regex",
  "tokio",
  "tracing",
  "tracing-subscriber",
@@ -279,9 +289,9 @@ checksum = "b06a4cde4c0f271a446782e3eff8de789548ce57dbc8eca9292c27f4a42004b4"
 
 [[package]]
 name = "memchr"
-version = "2.5.0"
+version = "2.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
+checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
 
 [[package]]
 name = "miniz_oxide"
@@ -436,6 +446,35 @@ dependencies = [
  "bitflags 1.3.2",
 ]
 
+[[package]]
+name = "regex"
+version = "1.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.23"

Cargo.toml

@@ -9,6 +9,7 @@ edition = "2021"
 anyhow = "1.0.72"
 clap = { version = "4.3.19", features = ["derive"] }
 rand = "0.8.5"
+regex = "1.11.0"
 tokio = { version = "1", features = ["full"] }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.17"

src/cli.rs

@@ -22,31 +22,15 @@ pub struct Cli {
     )]
     pub listen: String,
 
-    #[arg(
+    #[arg(short = 'd', long = "debug", help = "Enable debug logging")]
-        short = 'd',
-        long = "debug",
-        help = "Enable debug logging"
-    )]
     pub debug: bool,
 
-    #[arg(
+    #[arg(short = 'v', long = "verbose", help = "Enable verbose logging")]
-        short = 'v',
-        long = "verbose",
-        help = "Enable verbose logging"
-    )]
     pub verbose: bool,
 
-    #[arg(
+    #[arg(short = 'q', long = "quiet", help = "Enable quiet logging")]
-        short = 'q',
-        long = "quiet",
-        help = "Enable quiet logging"
-    )]
     pub quiet: bool,
 
-    #[arg(
+    #[arg(short = 'V', long = "version", help = "Print version information")]
-        short = 'V',
-        long = "version",
-        help = "Print version information"
-    )]
     pub version: bool,
 }

src/handler.rs

@@ -0,0 +1,195 @@
+use anyhow::{anyhow, Context, Result};
+use rand::Rng;
+use regex::Regex;
+use std::fs::File;
+use std::io::{BufRead, BufReader};
+
+#[derive(Debug, Clone)]
+pub enum Payload {
+    Raw(Vec<u8>),
+    Regex(String),
+}
+
+#[derive(Debug, Clone)]
+pub struct Signature {
+    pub payload: Payload,
+}
+
+pub fn parse_signatures(file_path: &str) -> Result<Vec<Signature>> {
+    let file = File::open(file_path).context("Failed to open signatures file")?;
+    let reader = BufReader::new(file);
+    let mut signatures = Vec::new();
+
+    for (index, line) in reader.lines().enumerate() {
+        let line = line.context("Failed to read line from signatures file")?;
+        if line.trim().is_empty() {
+            continue; // Skip empty lines
+        }
+
+        let payload = if line.contains('(') && line.contains(')') {
+            Payload::Regex(line)
+        } else {
+            Payload::Raw(
+                unescape_string(&line)
+                    .with_context(|| format!("Invalid payload on line {}", index + 1))?,
+            )
+        };
+
+        signatures.push(Signature { payload });
+    }
+
+    if signatures.is_empty() {
+        return Err(anyhow!("No valid signatures found in the file"));
+    }
+
+    Ok(signatures)
+}
+
+fn unescape_string(s: &str) -> Result<Vec<u8>> {
+    let mut result = Vec::new();
+    let mut chars = s.chars();
+
+    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');
+                            return String::new();
+                        });
+                    if !hex.is_empty() {
+                        if let Ok(byte) = u8::from_str_radix(&hex, 16) {
+                            result.push(byte);
+                        } else {
+                            result.push(b'\\');
+                            result.push(b'x');
+                            result.extend(hex.bytes());
+                        }
+                    }
+                }
+                Some('0') => result.push(0),
+                Some('n') => result.push(b'\n'),
+                Some('r') => result.push(b'\r'),
+                Some('t') => result.push(b'\t'),
+                Some(c) => result.push(c as u8),
+                None => result.push(b'\\'),
+            }
+        } else {
+            result.push(c as u8);
+        }
+    }
+
+    Ok(result)
+}
+
+pub fn generate_payload(signature: &Signature) -> Vec<u8> {
+    match &signature.payload {
+        Payload::Raw(v) => v.clone(),
+        Payload::Regex(r) => generate_regex_match(r),
+    }
+}
+
+fn generate_regex_match(regex_str: &str) -> Vec<u8> {
+    // Simplified regex matching that doesn't rely on the regex crate
+    let mut result = String::new();
+    let mut chars = regex_str.chars().peekable();
+
+    while let Some(c) = chars.next() {
+        match c {
+            '\\' => {
+                if let Some(next_char) = chars.next() {
+                    match next_char {
+                        'd' => result.push(rand::thread_rng().gen_range(b'0'..=b'9') as char),
+                        'w' => result.push(rand::thread_rng().gen_range(b'a'..=b'z') as char),
+                        'x' => {
+                            // Handle \x hex escapes
+                            let hex = chars
+                                .next()
+                                .and_then(|c1| chars.next().map(|c2| format!("{}{}", c1, c2)))
+                                .unwrap_or_else(|| "00".to_string());
+                            if let Ok(byte) = u8::from_str_radix(&hex, 16) {
+                                result.push(byte as char);
+                            }
+                        }
+                        _ => result.push(next_char),
+                    }
+                }
+            }
+            '[' => {
+                let mut class = String::new();
+                while let Some(class_char) = chars.next() {
+                    if class_char == ']' {
+                        break;
+                    }
+                    class.push(class_char);
+                }
+                if !class.is_empty() {
+                    result.push(
+                        class
+                            .chars()
+                            .nth(rand::thread_rng().gen_range(0..class.len()))
+                            .unwrap(),
+                    );
+                }
+            }
+            '(' => {
+                // Skip capturing groups
+                let mut depth = 1;
+                while let Some(group_char) = chars.next() {
+                    if group_char == '(' {
+                        depth += 1;
+                    }
+                    if group_char == ')' {
+                        depth -= 1;
+                    }
+                    if depth == 0 {
+                        break;
+                    }
+                }
+            }
+            '+' | '*' => {
+                if let Some(last_char) = result.chars().last() {
+                    let repeat = rand::thread_rng().gen_range(0..5);
+                    for _ in 0..repeat {
+                        result.push(last_char);
+                    }
+                }
+            }
+            '.' => result.push(rand::thread_rng().gen_range(b'!'..=b'~') as char),
+            _ => result.push(c),
+        }
+    }
+
+    result.into_bytes()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_generate_regex_match() {
+        let regex_str = r"Hello [\w]+, your lucky number is \d+";
+        let result = generate_regex_match(regex_str);
+        let result_str = String::from_utf8_lossy(&result);
+        assert!(result_str.starts_with("Hello "));
+        assert!(result_str.contains(", your lucky number is "));
+    }
+
+    #[test]
+    fn test_unescape_string() {
+        assert_eq!(unescape_string(r"Hello\nWorld").unwrap(), b"Hello\nWorld");
+        assert_eq!(unescape_string(r"Test\x41\x42\x43").unwrap(), b"TestABC");
+        assert_eq!(unescape_string(r"\0\r\n\t").unwrap(), b"\0\r\n\t");
+        assert_eq!(unescape_string(r"Incomplete\").unwrap(), b"Incomplete\\");
+        assert_eq!(unescape_string(r"Incomplete\x").unwrap(), b"Incomplete\\x");
+        assert_eq!(
+            unescape_string(r"Incomplete\x4").unwrap(),
+            b"Incomplete\\x4"
+        );
+    }
+}

src/main.rs

@@ -1,12 +1,14 @@
 use clap::Parser;
 use rand::seq::SliceRandom;
 use tokio::net::TcpListener;
-use tracing::{debug, info, Level};
+use tracing::{debug, error, info, Level};
-
-use cli::Cli;
 
 mod cli;
 mod config;
+mod handler;
+
+use cli::Cli;
+use handler::{generate_payload, parse_signatures, Signature};
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
@@ -16,17 +18,23 @@ async fn main() -> anyhow::Result<()> {
         .with_max_level(Level::DEBUG)
         .init();
 
-
     // Parse CLI
     let cli = Cli::parse();
     debug!("Parsed CLI flags");
+
     // Read signatures file
-    let signatures = config::read_signatures(&cli.signatures)?;
+    let signatures = match parse_signatures(&cli.signatures) {
-    debug!("Read signatures file");
+        Ok(sigs) => sigs,
+        Err(e) => {
+            error!("Failed to parse signatures file: {}", e);
+            return Err(e);
+        }
+    };
+    debug!("Read {} signatures", signatures.len());
 
     // Bind listener
     let listener = TcpListener::bind(&cli.listen).await?;
-    info!("Started listener");
+    info!("Started listener on {}", cli.listen);
 
     loop {
         // Accept connection
@@ -35,10 +43,8 @@ async fn main() -> anyhow::Result<()> {
             debug!("Accepted connection from {}", address);
         } else if cli.verbose {
             info!("Accepted connection from {}", address);
-        } else if cli.quiet {
-
         }
-        //debug!("Accepted connection");
+
         // Clone signatures
         let sigs = signatures.clone();
 
@@ -47,21 +53,28 @@ async fn main() -> anyhow::Result<()> {
             // Choose random signature
             let signature = sigs.choose(&mut rand::thread_rng());
 
-            // Write signature
+            if let Some(sig) = signature {
-            match stream.try_write(signature.expect("could not send signature").as_bytes()) {
+                // Generate payload
-                Ok(n) => {
+                let payload = generate_payload(sig);
-                    if cli.debug {
+
-                        debug!("Sent signature {:?} to {}", signature, address);
+                // Write payload
-                    } else if cli.verbose {
+                if let Err(e) = stream.try_write(&payload) {
-                        info!("Sent signature {:?} to {}", signature, address);
+                    error!("Failed to write payload to {}: {}", address, e);
-                    } else if cli.quiet {
                     return;
                 }
-                    //debug!("Sent signature {:?} to {}", signature, address);
+
-                    n
+                if cli.debug {
+                    debug!(
+                        "Sent payload to {}: {:?}",
+                        address,
+                        String::from_utf8_lossy(&payload)
+                    );
+                } else if cli.verbose {
+                    info!("Sent payload to {}", address);
+                }
+            } else {
+                debug!("No signature available");
             }
-                Err(_) => return,
-            };
         });
     }
 }