Merge pull request #795 from ReFirmLabs/entropy_stdout

Added --png argument; some entropy code clean up

src/cliparser.rs

@@ -11,7 +11,7 @@ pub struct CliArgs {
     #[arg(short, long)]
     pub stdin: bool,
 
-    /// Supress output to stdout
+    /// Supress normal stdout output
     #[arg(short, long)]
     pub quiet: bool,
 
@@ -35,9 +35,13 @@ pub struct CliArgs {
     #[arg(short = 'a', long)]
     pub search_all: bool,
 
-    /// Write entropy plot image to the ENTROPY file
+    /// Generate a PNG entropy graph
     #[arg(short = 'E', long, conflicts_with = "extract")]
-    pub entropy: Option<String>,
+    pub entropy: bool,
+
+    /// Specify an alternate file name for the PNG entropy graph
+    #[arg(short, long, requires = "entropy")]
+    pub png: Option<String>,
 
     /// Log JSON results to a file ('-' for stdout)
     #[arg(short, long)]

src/entropy.rs

@@ -23,34 +23,28 @@ pub struct FileEntropy {
 
 /// Splits the supplied data up into blocks and calculates the entropy of each block.
 fn blocks(data: &[u8]) -> Vec<BlockEntropy> {
-    const MIN_BLOCK_SIZE: usize = 1024;
-    const NUM_DATA_POINTS: usize = 4096 * 10;
+    const BLOCK_SIZE: usize = 1024 * 4;
 
     let mut offset: usize = 0;
+    let mut chunker = data.chunks(BLOCK_SIZE);
     let mut entropy_blocks: Vec<BlockEntropy> = vec![];
-    let mut block_size: usize = (data.len() as f64 / NUM_DATA_POINTS as f64).ceil() as usize;
 
-    if block_size < MIN_BLOCK_SIZE {
-        block_size = MIN_BLOCK_SIZE;
-    }
-
-    while offset < data.len() {
-        let mut block = BlockEntropy {
-            ..Default::default()
-        };
-
-        block.start = offset;
-        block.end = block.start + block_size;
-
-        if block.end > data.len() {
-            block.end = data.len();
+    loop {
+        match chunker.next() {
+            None => break,
+            Some(block_data) => {
+                let mut block = BlockEntropy {
+                    ..Default::default()
+                };
+
+                block.start = offset;
+                block.entropy = shannon_entropy(block_data);
+                block.end = block.start + block_data.len();
+
+                offset = block.end;
+                entropy_blocks.push(block);
+            }
         }
-
-        block.entropy = shannon_entropy(&data[block.start..block.end]);
-
-        entropy_blocks.push(block);
-
-        offset += block_size;
     }
 
     entropy_blocks
@@ -59,22 +53,22 @@ 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(
-    png_file_path: impl Into<String>,
     file_path: impl Into<String>,
     stdin: bool,
+    png_file_path: Option<String>,
 ) -> Result<FileEntropy, EntropyError> {
-    const FILE_EXTENSION: &str = ".png";
+    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);
-    }
+    // Use the specified output file path, else generate a default output file name
+    let png_path: String = match png_file_path {
+        Some(fpath) => fpath,
+        None => format!("{}.{}", target_file, FILE_EXTENSION),
+    };
 
     // Get the base name of the target file
     let target_file_name = path::Path::new(&target_file)

src/main.rs

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