diff --git a/yara-x-cli/src/commands/fmt.rs b/yara-x-cli/src/commands/fmt.rs index 131b0cd81..818b8dbdc 100644 --- a/yara-x-cli/src/commands/fmt.rs +++ b/yara-x-cli/src/commands/fmt.rs @@ -2,6 +2,7 @@ use std::fs; use std::fs::File; use std::io::{stdin, stdout, Cursor, Seek, Write}; use std::path::PathBuf; +use std::process; use clap::{arg, value_parser, ArgAction, ArgMatches, Command}; use yara_x_fmt::Formatter; @@ -18,29 +19,47 @@ pub fn fmt() -> Command { arg!(-w --write ... "Write output to source file instead of stdout") .action(ArgAction::SetTrue), ) + .arg( + arg!(-t --test ... "Exit with failure if reformatting changed the file") + .action(ArgAction::SetTrue), + ) } pub fn exec_fmt(args: &ArgMatches) -> anyhow::Result<()> { let files = args.get_many::("FILE"); let write = args.get_one::("write"); + let test = args.get_one::("test"); let formatter = Formatter::new(); if let Some(files) = files { + let mut changed_files = Vec::new(); + for file in files { let input = fs::read(file.as_path())?; - if *write.unwrap() { - let mut formatted = Cursor::new(Vec::new()); - formatter.format(input.as_slice(), &mut formatted)?; - formatted.rewind()?; + let mut formatted = Cursor::new(Vec::new()); + + formatter.format(input.as_slice(), &mut formatted)?; + formatted.rewind()?; - File::create(file.as_path())? - .write_all(formatted.into_inner().as_slice())?; + let output = formatted.into_inner(); + + if *test.unwrap() && input != output { + changed_files.push(file.display().to_string()); + } + + if *write.unwrap() { + File::create(file.as_path())?.write_all(output.as_slice())?; } else { - formatter.format(input.as_slice(), stdout())?; + print!("{}", String::from_utf8(output)?); }; } + + if changed_files.len() >= 1 { + eprintln!("File(s) to format: {}", changed_files.join(", ")); + process::exit(2) + } } else { formatter.format(stdin(), stdout())?; }