Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds ability for payload tool to emit the raw contents of blocks #1031

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Adds ability for payload tool to emit the raw contents of blocks #1031

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 34 additions & 17 deletions lading/src/bin/payloadtool.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
use std::io::Read;
use std::io::Write;
use std::time::Instant;
use std::{io, num::NonZeroU32};

use clap::Parser;
use lading::generator::http::Method;
use lading_payload::block;
use lading_payload::block::{self, Block};
use rand::{rngs::StdRng, SeedableRng};
use tracing::{debug, error, info, warn};
use tracing::{error, info, warn};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::{fmt::format::FmtSpan, util::SubscriberInitExt};

const UDP_PACKET_LIMIT_BYTES: u32 = 65_507;
Expand All @@ -17,6 +19,10 @@ struct Args {
/// Path to standard lading config file
config_path: String,

/// Emit blocks line by line on stdout
#[clap(short, long)]
emit_to_stdout: bool,

/// Optionally only run a single generator's payload
#[clap(short, long)]
generator_id: Option<String>,
Expand All @@ -41,21 +47,20 @@ fn generate_and_check(
seed: [u8; 32],
total_bytes: NonZeroU32,
max_block_size: byte_unit::Byte,
) -> Result<(), Error> {
) -> Result<Vec<Block>, Error> {
let mut rng = StdRng::from_seed(seed);
let start = Instant::now();
let blocks =
match block::Cache::fixed(&mut rng, total_bytes, max_block_size.get_bytes(), config)? {
block::Cache::Fixed { blocks, idx: _ } => blocks,
};
info!("Payload generation took {:?}", start.elapsed());
debug!("Payload: {:#?}", blocks);

let mut total_generated_bytes: u32 = 0;
for block in blocks.iter() {
total_generated_bytes += block.total_bytes.get();
}
if total_bytes.get() != total_generated_bytes {
if total_bytes.get() > total_generated_bytes {
let total_requested_bytes = byte_unit::Byte::from_bytes(total_bytes.get().into());
let total_requested_bytes_str = total_requested_bytes
.get_appropriate_unit(false)
Expand All @@ -67,11 +72,11 @@ fn generate_and_check(
warn!("Generator failed to generate {total_requested_bytes_str}, instead only found {total_generated_bytes_str} of data")
}

Ok(())
Ok(blocks)
}

fn check_generator(config: &lading::generator::Config) -> Result<(), Error> {
match &config.inner {
fn check_generator(config: &lading::generator::Config, emit_to_stdout: bool) -> Result<(), Error> {
let blocks = match &config.inner {
lading::generator::Inner::FileGen(_) => unimplemented!("FileGen not supported"),
lading::generator::Inner::UnixDatagram(g) => {
let max_block_size =
Expand All @@ -80,13 +85,13 @@ fn check_generator(config: &lading::generator::Config) -> Result<(), Error> {
let total_bytes =
NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(&g.variant, g.seed, total_bytes, max_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, max_block_size)?
}
lading::generator::Inner::Tcp(g) => {
let total_bytes =
NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?
}
lading::generator::Inner::Udp(g) => {
let total_bytes =
Expand All @@ -95,7 +100,7 @@ fn check_generator(config: &lading::generator::Config) -> Result<(), Error> {
let max_block_size =
byte_unit::Byte::from_unit(UDP_PACKET_LIMIT_BYTES.into(), byte_unit::ByteUnit::B)
.expect("valid bytes");
generate_and_check(&g.variant, g.seed, total_bytes, max_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, max_block_size)?
}
lading::generator::Inner::Http(g) => {
let (variant, max_prebuild_cache_size_bytes) = match &g.method {
Expand All @@ -107,32 +112,43 @@ fn check_generator(config: &lading::generator::Config) -> Result<(), Error> {
};
let total_bytes = NonZeroU32::new(max_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(variant, g.seed, total_bytes, g.maximum_block_size)?;
generate_and_check(variant, g.seed, total_bytes, g.maximum_block_size)?
}
lading::generator::Inner::SplunkHec(_) => unimplemented!("SplunkHec not supported"),
lading::generator::Inner::FileTree(_) => unimplemented!("FileTree not supported"),
lading::generator::Inner::Grpc(g) => {
let total_bytes =
NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?
}
lading::generator::Inner::UnixStream(g) => {
let total_bytes =
NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?
}
lading::generator::Inner::PassthruFile(g) => {
let total_bytes =
NonZeroU32::new(g.maximum_prebuild_cache_size_bytes.get_bytes() as u32)
.expect("Non-zero max prebuild cache size");
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?;
generate_and_check(&g.variant, g.seed, total_bytes, g.maximum_block_size)?
}
lading::generator::Inner::ProcessTree(_) => unimplemented!("ProcessTree not supported"),
lading::generator::Inner::ProcFs(_) => unimplemented!("ProcFs not supported"),
};

if emit_to_stdout {
for block in blocks {
match std::io::stdout().write_all(&block.bytes) {
Ok(_) => {}
Err(e) => {
error!("Failed to write block to stdout: {}", e);
}
}
}
}

Ok(())
}

Expand All @@ -141,6 +157,7 @@ async fn main() -> Result<(), Error> {
tracing_subscriber::fmt()
.with_span_events(FmtSpan::CLOSE)
.with_ansi(false)
.with_env_filter(EnvFilter::from_default_env())
.finish()
.init();

Expand Down Expand Up @@ -180,10 +197,10 @@ async fn main() -> Result<(), Error> {
error!("No generator found with id: {}", generator_id);
Error::InvalidArgs
})?;
check_generator(generator)?;
check_generator(generator, args.emit_to_stdout)?;
} else {
for generator in config.generator {
check_generator(&generator)?;
check_generator(&generator, args.emit_to_stdout)?;
}
}

Expand Down
Loading