brotoskyj
278 lines
9.3 KiB
Rust
278 lines
9.3 KiB
Rust
use chrono::{Duration, Local, NaiveDate};
|
|
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
|
|
use mongodb::bson::oid::ObjectId;
|
|
use pyo3::prelude::*;
|
|
use reqwest::{
|
|
Client,
|
|
header::{HeaderMap, HeaderName, HeaderValue},
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use std::{
|
|
collections::HashMap,
|
|
env,
|
|
fmt::Write,
|
|
fs::{self, File},
|
|
io::Read,
|
|
path::PathBuf,
|
|
str::FromStr,
|
|
};
|
|
|
|
#[derive(Debug, Deserialize, Serialize)]
|
|
struct ApiResponse {
|
|
error: String,
|
|
response: ExecHistories,
|
|
}
|
|
#[derive(Debug, Deserialize, Serialize)]
|
|
struct ExecHistories {
|
|
exechistories: Vec<Group>,
|
|
}
|
|
#[derive(Debug, Deserialize, Serialize, Clone)]
|
|
struct Group {
|
|
checkpoint: String,
|
|
#[serde(rename = "type")]
|
|
exectype: u8,
|
|
username: String,
|
|
hostname: String,
|
|
netdomain: String,
|
|
filename: String,
|
|
ppolicy: String,
|
|
policyname: String,
|
|
policyver: String,
|
|
commandline: String,
|
|
publisher: String,
|
|
pprocess: String,
|
|
gprocess: String,
|
|
sha256: String,
|
|
datetime: String,
|
|
md5: String,
|
|
sha128: String,
|
|
sha384: String,
|
|
sha512: String,
|
|
ip: String,
|
|
localip: String,
|
|
}
|
|
|
|
#[pyfunction]
|
|
pub fn pull_policy_exec_histories(
|
|
py: Python<'_>,
|
|
py_self: Py<PyAny>,
|
|
policy_names: String,
|
|
exec_types: String,
|
|
days: i64,
|
|
//output_json: bool,
|
|
) {
|
|
let file_path: PathBuf = format!(
|
|
"{}\\cache\\chunkinator.json",
|
|
get_base_directory().display()
|
|
)
|
|
.into();
|
|
let writeable_filepath = file_path.clone();
|
|
if !file_path.exists() {
|
|
if let Some(parent_dir) = file_path.parent()
|
|
&& !parent_dir.exists()
|
|
{
|
|
fs::create_dir_all(parent_dir).unwrap();
|
|
}
|
|
fs::File::create(file_path).unwrap();
|
|
}
|
|
let data = ApiResponse {
|
|
error: "Success".to_string(),
|
|
response: ExecHistories {
|
|
exechistories: vec![],
|
|
},
|
|
};
|
|
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
|
|
fs::write(writeable_filepath.clone(), data_write).unwrap();
|
|
let mut checkpoint_number: String = skipback(days).to_string();
|
|
let multi_progress = MultiProgress::new();
|
|
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
|
|
let data_bar = multi_progress.add(ProgressBar::new(10_000));
|
|
data_bar.set_style(
|
|
ProgressStyle::default_bar()
|
|
.template("Checkpoint Progress: [{bar:40.cyan/blue}] {pos}/{len} {msg}")
|
|
.unwrap(),
|
|
);
|
|
data_bar.set_message("Starting");
|
|
let progress_bar = multi_progress.add(ProgressBar::new(100));
|
|
progress_bar.set_style(
|
|
ProgressStyle::default_bar()
|
|
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
|
|
.unwrap(),
|
|
);
|
|
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
|
|
let api: Py<PyAny> = py_self;
|
|
loop {
|
|
let execution_histories =
|
|
history_logging(py, &api, &exec_types, &checkpoint_number, &policy_names);
|
|
let parsed_responses = execution_histories.response.exechistories;
|
|
if parsed_responses.is_empty() {
|
|
break;
|
|
}
|
|
data_bar.set_length(parsed_responses.len() as u64);
|
|
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
|
|
let mut f = File::open(&writeable_filepath).unwrap();
|
|
let mut contents = String::new();
|
|
f.read_to_string(&mut contents).unwrap();
|
|
let existing_data: ApiResponse =
|
|
serde_json::from_str(&contents).unwrap_or(ApiResponse {
|
|
error: "Success".to_string(),
|
|
response: ExecHistories {
|
|
exechistories: vec![],
|
|
},
|
|
});
|
|
existing_data
|
|
.response
|
|
.exechistories
|
|
.into_iter()
|
|
.map(|entry| {
|
|
(
|
|
(
|
|
entry.sha256.clone(),
|
|
entry.filename.clone(),
|
|
entry.hostname.clone(),
|
|
),
|
|
entry,
|
|
)
|
|
})
|
|
.collect()
|
|
} else {
|
|
HashMap::new()
|
|
};
|
|
for (index, executions) in parsed_responses.iter().enumerate() {
|
|
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
|
|
continue;
|
|
}
|
|
if index == parsed_responses.len() - 1 {
|
|
checkpoint_number = executions.checkpoint.clone();
|
|
data_bar.set_message(checkpoint_number.clone());
|
|
break;
|
|
}
|
|
let history_date = match NaiveDate::parse_from_str(
|
|
&executions.datetime.replace(" +0000 UTC", ""),
|
|
"%Y-%m-%dT%H:%M:%SZ",
|
|
) {
|
|
Ok(date) => date,
|
|
Err(_) => continue,
|
|
};
|
|
let cutoff = Local::now().naive_local() - Duration::days(days);
|
|
if history_date >= cutoff.into() {
|
|
let key = (
|
|
executions.sha256.clone(),
|
|
executions.filename.clone(),
|
|
executions.hostname.clone(),
|
|
);
|
|
seen.entry(key).or_insert(executions.clone());
|
|
}
|
|
data_bar.inc(1);
|
|
}
|
|
let final_response = ApiResponse {
|
|
error: "Success".to_string(),
|
|
response: ExecHistories {
|
|
exechistories: seen.values().cloned().collect(),
|
|
},
|
|
};
|
|
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
|
|
fs::write(&writeable_filepath, data_write).unwrap();
|
|
if let Some(last_item) = &final_response.response.exechistories.last()
|
|
&& let Ok(last_date) = NaiveDate::parse_from_str(
|
|
&last_item.datetime.replace(" +0000 UTC", ""),
|
|
"%Y-%m-%dT%H:%M:%SZ",
|
|
)
|
|
{
|
|
let date_diff = Local::now().naive_local().date() - last_date;
|
|
let percentage_diff =
|
|
((days + 10) - date_diff.num_days()) as f64 / (days + 10) as f64 * 100.0;
|
|
progress_bar.set_position(percentage_diff.round() as u64);
|
|
progress_bar.set_message("Total Percent Complete");
|
|
}
|
|
data_bar.set_position(0);
|
|
}
|
|
data_bar.finish_with_message("Finished Checkpoints");
|
|
progress_bar.finish_with_message("All Checkpoints Complete");
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn history_logging(
|
|
py: Python<'_>,
|
|
py_self: &Py<PyAny>,
|
|
exec_types: &String,
|
|
checkpoint_number: &String,
|
|
policy_names: &String,
|
|
) -> ApiResponse {
|
|
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
|
|
let headers = py_self.getattr(py, "headers").unwrap().to_string();
|
|
let headers_replace = headers.replace('\'', "\"");
|
|
let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
|
|
let mut header_map = HeaderMap::new();
|
|
if let Some(obj) = parsed.as_object() {
|
|
for (_key, value) in obj {
|
|
if let Some(v) = value.as_str() {
|
|
let val = HeaderValue::from_str(v).unwrap();
|
|
header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val);
|
|
}
|
|
}
|
|
}
|
|
let payload = format!(
|
|
r#"{{
|
|
"type": {},
|
|
"checkpoint": "{}",
|
|
"policy": ["{}"]
|
|
}}"#,
|
|
exec_types, checkpoint_number, policy_names
|
|
);
|
|
let client = Client::builder()
|
|
.danger_accept_invalid_certs(true)
|
|
.default_headers(header_map)
|
|
.timeout(std::time::Duration::from_secs(30))
|
|
.build()
|
|
.unwrap();
|
|
let res = client
|
|
.post(format!("{}/v1/logging/exechistories", base_url))
|
|
.body(payload)
|
|
.send()
|
|
.await;
|
|
match res {
|
|
Ok(res) => {
|
|
let first_response: ApiResponse = serde_json::from_str(&res.text().await.unwrap())
|
|
.expect("Failed to retrieve response from API");
|
|
return first_response;
|
|
}
|
|
Err(_res) => {
|
|
let failed_response: ApiResponse = ApiResponse {
|
|
error: "Failed".to_string(),
|
|
response: ExecHistories {
|
|
exechistories: vec![],
|
|
},
|
|
};
|
|
return failed_response;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn get_base_directory() -> PathBuf {
|
|
let home = env::var_os("HOME")
|
|
.map(PathBuf::from)
|
|
.or_else(|| env::var_os("USERPROFILE").map(PathBuf::from))
|
|
.expect("Could not find Home Directory");
|
|
let os = std::env::consts::OS;
|
|
match os {
|
|
"windows" => {
|
|
let appdata = env::var_os("APPDATA")
|
|
.map(PathBuf::from)
|
|
.unwrap_or_else(|| home.join("AppData").join("Roaming"));
|
|
appdata.join("AirlockTools")
|
|
}
|
|
_ => home.join(".local").join("share").join("AirlockTools"),
|
|
}
|
|
}
|
|
|
|
fn skipback(days: i64) -> ObjectId {
|
|
let date_days_ago = Local::now() - Duration::days(days);
|
|
let timestamp = date_days_ago.timestamp() as u32;
|
|
let mut hex_timestamp = String::new();
|
|
write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap();
|
|
let objectid_hex = format!("{}0000000000000000", hex_timestamp);
|
|
ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
|
|
}
|