diff --git a/src/cliparser.rs b/src/cliparser.rs index 0e9afe09a..6dd859c5a 100644 --- a/src/cliparser.rs +++ b/src/cliparser.rs @@ -35,9 +35,9 @@ pub struct CliArgs { #[arg(short = 'a', long)] pub search_all: bool, - /// Plot the entropy of the specified file + /// Write entropy plot image to the ENTROPY file #[arg(short = 'E', long, conflicts_with = "extract")] - pub entropy: bool, + pub entropy: Option, /// Log JSON results to a file ('-' for stdout) #[arg(short, long)] diff --git a/src/entropy.rs b/src/entropy.rs index 5c2c9a8e2..7b2e9bfb6 100644 --- a/src/entropy.rs +++ b/src/entropy.rs @@ -58,13 +58,23 @@ fn blocks(data: &[u8]) -> Vec { /// Generate a plot of a file's entropy. /// Will output a file to the current working directory with the name `.png`. -pub fn plot(file_path: impl Into, stdin: bool) -> Result { - const FILE_EXTENSION: &str = "png"; +pub fn plot( + png_file_path: impl Into, + file_path: impl Into, + stdin: bool, +) -> Result { + const FILE_EXTENSION: &str = ".png"; const SHANNON_MAX_VALUE: i32 = 8; const IMAGE_PIXEL_WIDTH: u32 = 2048; const IMAGE_PIXEL_HEIGHT: u32 = ((IMAGE_PIXEL_WIDTH as f64) * 0.6) as u32; let target_file: String = file_path.into(); + let mut png_path: String = png_file_path.into(); + + // Make sure the output file extension is .png + if !png_path.ends_with(FILE_EXTENSION) { + png_path = format!("{}{}", png_path, FILE_EXTENSION); + } // Get the base name of the target file let target_file_name = path::Path::new(&target_file) @@ -74,12 +84,10 @@ pub fn plot(file_path: impl Into, stdin: bool) -> Result