Skip to content

Commit

Permalink
Added ability to specify an output file name for entropy plots
Browse files Browse the repository at this point in the history
  • Loading branch information
devttys0 committed Dec 6, 2024
1 parent 6988c06 commit d751375
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 9 deletions.
4 changes: 2 additions & 2 deletions src/cliparser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
18 changes: 13 additions & 5 deletions src/entropy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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);
Expand Down
8 changes: 6 additions & 2 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down

0 comments on commit d751375

Please sign in to comment.