Compare commits

...

3 Commits

Author SHA1 Message Date
brotoskyj 0aabbfd36e Style Change
Build Library / Build Library (push) Successful in 6m25s
Cleaned up services.rs file and removed whitespaces
Specified Data Types were needed instead of allowing the compiler to select types
2025-12-05 11:12:29 -05:00
brotoskyj b19eeb6c96 Fixed Progress Bar
Build Library / Build Library (push) Successful in 6m22s
Changed date math so that the last date from the first response from the API is the minimum value for the progress bar. This gives true progress percentages from the first date to the last date on the last request. #36
2025-12-04 11:57:49 -05:00
brotoskyj 76bd3a6087 Feature added - MPSC channel
Build Library / Build Library (push) Successful in 5m44s
closes #38
2025-12-03 17:54:57 -05:00
8 changed files with 170 additions and 121 deletions
+51 -3
View File
@@ -26,10 +26,13 @@ dependencies = [
[[package]]
name = "airlock_libs"
version = "5.0.1"
version = "5.1.2"
dependencies = [
"chrono",
"crossbeam",
"flexi_logger",
"indicatif",
"log",
"mongodb",
"opentelemetry 0.18.0",
"opentelemetry-otlp",
@@ -504,6 +507,19 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b"
[[package]]
name = "crossbeam"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8"
dependencies = [
"crossbeam-channel",
"crossbeam-deque",
"crossbeam-epoch",
"crossbeam-queue",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-channel"
version = "0.5.15"
@@ -513,6 +529,16 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.18"
@@ -522,6 +548,15 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-queue"
version = "0.3.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.21"
@@ -766,6 +801,19 @@ version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
[[package]]
name = "flexi_logger"
version = "0.31.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31e5335674a3a259527f97e9176a3767dcc9b220b8e29d643daeb2d6c72caf8b"
dependencies = [
"chrono",
"log",
"nu-ansi-term",
"regex",
"thiserror 2.0.17",
]
[[package]]
name = "fnv"
version = "1.0.7"
@@ -1557,9 +1605,9 @@ dependencies = [
[[package]]
name = "log"
version = "0.4.28"
version = "0.4.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432"
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
dependencies = [
"value-bag",
]
+4 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "airlock_libs"
version = "5.0.1"
version = "5.1.2"
edition = "2024"
[lib]
@@ -25,6 +25,9 @@ tracing = "0.1.41"
tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0"
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
crossbeam = "0.8.4"
log = "0.4.29"
flexi_logger = "0.31.7"
[package.metadata.maturin]
generate-abi-stubs = true
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project]
name = "airlock_libs"
version = "5.0.1"
version = "5.1.2"
description = "Airlock Digital API Wrapper"
readme = "README.md"
license = { text = "AGPL-3.0-only" }
+1 -1
View File
@@ -1,7 +1,7 @@
use pyo3::prelude::*;
pub mod modules;
pub mod services;
pub mod prelude;
pub mod services;
#[pymodule]
fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
+109 -111
View File
@@ -1,5 +1,7 @@
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<'_>,
@@ -12,12 +14,12 @@ pub fn pull_policy_exec_histories(
ExtractedValues::Headers(h) => h,
ExtractedValues::BaseUrl(_) => std::process::abort(),
};
let base_url = match PyData::convert(py, &py_self, false) {
let base_url: String = match PyData::convert(py, &py_self, false) {
ExtractedValues::Headers(_) => std::process::abort(),
ExtractedValues::BaseUrl(b) => b,
};
let handle = std::thread::spawn(move || {
let rt = match tokio::runtime::Runtime::new() {
let handle: thread::JoinHandle<String> = std::thread::spawn(move || {
let rt: tokio::runtime::Runtime = match tokio::runtime::Runtime::new() {
Ok(rt) => rt,
Err(e) => {
println!("Failed to build Tokio Runtime: {:?}", e);
@@ -27,14 +29,13 @@ pub fn pull_policy_exec_histories(
rt.block_on(async {
let _ = init_tracer();
});
let tracer = global::tracer("global_tracer");
let _cx = Context::new();
let tracer: global::BoxedTracer = global::tracer("global_tracer");
let _cx: Context = Context::new();
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()
@@ -55,13 +56,14 @@ pub fn pull_policy_exec_histories(
}
}
}
let data = ApiResponse {
let data: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
};
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
let writeable_filepath: PathBuf = file_path.clone();
let data_write: String = serde_json::to_string_pretty(&data).expect("Failed to serialize");
match fs::write(writeable_filepath.clone(), data_write) {
Ok(_) => {}
Err(e) => {
@@ -70,17 +72,17 @@ pub fn pull_policy_exec_histories(
}
}
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
let multi_progress = MultiProgress::new();
multi_progress.set_draw_target(ProgressDrawTarget::stdout());
let progress_bar = multi_progress.add(ProgressBar::new(100));
let multi_progress: MultiProgress = MultiProgress::new();
multi_progress.set_draw_target(ProgressDrawTarget::stderr());
let progress_bar: ProgressBar = 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));
let client = tracer.in_span("Building HTTP Client", |cx| {
let client_result = build_client(headers);
let client: Client = tracer.in_span("Building HTTP Client", |cx| {
let client_result: Result<Client, reqwest::Error> = build_client(headers);
match client_result {
Ok(client_result) => {
cx.span().add_event(
@@ -107,26 +109,79 @@ 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();
}
let cutoff: chrono::NaiveDateTime = Local::now().naive_local() - Duration::days(days);
let (tx, rx) = unbounded::<Vec<Group>>();
thread::spawn(move || {
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
{
let contents: String = 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: Group| {
(
(
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;
}
let history_date: NaiveDate = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
) {
Ok(date) => date,
Err(_) => continue,
};
if history_date >= cutoff.into() {
let key: (String, String, String) = (
executions.sha256.clone(),
executions.filename.clone(),
executions.hostname.clone(),
);
seen.entry(key).or_insert(executions.clone());
}
}
let final_response: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: seen.values().cloned().collect(),
},
};
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
match fs::write(&writeable_filepath, data_write) {
Ok(_) => {}
Err(e) => {
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
}
}
}
});
let mut first_date: Option<NaiveDate> = None;
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
let span: opentelemetry::trace::SpanRef<'_> = 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,
@@ -141,103 +196,46 @@ pub fn pull_policy_exec_histories(
));
results
});
let parsed_responses = execution_histories.response.exechistories;
let parsed_responses: Vec<Group> = 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() {
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",
) {
Ok(date) => date,
Err(_) => continue,
};
if history_date >= cutoff.into() {
let key = (
executions.sha256.clone(),
executions.filename.clone(),
executions.hostname.clone(),
);
seen.entry(key).or_insert(executions.clone());
}
}
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();
match fs::write(&writeable_filepath, data_write) {
Ok(_) => {}
Err(e) => {
println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
}
}
if let Some(last_item) = &final_response.response.exechistories.last()
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",
)
{
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);
progress_bar.set_message("Total Percent Complete");
if first_date.is_none() {
first_date = Some(last_date);
}
if let Some(base_date) = first_date {
let date_diff: chrono::TimeDelta = last_date - base_date;
let total_span: i64 = (Local::now().naive_local().date() - base_date).num_days();
let percentage: u64 = ((date_diff.num_days() as f64 / total_span as f64) * 100.0)
.clamp(0.0, 100.0)
.round() as u64;
progress_bar.set_position(percentage);
}
}
}
});
progress_bar.finish_with_message("All Checkpoints Complete");
let return_data = match fs::read_to_string(&writeable_filepath) {
let return_data: String = 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()
});
let gil_value = handle.join().unwrap();
Python::attach(|py| PyString::new(py, &gil_value).into())
let gil_value: String = handle.join().unwrap();
Python::attach(|py: Python<'_>| PyString::new(py, &gil_value).into())
}
fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
@@ -264,7 +262,7 @@ async fn history_logging(
}}"#,
exec_types, checkpoint_number, policy_names
);
let res = client
let res: Result<reqwest::Response, reqwest::Error> = client
.post(format!("{}/v1/logging/exechistories", base_url))
.body(payload)
.send()
@@ -305,13 +303,13 @@ pub fn get_base_directory() -> PathBuf {
}
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
let cfg = TelemetryConfig::load();
let cfg: TelemetryConfig = TelemetryConfig::load();
if !cfg.TELEMETRY {
global::set_tracer_provider(NoopTracerProvider::new());
return Ok(None);
}
let endpoint = cfg.TELEM_URL.unwrap_or_default();
let tracer =
let endpoint: String = cfg.TELEM_URL.unwrap_or_default();
let tracer: sdktrace::Tracer =
opentelemetry_otlp::new_pipeline()
.tracing()
.with_exporter(
+1 -1
View File
@@ -11,4 +11,4 @@ urllib3==2.5.0
pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==5.0.1
airlock_libs==5.1.2