Merge pull request ReFirmLabs#794 from ReFirmLabs/entropy_output_file

Added ability to specify an output file name for entropy plots

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<String>,
 
     /// Log JSON results to a file ('-' for stdout)
     #[arg(short, long)]

src/entropy.rs

@@ -58,13 +58,23 @@ fn blocks(data: &[u8]) -> Vec<BlockEntropy> {
 
 /// Generate a plot of a file's entropy.
 /// Will output a file to the current working directory with the name `<file_name>.png`.
-pub fn plot(file_path: impl Into<String>, stdin: bool) -> Result<FileEntropy, EntropyError> {
-    const FILE_EXTENSION: &str = "png";
+pub fn plot(
+    png_file_path: impl Into<String>,
+    file_path: impl Into<String>,
+    stdin: bool,
+) -> Result<FileEntropy, EntropyError> {
+    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<String>, stdin: bool) -> Result<FileEntropy, En
         .unwrap();
 
     let mut file_entropy = FileEntropy {
-        file: format!("{}.{}", target_file_name, FILE_EXTENSION),
+        file: png_path.clone(),
         ..Default::default()
     };
 
-    let png_path = file_entropy.file.clone();
-
     // Make sure the output file doesn't already exist
     if path::Path::new(&png_path).exists() {
         error!("Cannot create entropy graph {}: File exists", png_path);

src/main.rs

@@ -74,10 +74,14 @@ fn main() {
     let mut json_logger = json::JsonLogger::new(cliargs.log);
 
     // If entropy analysis was requested, generate the entropy graph and return
-    if cliargs.entropy {
+    if let Some(entropy_output_file) = cliargs.entropy {
         display::print_plain(cliargs.quiet, "Calculating file entropy...");
 
-        if let Ok(entropy_results) = entropy::plot(cliargs.file_name.unwrap(), cliargs.stdin) {
+        if let Ok(entropy_results) = entropy::plot(
+            &entropy_output_file,
+            cliargs.file_name.unwrap(),
+            cliargs.stdin,
+        ) {
             // Log entropy results to JSON file, if requested
             json_logger.log(json::JSONType::Entropy(entropy_results.clone()));
             json_logger.close();