From f1e1d9ec905e7cc7ee7c7909e1ce144353613574 Mon Sep 17 00:00:00 2001 From: sad Date: Fri, 4 Oct 2024 17:06:29 +0000 Subject: [PATCH] implement regex parsing --- Cargo.lock | 43 ++++++++++- Cargo.toml | 1 + src/cli.rs | 24 +----- src/handler.rs | 195 +++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 61 ++++++++++------ 5 files changed, 278 insertions(+), 46 deletions(-) create mode 100644 src/handler.rs diff --git a/Cargo.lock b/Cargo.lock index cc93fb5..0889b13 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index a8da4c7..1d8ae4e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/cli.rs b/src/cli.rs index de0badf..9d1ac3a 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -22,31 +22,15 @@ pub struct Cli { )] pub listen: String, - #[arg( - short = 'd', - long = "debug", - help = "Enable debug logging" - )] + #[arg(short = 'd', long = "debug", help = "Enable debug logging")] pub debug: bool, - #[arg( - short = 'v', - long = "verbose", - help = "Enable verbose logging" - )] + #[arg(short = 'v', long = "verbose", help = "Enable verbose logging")] pub verbose: bool, - #[arg( - short = 'q', - long = "quiet", - help = "Enable quiet logging" - )] + #[arg(short = 'q', long = "quiet", help = "Enable quiet logging")] pub quiet: bool, - #[arg( - short = 'V', - long = "version", - help = "Print version information" - )] + #[arg(short = 'V', long = "version", help = "Print version information")] pub version: bool, } diff --git a/src/handler.rs b/src/handler.rs new file mode 100644 index 0000000..9135e7c --- /dev/null +++ b/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), + Regex(String), +} + +#[derive(Debug, Clone)] +pub struct Signature { + pub payload: Payload, +} + +pub fn parse_signatures(file_path: &str) -> Result> { + 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> { + 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 { + match &signature.payload { + Payload::Raw(v) => v.clone(), + Payload::Regex(r) => generate_regex_match(r), + } +} + +fn generate_regex_match(regex_str: &str) -> Vec { + // 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" + ); + } +} diff --git a/src/main.rs b/src/main.rs index 84db14b..bd6c508 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,14 @@ use clap::Parser; use rand::seq::SliceRandom; use tokio::net::TcpListener; -use tracing::{debug, info, Level}; - -use cli::Cli; +use tracing::{debug, error, info, Level}; 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)?; - debug!("Read signatures file"); + let signatures = match parse_signatures(&cli.signatures) { + 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 - match stream.try_write(signature.expect("could not send signature").as_bytes()) { - Ok(n) => { - if cli.debug { - debug!("Sent signature {:?} to {}", signature, address); - } else if cli.verbose { - info!("Sent signature {:?} to {}", signature, address); - } else if cli.quiet { - return; - } - //debug!("Sent signature {:?} to {}", signature, address); - n + 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; } - Err(_) => return, - }; + + 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"); + } }); } }