mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Feature added - MPSC channel
Build Library / Build Library (push) Successful in 5m44s
Details

Browse Source 
closes #38
This commit is contained in:
brotoskyj
2025-12-03 17:54:57 -05:00
parent 6ab9de413f
commit 76bd3a6087
9 changed files with 115 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
airlock_libs/src/services.rs
+73 -80
View File
@@ -1,5 +1,9 @@
use crate::prelude::*;
use std::thread;
use crossbeam::channel::unbounded;
use crate::modules::datatypes::*;
use crate::prelude::*;
#[pyfunction]
pub fn pull_policy_exec_histories(
py: Python<'_>,
@@ -34,7 +38,6 @@ pub fn pull_policy_exec_histories(
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()
@@ -61,6 +64,7 @@ pub fn pull_policy_exec_histories(
exechistories: vec![],
},
};
let writeable_filepath = file_path.clone();
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
match fs::write(writeable_filepath.clone(), data_write) {
Ok(_) => {}
@@ -75,7 +79,7 @@ pub fn pull_policy_exec_histories(
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}")
.template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len} {message}")
.unwrap(),
);
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
@@ -108,80 +112,41 @@ pub fn pull_policy_exec_histories(
}
});
let cutoff = Local::now().naive_local() - Duration::days(days);
let mut f = match File::open(&writeable_filepath) {
Ok(f) => f,
Err(e) => {
println!("Failed to Access {:?}: {}", &writeable_filepath, e);
std::process::abort();
}
};
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
span.set_attribute(Key::new("Days").string(days.to_string()));
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
loop {
match f.seek(SeekFrom::Start(0)) {
Ok(_) => {}
Err(e) => {
println!("Failed to seek start of {:?}: {}", f, e);
std::process::abort();
}
}
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
&base_url,
&exec_types,
&checkpoint_number,
&policy_names,
&client,
);
cx.span().set_attribute(KeyValue::new(
"items_in_response",
results.response.exechistories.len().to_string(),
));
results
});
let parsed_responses = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
let mut seen: HashMap<(String, String, String), Group> =
if writeable_filepath.exists() {
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() {
let (tx, rx) = unbounded::<Vec<Group>>();
thread::spawn(move || {
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
{
let contents = fs::read_to_string(&writeable_filepath).unwrap_or_default();
let existing: ApiResponse =
serde_json::from_str(&contents).unwrap_or(ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
});
existing
.response
.exechistories
.into_iter()
.map(|entry| {
(
(
entry.sha256.clone(),
entry.filename.clone(),
entry.hostname.clone(),
),
entry,
)
})
.collect()
} else {
HashMap::new()
};
while let Ok(parsed_responses) = rx.recv() {
for executions in parsed_responses {
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
continue;
}
if index == parsed_responses.len() - 1 {
checkpoint_number = executions.checkpoint.clone();
break;
}
let history_date = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
@@ -211,7 +176,34 @@ pub fn pull_policy_exec_histories(
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
}
}
if let Some(last_item) = &final_response.response.exechistories.last()
}
});
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
span.set_attribute(Key::new("Days").string(days.to_string()));
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
loop {
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging(
&base_url,
&exec_types,
&checkpoint_number,
&policy_names,
&client,
);
cx.span().set_attribute(KeyValue::new(
"items_in_response",
results.response.exechistories.len().to_string(),
));
results
});
let parsed_responses = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
tx.send(parsed_responses.clone()).unwrap();
checkpoint_number = parsed_responses.last().unwrap().checkpoint.clone();
if let Some(last_item) = parsed_responses.last()
&& let Ok(last_date) = NaiveDate::parse_from_str(
&last_item.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
@@ -219,20 +211,21 @@ pub fn pull_policy_exec_histories(
{
let date_diff = Local::now().naive_local().date() - last_date;
let percentage_diff =
(days - date_diff.num_days()) as f64 / days as f64 * 100.0;
progress_bar.set_position(percentage_diff.round() as u64);
((days - date_diff.num_days()) as f64 / days as f64 * 100.0).round() as u64;
progress_bar.set_position(percentage_diff);
progress_bar.set_message("Total Percent Complete");
}
}
});
progress_bar.finish_with_message("All Checkpoints Complete");
let return_data = match fs::read_to_string(&writeable_filepath) {
let return_data = match fs::read_to_string(file_path.clone()) {
Ok(return_data) => return_data,
Err(e) => {
println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
println!("Failed to read data from: {:?}: {}", &file_path, e);
std::process::abort();
}
};
drop(tx);
shutdown_tracer_provider();
return_data.to_string()
});
@@ -325,4 +318,4 @@ fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
.install_simple()
.unwrap();
Ok(Some(tracer))
}
}