92 lines
2.7 KiB
Rust
92 lines
2.7 KiB
Rust
use clap::Parser;
|
|
use rand::seq::SliceRandom;
|
|
use tokio::net::TcpListener;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tracing::{debug, error, info, Level};
|
|
use tracing_subscriber::FmtSubscriber;
|
|
|
|
mod cli;
|
|
mod handler;
|
|
|
|
use cli::Cli;
|
|
use handler::{generate_payload, parse_signatures};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
// Parse CLI
|
|
let cli = Cli::parse();
|
|
|
|
// Setup logger
|
|
let subscriber = FmtSubscriber::builder()
|
|
.with_max_level(if cli.debug {
|
|
Level::DEBUG
|
|
} else if cli.verbose {
|
|
Level::INFO
|
|
} else {
|
|
Level::ERROR
|
|
})
|
|
.without_time()
|
|
.finish();
|
|
|
|
tracing::subscriber::set_global_default(subscriber)
|
|
.expect("setting default subscriber failed");
|
|
|
|
debug!("Parsed CLI flags");
|
|
|
|
// 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());
|
|
|
|
let listener = TcpListener::bind(&cli.listen).await?;
|
|
info!("Started listener on {}", cli.listen);
|
|
|
|
loop {
|
|
// Accept connection
|
|
let (mut stream, address) = listener.accept().await?;
|
|
debug!("Accepted connection from {}", address);
|
|
|
|
let sigs = signatures.clone();
|
|
let cli_clone = cli.clone();
|
|
|
|
// Spawn async thread
|
|
tokio::spawn(async move {
|
|
// Choose random signature
|
|
let signature = sigs.choose(&mut rand::thread_rng());
|
|
|
|
if let Some(sig) = signature {
|
|
let payload = generate_payload(sig);
|
|
|
|
// Write payload
|
|
match stream.write_all(&payload).await {
|
|
Ok(()) => {
|
|
debug!(
|
|
"Sent payload to {}: {:?} ({} bytes)",
|
|
address,
|
|
String::from_utf8_lossy(&payload),
|
|
payload.len()
|
|
);
|
|
if cli_clone.verbose {
|
|
info!("Sent payload ({} bytes) to {}", payload.len(), address);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
if e.kind() == std::io::ErrorKind::ConnectionReset {
|
|
debug!("Connection reset by peer: {}", address);
|
|
} else {
|
|
error!("Failed to write payload to {}: {}", address, e);
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
debug!("No signature available");
|
|
}
|
|
});
|
|
}
|
|
}
|