implement regex parsing

This commit is contained in:
sad 2024-10-04 17:06:29 +00:00
parent fecd1cc899
commit f1e1d9ec90
5 changed files with 278 additions and 46 deletions

43
Cargo.lock generated
View File

@ -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"

View File

@ -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"

View File

@ -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
View 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"
);
}
}

View File

@ -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 {
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;
}
//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,
};
});
}
}