hayabusa/src/afterfact.rs

use crate::detections::configs;
use crate::detections::print;
use crate::detections::print::AlertMessage;
use chrono::{DateTime, Local, TimeZone, Utc};
use serde::Serialize;
use std::error::Error;
use std::fs::File;
use std::io;
use std::process;

#[derive(Debug, Serialize)]
#[serde(rename_all = "PascalCase")]
pub struct CsvFormat<'a> {
    time: &'a str,
    computername: &'a str,
    eventid: &'a str,
    level: &'a str,
    alert: &'a str,
    details: &'a str,
    rulepath: &'a str,
    filepath: &'a str,
}

pub fn after_fact() {
    let fn_emit_csv_err = |err: Box<dyn Error>| {
        AlertMessage::alert(
            &mut std::io::stderr().lock(),
            format!("Failed to write CSV. {}", err),
        )
        .ok();
        process::exit(1);
    };

    let mut target: Box<dyn io::Write> = if let Some(csv_path) = configs::CONFIG
        .read()
        .unwrap()
        .args
        .value_of("csv-timeline")
    {
        // ファイル出力する場合
        match File::create(csv_path) {
            Ok(file) => Box::new(file),
            Err(err) => {
                AlertMessage::alert(
                    &mut std::io::stderr().lock(),
                    format!("Failed to open file. {}", err),
                )
                .ok();
                process::exit(1);
            }
        }
    } else {
        // 標準出力に出力する場合
        Box::new(io::stdout())
    };

    if let Err(err) = emit_csv(&mut target) {
        fn_emit_csv_err(err);
    }
}

fn emit_csv<W: std::io::Write>(writer: &mut W) -> Result<(), Box<dyn Error>> {
    let mut wtr = csv::WriterBuilder::new().from_writer(writer);
    let messages = print::MESSAGES.lock().unwrap();
    let mut detect_count = 0;
    for (time, detect_infos) in messages.iter() {
        for detect_info in detect_infos {
            wtr.serialize(CsvFormat {
                time: &format_time(time),
                filepath: &detect_info.filepath,
                rulepath: &detect_info.rulepath,
                level: &detect_info.level,
                computername: &detect_info.computername,
                eventid: &detect_info.eventid,
                alert: &detect_info.alert,
                details: &detect_info.detail,
            })?;
        }
        detect_count += detect_infos.len();
    }
    println!("");

    wtr.flush()?;
    println!("");
    println!("Events Detected:{:?}", detect_count);
    Ok(())
}

fn format_time(time: &DateTime<Utc>) -> String {
    if configs::CONFIG.read().unwrap().args.is_present("utc") {
        format_rfc(time)
    } else {
        format_rfc(&time.with_timezone(&Local))
    }
}

fn format_rfc<Tz: TimeZone>(time: &DateTime<Tz>) -> String
where
    Tz::Offset: std::fmt::Display,
{
    if configs::CONFIG.read().unwrap().args.is_present("rfc-2822") {
        return time.to_rfc2822();
    } else if configs::CONFIG.read().unwrap().args.is_present("rfc-3339") {
        return time.to_rfc3339();
    } else {
        return time.format("%Y-%m-%d %H:%M:%S%.3f %:z").to_string();
    }
}

#[test]
fn test_emit_csv() {
    use serde_json::Value;
    use std::fs::{read_to_string, remove_file};
    let testfilepath: &str = "test.evtx";
    let testrulepath: &str = "test-rule.yml";
    let test_title = "test_title";
    let test_level = "high";
    let test_computername = "testcomputer";
    let test_eventid = "1111";
    let output = "pokepoke";
    {
        let mut messages = print::MESSAGES.lock().unwrap();

        let val = r##"
        {
            "Event": {
                "EventData": {
                    "CommandRLine": "hoge"
                },
                "System": {
                    "TimeCreated_attributes": {
                        "SystemTime": "1996-02-27T01:05:01Z"
                    }
                }
            }
        }
    "##;
        let event: Value = serde_json::from_str(val).unwrap();
        messages.insert(
            testfilepath.to_string(),
            testrulepath.to_string(),
            &event,
            test_level.to_string(),
            test_computername.to_string(),
            test_eventid.to_string(),
            test_title.to_string(),
            output.to_string(),
        );
    }

    let expect_time = Utc
        .datetime_from_str("1996-02-27T01:05:01Z", "%Y-%m-%dT%H:%M:%SZ")
        .unwrap();
    let expect_tz = expect_time.with_timezone(&Local);
    let expect = "Time,Computername,Eventid,Level,Alert,Details,Rulepath,Filepath\n".to_string()
        + &expect_tz
            .clone()
            .format("%Y-%m-%d %H:%M:%S%.3f %:z")
            .to_string()
        + ","
        + test_computername
        + ","
        + test_eventid
        + ","
        + test_level
        + ","
        + test_title
        + ","
        + output
        + ","
        + testrulepath
        + ","
        + &testfilepath.replace(".evtx", "").to_string()
        + "\n";

    let mut file: Box<dyn io::Write> =
        Box::new(File::create("./test_emit_csv.csv".to_string()).unwrap());
    assert!(emit_csv(&mut file).is_ok());

    match read_to_string("./test_emit_csv.csv") {
        Err(_) => panic!("Failed to open file"),
        Ok(s) => {
            assert_eq!(s, expect);
        }
    };
    assert!(remove_file("./test_emit_csv.csv").is_ok());
}