add clio support for evtx2bodyfile

src/bin/es4forensics/cli.rs

@@ -1,27 +1,29 @@
-use clap::Parser;
-use log::LevelFilter;
 use crate::Protocol;
+use clap::{Parser, ValueHint};
+use clio::Input;
 use dfir_toolkit::common::HasVerboseFlag;
+use log::LevelFilter;
 
 #[cfg(feature = "gzip")]
-const INPUTFILE_HELP: &str = "path to input file or '-' for stdin (files ending with .gz will be treated as being gzipped)";
+const INPUTFILE_HELP: &str =
+    "path to input file or '-' for stdin (files ending with .gz will be treated as being gzipped)";
 #[cfg(not(feature = "gzip"))]
 const INPUTFILE_HELP: &str = "path to input file or '-' for stdin";
 
-#[derive(clap::Subcommand)]
-pub (crate) enum Action {
+#[derive(clap::Subcommand, Clone)]
+pub(crate) enum Action {
     // create a new index
     CreateIndex,
 
     // import timeline data
     Import {
-        #[clap(default_value="-", help=INPUTFILE_HELP)]
-        input_file: String,
+        #[clap(default_value="-", help=INPUTFILE_HELP, value_hint=ValueHint::FilePath)]
+        input_file: Input,
 
         /// number of timeline entries to combine in one bulk operation
-        #[clap(long("bulk-size"), default_value_t=1000)]
-        bulk_size: usize
-    }
+        #[clap(long("bulk-size"), default_value_t = 1000)]
+        bulk_size: usize,
+    },
 }
 
 #[derive(Parser)]
@@ -39,7 +41,7 @@ pub struct Cli {
     pub(crate) index_name: String,
 
     /// server name or IP address of elasticsearch server
-    #[clap( 
+    #[clap(
         short('H'),
         long("host"),
         display_order = 810,
@@ -56,7 +58,7 @@ pub struct Cli {
     pub(crate) protocol: Protocol,
 
     /// omit certificate validation
-    #[clap( 
+    #[clap(
         short('k'),
         long("insecure"),
         display_order = 840,
@@ -77,7 +79,7 @@ pub struct Cli {
 }
 
 impl HasVerboseFlag for Cli {
-    fn log_level_filter(&self)-> LevelFilter {
+    fn log_level_filter(&self) -> LevelFilter {
         self.verbose.log_level_filter()
     }
-}
+}

src/bin/es4forensics/main.rs

@@ -7,41 +7,48 @@ use anyhow::{Result, anyhow};
 use cli::{Cli, Action};
 use elasticsearch::auth::Credentials;
 use dfir_toolkit::es4forensics::*;
-use dfir_toolkit::common::FancyParser;
+use dfir_toolkit::common::{FancyParser, FileInput};
 
 #[tokio::main]
 async fn main() -> Result<()> {
     let cli: Cli = Cli::parse_cli();
-
+
+    let action = cli.action.clone();
     let e4f: Es4Forensics = cli.into();
-    e4f.run().await
+    e4f.run(action).await
 }
 
 struct Es4Forensics {
-    cli: Cli
+    strict_mode: bool,
+    index_name: String,
+    host: String,
+    port: u16,
+    protocol: Protocol,
+    omit_certificate_validation: bool,
+    username: String,
+    password: String,
 }
 
 impl Es4Forensics {
-    pub async fn run(self) -> Result<()> {
+    pub async fn run(self, action: Action) -> Result<()> {
 
         let builder = self.create_index_builder()?;
 
-        match &self.cli.action {
+        match action {
             Action::CreateIndex => {
                 if builder.index_exists().await? {
-                    return Err(anyhow!("index '{}' exists already", self.cli.index_name));
+                    return Err(anyhow!("index '{}' exists already", self.index_name));
                 }
                 builder.create_index().await?;
                 Ok(())
             }
             Action::Import{input_file, bulk_size} => {
-                let source = StreamSource::from(input_file)?;
-                self.import(builder, source.into(), *bulk_size).await
+                self.import(builder, input_file.into(), bulk_size).await
             }
         }
     }
 
-    async fn import(&self, builder: IndexBuilder, reader: Box<dyn BufRead + Send>, bulk_size: usize) -> Result<()> {
+    async fn import(&self, builder: IndexBuilder, reader: FileInput, bulk_size: usize) -> Result<()> {
         let mut index = builder.connect().await?;
         index.set_cache_size(bulk_size).await?;
 
@@ -50,7 +57,7 @@ impl Es4Forensics {
             let value = match serde_json::from_str(&line) {
                 Ok(v) => v,
                 Err(why) => {
-                    if self.cli.strict_mode {
+                    if self.strict_mode {
                         return Err(anyhow!(why))
                     } else {
                         ::log::error!("error while parsing: {}", why);
@@ -67,16 +74,16 @@ impl Es4Forensics {
     }    
 
     fn create_index_builder(&self) -> Result<IndexBuilder> {
-        let mut builder = IndexBuilder::with_name(self.cli.index_name.clone())
-            .with_host(self.cli.host.clone())
-            .with_port(self.cli.port)
+        let mut builder = IndexBuilder::with_name(self.index_name.clone())
+            .with_host(self.host.clone())
+            .with_port(self.port)
             .with_credentials(Credentials::Basic(
-                self.cli.username.clone(),
-                self.cli.password.clone(),
+                self.username.clone(),
+                self.password.clone(),
             ))
-            .with_protocol(self.cli.protocol.clone());
+            .with_protocol(self.protocol.clone());
 
-        if self.cli.omit_certificate_validation {
+        if self.omit_certificate_validation {
             ::log::warn!("disabling certificate validation");
             builder = builder.without_certificate_validation();
         }
@@ -88,7 +95,14 @@ impl Es4Forensics {
 impl From<Cli> for Es4Forensics {
     fn from(cli: Cli) -> Self {
         Self {
-            cli
+            strict_mode: cli.strict_mode,
+            host: cli.host.clone(),
+            port: cli.port,
+            username: cli.username.clone(),
+            password: cli.password.clone(),
+            index_name: cli.index_name.clone(),
+            protocol: cli.protocol.clone(),
+            omit_certificate_validation: cli.omit_certificate_validation,
         }
     }
 }

src/bin/evtx2bodyfile/evtx2bodyfile_app.rs

@@ -1,7 +1,8 @@
 
 use crate::{bf_data::*, evtx_file::EvtxFile};
 use anyhow::Result;
-use clap::Parser;
+use clap::{Parser, ValueHint};
+use clio::Input;
 use dfir_toolkit::common::HasVerboseFlag;
 use evtx::SerializedEvtxRecord;
 use getset::Getters;
@@ -12,7 +13,8 @@ use serde_json::Value;
 #[clap(name=env!("CARGO_BIN_NAME"), author, version, about, long_about = None)]
 pub(crate) struct Evtx2BodyfileApp {
     /// names of the evtx files
-    evtx_files: Vec<String>,
+    #[clap(value_hint=ValueHint::FilePath)]
+    evtx_files: Vec<Input>,
 
     /// output json for elasticsearch instead of bodyfile
     #[clap(short('J'), long("json"))]
@@ -29,8 +31,8 @@ pub(crate) struct Evtx2BodyfileApp {
 
 impl Evtx2BodyfileApp {
     pub(crate) fn handle_evtx_files(&self) -> Result<()> {
-        for file in self.evtx_files.iter() {
-            self.handle_evtx_file((&file[..]).try_into()?);
+        for file in self.evtx_files.clone().into_iter() {
+            self.handle_evtx_file(file.into());
         }
         Ok(())
     }

src/bin/evtx2bodyfile/evtx_file.rs

@@ -1,18 +1,15 @@
-use std::{path::PathBuf, fs::File};
 use anyhow::Result;
-use evtx::{EvtxParser, SerializedEvtxRecord, err::EvtxError};
+use clio::Input;
+use evtx::{EvtxParser, SerializedEvtxRecord};
 use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
 use serde_json::Value;
 use ouroboros::self_referencing;
 
-pub (crate) struct EvtxFile {
-    filename: String,
-    fp: PathBuf,
-}
+pub (crate) struct EvtxFile (Input);
 
 #[self_referencing()]
 pub (crate) struct EvtxFileIterator {
-    parser: EvtxParser<File>,
+    parser: EvtxParser<Input>,
 
     #[borrows(mut parser)]
     #[not_covariant]
@@ -35,39 +32,31 @@ impl Iterator for EvtxFileIterator {
     }
 }
 
-impl<'a> IntoIterator for &'a EvtxFile {
+impl IntoIterator for EvtxFile {
     type Item = SerializedEvtxRecord<Value>;
 
     type IntoIter = EvtxFileIterator;
 
     fn into_iter(self) -> Self::IntoIter {
-        let parser = EvtxParser::from_path(&self.fp).expect("unable to create parser");
+        let parser = EvtxParser::from_read_seek(self.0).expect("unable to create parser");
 
         EvtxFileIteratorBuilder {
             parser,
-            inner_iterator_builder: |parser: &mut EvtxParser<File>| Box::new(parser.records_json_value())
+            inner_iterator_builder: |parser: &mut EvtxParser<Input>| Box::new(parser.records_json_value())
         }.build()
     }
 }
 
-impl TryFrom<&str> for EvtxFile {
-    type Error = EvtxError;
-    fn try_from(file: &str) -> std::result::Result<Self, Self::Error> {
-        let fp = PathBuf::from(file);
-        let filename = fp.file_name().unwrap().to_str().unwrap().to_owned();
-
-        Ok(Self {
-            filename,
-            fp: PathBuf::from(file),
-        })
+impl From<Input> for EvtxFile {
+    fn from(file: Input) -> Self {
+        Self(file)
     }
-
 }
 
 
 impl EvtxFile {
     pub (crate) fn count_records(&self) -> Result<usize> {
-        let mut parser = EvtxParser::from_path(&self.fp)?;
+        let mut parser = EvtxParser::from_read_seek(self.0.clone())?;
         Ok(parser.serialized_records(|r| r.and(Ok(()))).count())
     }
 
@@ -77,7 +66,6 @@ impl EvtxFile {
         let bar = ProgressBar::new(count as u64);
         let target = ProgressDrawTarget::stderr_with_hz(10);
         bar.set_draw_target(target);
-        bar.set_message(self.filename.clone());
 
         let progress_style = ProgressStyle::default_bar()
             .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>9}/{len:9}({percent}%) {msg}")?

src/common/file_input.rs

@@ -0,0 +1,72 @@
+use std::io::{BufRead, BufReader, Read};
+
+use clio::Input;
+use flate2::bufread::GzDecoder;
+
+pub struct FileInput {
+    stream: Box<dyn BufRead>,
+    input: Input
+}
+
+impl From<Input> for FileInput {
+    fn from(value: Input) -> Self {
+        let cloned_input = value.clone();
+        if let Some(extension) = value.path().extension() {
+            if extension == "gz" {
+                return Self {
+                    stream: Box::new(BufReader::new(GzDecoder::new(BufReader::new(value)))),
+                    input: cloned_input
+                };
+            }
+        }
+
+        Self {
+            stream: Box::new(BufReader::new(value)),
+            input: cloned_input
+        }
+    }
+}
+
+impl From<&Input> for FileInput {
+    fn from(value: &Input) -> Self {
+        Self::from(value.clone())
+    }
+}
+
+
+impl From<&mut Input> for FileInput {
+    fn from(value: &mut Input) -> Self {
+        Self::from(value.clone())
+    }
+}
+
+impl Clone for FileInput {
+    fn clone(&self) -> Self {
+        self.input.clone().into()
+    }
+}
+
+impl Read for FileInput {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.stream.read(buf)
+    }
+}
+
+impl BufRead for FileInput {
+    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
+        self.stream.fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.stream.consume(amt)
+    }
+}
+
+impl<'a> TryFrom<&'a str> for FileInput {
+    type Error = <Input as TryFrom<&'a str>>::Error;
+
+    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
+        let res = Input::try_from(value)?;
+        Ok(res.into())
+    }
+}

src/common/mod.rs

@@ -1,6 +1,8 @@
 mod rfc3339_datetime;
 pub mod bodyfile;
 mod parse_cli;
+mod file_input;
 
 pub use rfc3339_datetime::*;
-pub use parse_cli::*;
+pub use parse_cli::*;
+pub use file_input::*;