implement regex parsing
This commit is contained in:
parent
fecd1cc899
commit
f1e1d9ec90
43
Cargo.lock
generated
43
Cargo.lock
generated
@ -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"
|
||||
|
@ -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"
|
||||
|
24
src/cli.rs
24
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,
|
||||
}
|
||||
|
195
src/handler.rs
Normal file
195
src/handler.rs
Normal file
@ -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"
|
||||
);
|
||||
}
|
||||
}
|
61
src/main.rs
61
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");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user