Implemented Threaded Rust Process

It still locks the GIL, but this is the ground work for the non blocking GUI. Python objects are converted to Rust objects in the beginning of the function call so they can be sent to threads safely.

airlock_libs/.gitignore

@@ -1,3 +1,4 @@
 /target
 build.sh
-pythontest.py
+pythontest.py
+changelog.md

airlock_libs/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "airlock_libs"
-version = "3.1.2"
+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 = "3.1.2"
+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},
@@ -25,12 +26,34 @@ use std::{
     str::FromStr,
 };
 
+#[allow(non_snake_case)]
 #[derive(Deserialize, Debug)]
 struct TelemetryConfig {
     TELEMETRY: bool,
     TELEM_URL: Option<String>,
 }
 
+impl TelemetryConfig {
+    pub fn load() -> Self {
+        let cfg_path = get_base_directory().join("config\\user_config.json");
+        if !cfg_path.exists() {
+            return Self {
+                TELEMETRY: false,
+                TELEM_URL: None,
+            };
+        }
+        match fs::read_to_string(&cfg_path) {
+            Ok(contents) => serde_json::from_str::<Self>(&contents).unwrap_or(Self {
+                TELEMETRY: false,
+                TELEM_URL: None,
+            }),
+            Err(_) => Self {
+                TELEMETRY: false,
+                TELEM_URL: None,
+            },
+        }
+    }
+}
 #[derive(Debug, Deserialize, Serialize)]
 struct ApiResponse {
     error: String,
@@ -66,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<'_>,
@@ -74,209 +131,254 @@ pub fn pull_policy_exec_histories(
     exec_types: String,
     days: i64,
 ) -> Py<PyString> {
-    let rt = tokio::runtime::Runtime::new().unwrap();
-    rt.block_on(async {
-        let _ = init_tracer();
-    });
-    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()
-        {
-            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 headers: HeaderMap = match PyData::convert(py, &py_self, true) {
+        ExtractedValues::Headers(h) => h,
+        ExtractedValues::BaseUrl(_) => std::process::abort(),
     };
-    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 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 client = tracer.in_span("Building HTTP Client", |cx| {
-        let client_result = build_client(py, &py_self);
-        match client_result {
-            Ok(client_result) => {
-                cx.span().add_event(
-                    "info",
-                    vec![KeyValue::new(
-                        "Client Built Successfully",
-                        format!("{:?}", client_result),
-                    )],
-                );
-                client_result
+    let base_url = 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() {
+            Ok(rt) => rt,
+            Err(e) => {
+                println!("Failed to build Tokio Runtime: {:?}", e);
+                std::process::abort();
             }
-            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"));
-                panic!("Failed to Build Client: {:?}", client_result);
+        };
+        rt.block_on(async {
+            let _ = init_tracer();
+        });
+        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 {:?}: {}", &file_path, e);
+                    std::process::abort();
+                }
             }
         }
-    });
-    let api: Py<PyAny> = py_self;
-    let cutoff = Local::now().naive_local() - Duration::days(days);
-    let mut f = File::open(&writeable_filepath).unwrap();
-    tracer.in_span("Airlock Data Retreival", |cx| {
-        let span = cx.span();
-        span.set_attribute(Key::new("Days").string(days.to_string().to_string()));
-        loop {
-            f.seek(SeekFrom::Start(0)).unwrap();
-            let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
-                let results: ApiResponse = history_logging(
-                    py,
-                    &api,
-                    &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 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 write to: {:?}: {}", &writeable_filepath, e);
+                std::process::abort();
             }
-            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;
+        }
+        let mut checkpoint_number: String = skipback(days).to_string();
+        let multi_progress = MultiProgress::new();
+        multi_progress.set_draw_target(ProgressDrawTarget::stdout());
+        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 client = tracer.in_span("Building HTTP Client", |cx| {
+            let client_result = build_client(headers);
+            match client_result {
+                Ok(client_result) => {
+                    cx.span().add_event(
+                        "info",
+                        vec![KeyValue::new(
+                            "Client Built Successfully",
+                            format!("{:?}", client_result),
+                        )],
+                    );
+                    client_result
                 }
-                if index == parsed_responses.len() - 1 {
-                    checkpoint_number = executions.checkpoint.clone();
+                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 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 history_date = match NaiveDate::parse_from_str(
-                    &executions.datetime.replace(" +0000 UTC", ""),
-                    "%Y-%m-%dT%H:%M:%SZ",
-                ) {
-                    Ok(date) => date,
-                    Err(_) => continue,
+                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(),
+                    },
                 };
-                if history_date >= cutoff.into() {
-                    let key = (
-                        executions.sha256.clone(),
-                        executions.filename.clone(),
-                        executions.hostname.clone(),
-                    );
-                    seen.entry(key).or_insert(executions.clone());
+                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()
+                    && 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(),
-                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 - 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 return_data = match fs::read_to_string(&writeable_filepath) {
+            Ok(return_data) => return_data,
+            Err(e) => {
+                println!("Failed to read data from: {:?}: {}", &writeable_filepath, e);
+                std::process::abort();
             }
-        }
+        };
+        shutdown_tracer_provider();
+        return_data.to_string()
     });
-    progress_bar.finish_with_message("All Checkpoints Complete");
-    let return_data = fs::read_to_string(&writeable_filepath).unwrap();
-    shutdown_tracer_provider();
-    PyString::new(py, &return_data).into()
+    let gil_value = handle.join().unwrap();
+    Python::attach(|py| PyString::new(py, &gil_value).into())
 }
 
-fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> 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);
-            }
-        }
-    }
-
+fn build_client(headers: HeaderMap) -> Result<reqwest::Client, reqwest::Error> {
     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<'_>,
-    py_self: &Py<PyAny>,
+    base_url: &String,
     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": {},
@@ -334,30 +436,8 @@ fn skipback(days: i64) -> ObjectId {
     ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
 }
 
-fn load_telemetry_config() -> TelemetryConfig {
-    let cfg_path = get_base_directory().join("config\\user_config.json");
-    if !cfg_path.exists() {
-        return TelemetryConfig {
-            TELEMETRY: false,
-            TELEM_URL: None,
-        };
-    }
-    match fs::read_to_string(&cfg_path) {
-        Ok(contents) => {
-            serde_json::from_str::<TelemetryConfig>(&contents).unwrap_or(TelemetryConfig {
-                TELEMETRY: false,
-                TELEM_URL: None,
-            })
-        }
-        Err(_) => TelemetryConfig {
-            TELEMETRY: false,
-            TELEM_URL: None,
-        },
-    }
-}
-
 fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
-    let cfg = load_telemetry_config();
+    let cfg = TelemetryConfig::load();
     if !cfg.TELEMETRY {
         global::set_tracer_provider(NoopTracerProvider::new());
         return Ok(None);

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