From 960b7ff5e47f50158d7e751b033f68224538e81d Mon Sep 17 00:00:00 2001 From: brotoskyj Date: Tue, 4 Nov 2025 15:45:36 -0500 Subject: [PATCH] Working Library. Next commit will have it implemented in the script --- airlock_libs/Cargo.lock | 2 +- airlock_libs/src/services.rs | 104 ++++++++++++++++++++++++++++------- 2 files changed, 84 insertions(+), 22 deletions(-) diff --git a/airlock_libs/Cargo.lock b/airlock_libs/Cargo.lock index f5367f6..db62a6a 100644 --- a/airlock_libs/Cargo.lock +++ b/airlock_libs/Cargo.lock @@ -17,7 +17,7 @@ dependencies = [ [[package]] name = "airlock_libs" -version = "0.1.1" +version = "1.0.0-alpha" dependencies = [ "chrono", "indicatif", diff --git a/airlock_libs/src/services.rs b/airlock_libs/src/services.rs index aad6a80..cdca8bc 100644 --- a/airlock_libs/src/services.rs +++ b/airlock_libs/src/services.rs @@ -1,5 +1,5 @@ use chrono::{Duration, Local, NaiveDate}; -use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; +use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; use mongodb::bson::oid::ObjectId; use pyo3::prelude::*; use reqwest::{ @@ -8,7 +8,15 @@ use reqwest::{ }; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::{env, fmt::Write, fs, path::PathBuf, str::FromStr}; +use std::{ + collections::HashMap, + env, + fmt::Write, + fs::{self, File}, + io::Read, + path::PathBuf, + str::FromStr, +}; #[derive(Debug, Deserialize, Serialize)] struct ApiResponse { @@ -62,34 +70,37 @@ pub fn pull_policy_exec_histories( 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(); - } + && !parent_dir.exists() + { + fs::create_dir_all(parent_dir).unwrap(); + } fs::File::create(file_path).unwrap(); } - let mut data = ApiResponse { + 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, data_write).unwrap(); + fs::write(writeable_filepath.clone(), data_write).unwrap(); let mut checkpoint_number: String = skipback(days).to_string(); - let data_bar = - ProgressBar::with_draw_target(Some(10_000), ProgressDrawTarget::stdout_with_hz(255)); + 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: {msg} [{bar:40.cyan/blue}] {pos}/{len}") + .template("Checkpoint Progress: [{bar:40.cyan/blue}] {pos}/{len} {msg}") .unwrap(), ); data_bar.set_message("Starting"); - let progress_bar = ProgressBar::new(100); + let progress_bar = multi_progress.add(ProgressBar::new(100)); progress_bar.set_style( ProgressStyle::default_bar() - .template("Total Completion: {msg} [{bar:40.cyan/blue}] {pos}/{len}") + .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 = py_self; loop { let execution_histories = @@ -99,9 +110,36 @@ pub fn pull_policy_exec_histories( break; } data_bar.set_length(parsed_responses.len() as u64); - let mut batch = 0; + 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() { - batch += 1; if executions.checkpoint.is_empty() || executions.datetime.is_empty() { continue; } @@ -119,17 +157,41 @@ pub fn pull_policy_exec_histories( }; let cutoff = Local::now().naive_local() - Duration::days(days); if history_date >= cutoff.into() { - data.response.exechistories.push(executions.clone()); - } - if batch >= 50 { - data_bar.inc(batch); - batch = 0; + 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"); } - batch = 0; 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<'_>, @@ -146,7 +208,7 @@ async fn history_logging( 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(); + let val = HeaderValue::from_str(v).unwrap(); header_map.insert(HeaderName::from_str("X-APIKey").unwrap(), val); } }