Skip to content

Commit

Permalink
add clio support for evtx2bodyfile
Browse files Browse the repository at this point in the history
  • Loading branch information
Jan Starke committed Sep 21, 2023
1 parent 6d24e80 commit d2f3f87
Show file tree
Hide file tree
Showing 6 changed files with 141 additions and 61 deletions.
30 changes: 16 additions & 14 deletions src/bin/es4forensics/cli.rs
Original file line number Diff line number Diff line change
@@ -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)]
Expand All @@ -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,
Expand All @@ -56,7 +58,7 @@ pub struct Cli {
pub(crate) protocol: Protocol,

/// omit certificate validation
#[clap(
#[clap(
short('k'),
long("insecure"),
display_order = 840,
Expand All @@ -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()
}
}
}
52 changes: 33 additions & 19 deletions src/bin/es4forensics/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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?;

Expand All @@ -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);
Expand All @@ -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();
}
Expand All @@ -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,
}
}
}
10 changes: 6 additions & 4 deletions src/bin/evtx2bodyfile/evtx2bodyfile_app.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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"))]
Expand All @@ -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(())
}
Expand Down
34 changes: 11 additions & 23 deletions src/bin/evtx2bodyfile/evtx_file.rs
Original file line number Diff line number Diff line change
@@ -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]
Expand All @@ -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())
}

Expand All @@ -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}")?
Expand Down
72 changes: 72 additions & 0 deletions src/common/file_input.rs
Original file line number Diff line number Diff line change
@@ -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())
}
}
4 changes: 3 additions & 1 deletion src/common/mod.rs
Original file line number Diff line number Diff line change
@@ -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::*;

0 comments on commit d2f3f87

Please sign in to comment.