airlock_libs/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "airlock_libs"
-version = "4.0.3"
+version = "5.0.0"
 edition = "2024"
 
 [lib]
@@ -24,6 +24,7 @@ tonic = { version = "0.8.2", features = ["tls-roots"] }
 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"] }
 
 [package.metadata.maturin]
 generate-abi-stubs = true

airlock_libs/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "airlock_libs"
-version = "4.0.3"
+version = "5.0.0"
 description = "Airlock Digital API Wrapper"
 readme = "README.md"
 license = { text = "AGPL-3.0-only" }

airlock_libs/src/services.rs

@@ -9,6 +9,7 @@ use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
 use opentelemetry::{Key, global};
 use opentelemetry_otlp::WithExportConfig;
 use pyo3::{prelude::*, types::PyString};
+use pyo3_async_runtimes::async_std;
 use reqwest::{
     Client,
     header::{HeaderMap, HeaderName, HeaderValue},
@@ -88,6 +89,40 @@ struct Group {
     localip: String,
 }
 
+enum ExtractedValues {
+    Headers(reqwest::header::HeaderMap),
+    BaseUrl(String),
+}
+
+trait Converter {
+    fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues;
+}
+
+struct PyData;
+
+impl Converter for PyData {
+    fn convert(py: Python<'_>, py_self: &Py<PyAny>, extract_headers: bool) -> ExtractedValues {
+        if extract_headers {
+            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);
+                    }
+                }
+            }
+            ExtractedValues::Headers(header_map)
+        } else {
+            let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
+            ExtractedValues::BaseUrl(base_url)
+        }
+    }
+}
+
 #[pyfunction]
 pub fn pull_policy_exec_histories(
     py: Python<'_>,
@@ -96,258 +131,254 @@ pub fn pull_policy_exec_histories(
     exec_types: String,
     days: i64,
 ) -> Py<PyString> {
-    let rt = match tokio::runtime::Runtime::new() {
+    let headers: HeaderMap = match PyData::convert(py, &py_self, true) {
-        Ok(rt) => rt,
+        ExtractedValues::Headers(h) => h,
-        Err(e) => {
+        ExtractedValues::BaseUrl(_) => std::process::abort(),
-            println!("Failed to build Tokio Runtime: {:?}", e);
-            std::process::abort();
-        }
     };
-    rt.block_on(async {
+    let base_url = match PyData::convert(py, &py_self, false) {
-        let _ = init_tracer();
+        ExtractedValues::Headers(_) => std::process::abort(),
-    });
+        ExtractedValues::BaseUrl(b) => b,
-    let tracer = global::tracer("global_tracer");
+    };
-    let _cx = Context::new();
+    let handle = std::thread::spawn(move || {
-    let file_path: PathBuf = format!(
+        let rt = match tokio::runtime::Runtime::new() {
-        "{}\\cache\\chunkinator.json",
+            Ok(rt) => rt,
-        get_base_directory().display()
+            Err(e) => {
-    )
+                println!("Failed to build Tokio Runtime: {:?}", e);
-    .into();
+                std::process::abort();
-    let writeable_filepath = file_path.clone();
+            }
-    if !&file_path.exists() {
+        };
-        if let Some(parent_dir) = &file_path.parent()
+        rt.block_on(async {
-            && !parent_dir.exists()
+            let _ = init_tracer();
-        {
+        });
-            match fs::create_dir_all(parent_dir) {
+        let tracer = global::tracer("global_tracer");
+        let _cx = 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()
+            {
+                match fs::create_dir_all(parent_dir) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        println!("Failed to Create Directory {:?}: {}", parent_dir, e);
+                        std::process::abort();
+                    }
+                }
+            }
+            match fs::File::create(&file_path) {
                 Ok(_) => {}
                 Err(e) => {
-                    println!("Failed to Create Directory {:?}: {}", parent_dir, e);
+                    println!("Failed to Create Directory {:?}: {}", &file_path, e);
                     std::process::abort();
                 }
             }
         }
-        match fs::File::create(&file_path) {
+        let data = ApiResponse {
+            error: "Success".to_string(),
+            response: ExecHistories {
+                exechistories: vec![],
+            },
+        };
+        let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
+        match fs::write(writeable_filepath.clone(), data_write) {
             Ok(_) => {}
             Err(e) => {
-                println!("Failed to Create Directory {:?}: {}", &file_path, e);
+                println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
                 std::process::abort();
             }
         }
-    }
+        let mut checkpoint_number: String = skipback(days).to_string();
-    let data = ApiResponse {
+        let multi_progress = MultiProgress::new();
-        error: "Success".to_string(),
+        multi_progress.set_draw_target(ProgressDrawTarget::stdout());
-        response: ExecHistories {
+        let progress_bar = multi_progress.add(ProgressBar::new(100));
-            exechistories: vec![],
+        progress_bar.set_style(
-        },
+            ProgressStyle::default_bar()
-    };
+                .template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
-    let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
+                .unwrap(),
-    match fs::write(writeable_filepath.clone(), data_write) {
+        );
-        Ok(_) => {}
+        progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
-        Err(e) => {
+        let client = tracer.in_span("Building HTTP Client", |cx| {
-            println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
+            let client_result = build_client(headers);
-            std::process::abort();
+            match client_result {
-        }
+                Ok(client_result) => {
-    }
+                    cx.span().add_event(
-    let mut checkpoint_number: String = skipback(days).to_string();
+                        "info",
-    let multi_progress = MultiProgress::new();
+                        vec![KeyValue::new(
-    multi_progress.set_draw_target(ProgressDrawTarget::stdout());
+                            "Client Built Successfully",
-    let progress_bar = multi_progress.add(ProgressBar::new(100));
+                            format!("{:?}", client_result),
-    progress_bar.set_style(
+                        )],
-        ProgressStyle::default_bar()
+                    );
-            .template("Total Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}")
+                    client_result
-            .unwrap(),
+                }
-    );
+                Err(client_result) => {
-    progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
+                    cx.span().add_event(
-    let client = tracer.in_span("Building HTTP Client", |cx| {
+                        "warn",
-        let client_result = build_client(py, &py_self);
+                        vec![KeyValue::new(
-        match client_result {
+                            "Client Failed to Build",
-            Ok(client_result) => {
+                            format!("{:?}", &client_result),
-                cx.span().add_event(
+                        )],
-                    "info",
+                    );
-                    vec![KeyValue::new(
+                    cx.span()
-                        "Client Built Successfully",
+                        .set_status(Status::error("Client Failed to Build"));
-                        format!("{:?}", client_result),
+                    println!("Failed to Build Client: {:?}", client_result);
-                    )],
-                );
-                client_result
-            }
-            Err(client_result) => {
-                cx.span().add_event(
-                    "warn",
-                    vec![KeyValue::new(
-                        "Client Failed to Build",
-                        format!("{:?}", &client_result),
-                    )],
-                );
-                cx.span()
-                    .set_status(Status::error("Client Failed to Build"));
-                println!("Failed to Build Client: {:?}", client_result);
-                std::process::abort();
-            }
-        }
-    });
-    let api: Py<PyAny> = py_self;
-    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(
+        let cutoff = Local::now().naive_local() - Duration::days(days);
-                    py,
+        let mut f = match File::open(&writeable_filepath) {
-                    &api,
+            Ok(f) => f,
-                    &exec_types,
+            Err(e) => {
-                    &checkpoint_number,
+                println!("Failed to Access {:?}: {}", &writeable_filepath, e);
-                    &policy_names,
+                std::process::abort();
-                    &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()
+        };
-            {
+        tracer.in_span("Airlock Data Retreival", |cx| {
-                let mut contents = String::new();
+            let span = cx.span();
-                f.read_to_string(&mut contents).unwrap();
+            span.set_attribute(Key::new("Days").string(days.to_string()));
-                let existing_data: ApiResponse =
+            span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
-                    serde_json::from_str(&contents).unwrap_or(ApiResponse {
+            loop {
-                        error: "Success".to_string(),
+                match f.seek(SeekFrom::Start(0)) {
-                        response: ExecHistories {
+                    Ok(_) => {}
-                            exechistories: vec![],
+                    Err(e) => {
-                        },
+                        println!("Failed to seek start of {:?}: {}", f, e);
-                    });
+                        std::process::abort();
-                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 {
+                let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
-                    checkpoint_number = executions.checkpoint.clone();
+                    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 history_date = match NaiveDate::parse_from_str(
+                let mut seen: HashMap<(String, String, String), Group> =
-                    &executions.datetime.replace(" +0000 UTC", ""),
+                    if writeable_filepath.exists() {
-                    "%Y-%m-%dT%H:%M:%SZ",
+                        let mut contents = String::new();
-                ) {
+                        f.read_to_string(&mut contents).unwrap();
-                    Ok(date) => date,
+                        let existing_data: ApiResponse =
-                    Err(_) => continue,
+                            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(),
+                    },
                 };
-                if history_date >= cutoff.into() {
+                let data_write = serde_json::to_string_pretty(&final_response).unwrap();
-                    let key = (
+                match fs::write(&writeable_filepath, data_write) {
-                        executions.sha256.clone(),
+                    Ok(_) => {}
-                        executions.filename.clone(),
+                    Err(e) => {
-                        executions.hostname.clone(),
+                        println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
-                    );
+                    }
-                    seen.entry(key).or_insert(executions.clone());
+                }
+                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 - 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");
                 }
             }
-            let final_response = ApiResponse {
+        });
-                error: "Success".to_string(),
+        progress_bar.finish_with_message("All Checkpoints Complete");
-                response: ExecHistories {
+        let return_data = match fs::read_to_string(&writeable_filepath) {
-                    exechistories: seen.values().cloned().collect(),
+            Ok(return_data) => return_data,
-                },
+            Err(e) => {
-            };
+                println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
-            let data_write = serde_json::to_string_pretty(&final_response).unwrap();
+                std::process::abort();
-            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()
+        };
-                && let Ok(last_date) = NaiveDate::parse_from_str(
+        shutdown_tracer_provider();
-                    &last_item.datetime.replace(" +0000 UTC", ""),
+        return_data.to_string()
-                    "%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");
-            }
-        }
     });
-    progress_bar.finish_with_message("All Checkpoints Complete");
+    let gil_value = handle.join().unwrap();
-    let return_data = match fs::read_to_string(&writeable_filepath) {
+    Python::attach(|py| PyString::new(py, &gil_value).into())
-        Ok(return_data) => return_data,
-        Err(e) => {
-            println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
-            std::process::abort();
-        }
-    };
-    shutdown_tracer_provider();
-    PyString::new(py, &return_data).into()
 }
 
-fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Result<reqwest::Client, reqwest::Error> {
+fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
-    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);
-            }
-        }
-    }
-
     Client::builder()
         .danger_accept_invalid_certs(true)
-        .default_headers(header_map)
+        .default_headers(headers)
         .timeout(std::time::Duration::from_secs(300))
         .build()
 }
 
 #[tokio::main]
 async fn history_logging(
-    py: Python<'_>,
+    base_url: &String,
-    py_self: &Py<PyAny>,
     exec_types: &String,
     checkpoint_number: &String,
     policy_names: &String,
     client: &Client,
 ) -> ApiResponse {
-    let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
     let payload = format!(
         r#"{{
         "type": {},

requirements.txt

@@ -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==4.0.3
+airlock_libs==5.0.0