diff --git a/src/config.rs b/src/config.rs deleted file mode 100644 index 6be5365..0000000 --- a/src/config.rs +++ /dev/null @@ -1,10 +0,0 @@ -use anyhow::{Context, Result}; -use std::fs; - -// Read signatures file -pub fn read_signatures(path: &str) -> Result> { - // Read path - let content = fs::read_to_string(path) - .with_context(|| format!("Could not read signatures file {}", path))?; - Ok(content.lines().map(String::from).collect()) -} diff --git a/src/handler.rs b/src/handler.rs index 9135e7c..ee83edb 100644 --- a/src/handler.rs +++ b/src/handler.rs @@ -1,8 +1,7 @@ -use anyhow::{anyhow, Context, Result}; -use rand::Rng; -use regex::Regex; use std::fs::File; use std::io::{BufRead, BufReader}; +use rand::Rng; +use anyhow::{Context, Result, anyhow}; #[derive(Debug, Clone)] pub enum Payload { @@ -23,16 +22,14 @@ pub fn parse_signatures(file_path: &str) -> Result> { 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 + continue; } 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))?, - ) + Payload::Raw(unescape_string(&line) + .with_context(|| format!("Invalid payload on line {}", index + 1))?) }; signatures.push(Signature { payload }); @@ -53,14 +50,13 @@ fn unescape_string(s: &str) -> Result> { 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(); - }); + 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() { if let Ok(byte) = u8::from_str_radix(&hex, 16) { result.push(byte); @@ -70,7 +66,7 @@ fn unescape_string(s: &str) -> Result> { result.extend(hex.bytes()); } } - } + }, Some('0') => result.push(0), Some('n') => result.push(b'\n'), Some('r') => result.push(b'\r'), @@ -94,7 +90,6 @@ pub fn generate_payload(signature: &Signature) -> Vec { } 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(); @@ -106,10 +101,7 @@ fn generate_regex_match(regex_str: &str) -> Vec { '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))) + 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); @@ -118,39 +110,25 @@ fn generate_regex_match(regex_str: &str) -> Vec { _ => result.push(next_char), } } - } + }, '[' => { let mut class = String::new(); - while let Some(class_char) = chars.next() { - if class_char == ']' { - break; - } + for class_char in chars.by_ref() { + 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(), - ); + 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; - } + for group_char in chars.by_ref() { + 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); @@ -158,7 +136,7 @@ fn generate_regex_match(regex_str: &str) -> Vec { result.push(last_char); } } - } + }, '.' => result.push(rand::thread_rng().gen_range(b'!'..=b'~') as char), _ => result.push(c), } @@ -187,9 +165,6 @@ mod tests { 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" - ); + assert_eq!(unescape_string(r"Incomplete\x4").unwrap(), b"Incomplete\\x4"); } } diff --git a/src/main.rs b/src/main.rs index bd6c508..f45138d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,11 +4,10 @@ use tokio::net::TcpListener; use tracing::{debug, error, info, Level}; mod cli; -mod config; mod handler; use cli::Cli; -use handler::{generate_payload, parse_signatures, Signature}; +use handler::{generate_payload, parse_signatures}; #[tokio::main] async fn main() -> anyhow::Result<()> {