print_file_names & pretty_print command line optional arguments

shindan-io · Jul 2, 2024 · 3a523d8 · 3a523d8
1 parent e5db767
commit 3a523d8
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 23 deletions.
diff --git a/py_scnr/src/lib.rs b/py_scnr/src/lib.rs
@@ -41,41 +41,45 @@ fn activate_verbose(verbose: bool) {
 }
 
 #[pyfunction]
-#[pyo3(signature = (*, input = DEFAULT_INPUT.to_string(), filter=vec![], starter=vec![], cfg=vec![], profile=CfgProfile::default(), verbose=false))]
+#[pyo3(signature = (*, input = DEFAULT_INPUT.to_string(), filter=vec![], starter=vec![], cfg=vec![], profile=CfgProfile::default(), print_file_names=false, pretty_print=false, verbose=false))]
 fn scan(
   input: String,
   filter: Vec<String>,
   starter: Vec<Plugin>,
   cfg: Vec<(String, Plugin)>,
   profile: CfgProfile,
+  print_file_names: bool,
+  pretty_print: bool,
   verbose: bool,
 ) -> Result<ScanResultIterator, PyScnrError> {
   activate_verbose(verbose);
   let starter = to_scnr_starter(starter);
   let cfg = to_scnr_cfg(cfg);
   let profile = profile.into();
-  let common = CommonArgs { input, filter, starter, cfg, profile };
+  let common = CommonArgs { input, filter, starter, cfg, profile, print_file_names, pretty_print };
   let scanner = scnr::get_scanner_from_options(&common)?;
   let result = scanner.scan()?;
   Ok(result.into())
 }
 
 #[pyfunction]
-#[pyo3(signature = (*, input = DEFAULT_INPUT.to_string(), query = DEFAULT_JQ_QUERY, filter=vec![], starter=vec![], cfg=vec![], profile=CfgProfile::default(), verbose=false))]
+#[pyo3(signature = (*, input = DEFAULT_INPUT.to_string(), query = DEFAULT_JQ_QUERY, filter=vec![], starter=vec![], cfg=vec![], profile=CfgProfile::default(), print_file_names=false, pretty_print=false, verbose=false))]
 fn jq(
   input: String,
   query: &str,
   filter: Vec<String>,
   starter: Vec<Plugin>,
   cfg: Vec<(String, Plugin)>,
   profile: CfgProfile,
+  print_file_names: bool,
+  pretty_print: bool,
   verbose: bool,
 ) -> Result<JqIterator, PyScnrError> {
   activate_verbose(verbose);
   let starter = to_scnr_starter(starter);
   let cfg = to_scnr_cfg(cfg);
   let profile = profile.into();
-  let common = CommonArgs { input, filter, starter, cfg, profile };
+  let common = CommonArgs { input, filter, starter, cfg, profile, print_file_names, pretty_print };
   let scanner = scnr::get_scanner_from_options(&common)?;
   let result = scanner.scan()?;
   let iterator = JqIterator::new(result, query)?;

diff --git a/scnr/src/main.rs b/scnr/src/main.rs
@@ -1,8 +1,8 @@
 #![allow(clippy::default_trait_access, clippy::module_name_repetitions, clippy::wildcard_imports)]
 #![deny(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
 
-use scnr_core::{bin_repr, jq, Scanner};
-use std::io::Write;
+use scnr_core::{bin_repr, jq, Content, Scanner};
+use std::{io::Write, path::PathBuf};
 
 use scnr::options::*;
 
@@ -26,6 +26,32 @@ fn main() -> anyhow::Result<()> {
   Ok(())
 }
 
+fn print_path(out: &mut impl Write, path: &PathBuf, options: &CommonArgs) -> anyhow::Result<()> {
+  if options.print_file_names {
+    writeln!(out, "{}", path.display())?;
+  }
+  Ok(())
+}
+
+fn print_content(out: &mut impl Write, content: &Content, options: &CommonArgs) -> anyhow::Result<()> {
+  match &content {
+    scnr_core::Content::Json(json) => {
+      let consume_out = &mut *out;
+      if options.pretty_print {
+        serde_json::to_writer_pretty(consume_out, &json)?;
+      } else {
+        serde_json::to_writer(consume_out, &json)?;
+      }
+    }
+    scnr_core::Content::Text(text) => writeln!(out, "{text}")?,
+    scnr_core::Content::Bytes(bytes) => writeln!(out, "{}", bin_repr::BinRepr::Base64.to_string(&bytes))?,
+  }
+
+  writeln!(out)?;
+
+  Ok(())
+}
+
 #[tracing::instrument(skip(scanner), err)]
 fn scan(scanner: Scanner, args: ScanArgs) -> anyhow::Result<()> {
   let stdout = std::io::stdout();
@@ -36,12 +62,8 @@ fn scan(scanner: Scanner, args: ScanArgs) -> anyhow::Result<()> {
   for content in iter {
     match content {
       Ok(content) => {
-        writeln!(lock, "{}", content.rel_path.display())?;
-        match content.content {
-          scnr_core::Content::Json(json) => serde_json::to_writer_pretty(&mut lock, &json)?,
-          scnr_core::Content::Text(text) => writeln!(lock, "{text}")?,
-          scnr_core::Content::Bytes(bytes) => writeln!(lock, "{}", bin_repr::BinRepr::Base64.to_string(&bytes))?,
-        }
+        print_path(&mut lock, &content.rel_path, &args.common)?;
+        print_content(&mut lock, &content.content, &args.common)?;
       }
       Err(err) => tracing::error!("{err:?}"),
     }
@@ -63,13 +85,9 @@ fn jq(scanner: Scanner, args: JqArgs) -> anyhow::Result<()> {
     match content {
       Ok(content) => {
         if let Some(json) = content.content.json() {
+          print_path(&mut lock, &content.rel_path, &args.common)?;
           for element in jq_filter.run(json)? {
-            if args.no_pretty_print {
-              serde_json::to_writer(&mut lock, &element)?;
-            } else {
-              serde_json::to_writer_pretty(&mut lock, &element)?;
-            }
-            writeln!(lock)?;
+            print_content(&mut lock, &Content::Json(element), &args.common)?;
           }
         }
       }

diff --git a/scnr/src/options.rs b/scnr/src/options.rs
@@ -47,11 +47,25 @@ pub struct CommonArgs {
 
   #[arg(short, long, default_value_t = CfgProfile::default(),  help = "Plugins configuration profile to start with. Profiles are cfg bundles and can be then overridden by cfg args")]
   pub profile: CfgProfile,
+
+  #[arg(long, short = 'n', help = "DO print the file names (before the content)")]
+  pub print_file_names: bool,
+
+  #[arg(long, short = 'b', help = "DO pretty(beautiful) print the output")]
+  pub pretty_print: bool,
 }
 
 impl Default for CommonArgs {
   fn default() -> Self {
-    CommonArgs { input: DEFAULT_INPUT.to_string(), filter: vec![], profile: CfgProfile::default(), cfg: vec![], starter: vec![] }
+    CommonArgs {
+      input: DEFAULT_INPUT.to_string(),
+      filter: vec![],
+      profile: CfgProfile::default(),
+      cfg: vec![],
+      starter: vec![],
+      print_file_names: false,
+      pretty_print: false,
+    }
   }
 }
 
@@ -139,9 +153,6 @@ pub struct JqArgs {
 
   #[arg(long, short, help = "Jq query to apply to all 'json-ed' results")]
   pub query: String,
-
-  #[arg(long, short = 'n', help = "Do NOT pretty print the json output")]
-  pub no_pretty_print: bool,
 }
 
 // =================================================================================================
@@ -182,7 +193,7 @@ mod tests {
   #[test]
   fn parse_cmd_2() {
     let cmd =
-      "scnr -v extract --output /tmp -f *.json --filter=**/*.xml --force -p sysdiagnose --cfg img.svg=json --cfg *.toml=text -s file-system";
+      "scnr -v extract --output /tmp -f *.json --filter=**/*.xml --force -p sysdiagnose --cfg img.svg=json --cfg *.toml=text -s file-system -nb";
     let opts = Opts::parse_from(cmd.split(' '));
     assert!(opts.verbose);
     assert_eq!(
@@ -194,6 +205,8 @@ mod tests {
           profile: CfgProfile::Sysdiagnose,
           cfg: vec![("img.svg".into(), Plugin::Json), ("*.toml".into(), Plugin::Text)],
           starter: vec![Plugin::FileSystem],
+          print_file_names: true,
+          pretty_print: true
         },
         output: PathBuf::from("/tmp"),
         force: true,