Merge pull request 'RustImplementation' (#28) from RustImplementation into master

Reviewed-on: brotoskyj/AirlockTools#28

airlock_libs/.gitignore

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

airlock_libs/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "airlock_libs"
-version = "5.0.0"
+version = "3.1.2"
 edition = "2024"
 
 [lib]
@@ -24,7 +24,6 @@ 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 = "5.0.0"
+version = "3.1.2"
 description = "Airlock Digital API Wrapper"
 readme = "README.md"
 license = { text = "AGPL-3.0-only" }

airlock_libs/src/services.rs

@@ -9,7 +9,6 @@ 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},
@@ -26,34 +25,12 @@ 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,
@@ -89,40 +66,6 @@ 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<'_>,
@@ -131,22 +74,7 @@ pub fn pull_policy_exec_histories(
     exec_types: String,
     days: i64,
 ) -> Py<PyString> {
-    let headers: HeaderMap = match PyData::convert(py, &py_self, true) {
+    let rt = tokio::runtime::Runtime::new().unwrap();
-        ExtractedValues::Headers(h) => h,
-        ExtractedValues::BaseUrl(_) => std::process::abort(),
-    };
-    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();
-            }
-        };
     rt.block_on(async {
         let _ = init_tracer();
     });
@@ -158,25 +86,13 @@ pub fn pull_policy_exec_histories(
     )
     .into();
     let writeable_filepath = file_path.clone();
-        if !&file_path.exists() {
+    if !file_path.exists() {
-            if let Some(parent_dir) = &file_path.parent()
+        if let Some(parent_dir) = file_path.parent()
             && !parent_dir.exists()
         {
-                match fs::create_dir_all(parent_dir) {
+            fs::create_dir_all(parent_dir).unwrap();
-                    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();
-                }
         }
+        fs::File::create(file_path).unwrap();
     }
     let data = ApiResponse {
         error: "Success".to_string(),
@@ -185,13 +101,7 @@ pub fn pull_policy_exec_histories(
         },
     };
     let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
-        match fs::write(writeable_filepath.clone(), data_write) {
+    fs::write(writeable_filepath.clone(), data_write).unwrap();
-            Ok(_) => {}
-            Err(e) => {
-                println!("Failed to write to: {:?}: {}", &writeable_filepath, e);
-                std::process::abort();
-            }
-        }
     let mut checkpoint_number: String = skipback(days).to_string();
     let multi_progress = MultiProgress::new();
     multi_progress.set_draw_target(ProgressDrawTarget::stdout());
@@ -203,7 +113,7 @@ pub fn pull_policy_exec_histories(
     );
     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_result = build_client(py, &py_self);
         match client_result {
             Ok(client_result) => {
                 cx.span().add_event(
@@ -225,34 +135,22 @@ pub fn pull_policy_exec_histories(
                 );
                 cx.span()
                     .set_status(Status::error("Client Failed to Build"));
-                    println!("Failed to Build Client: {:?}", client_result);
+                panic!("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) {
+    let mut f = File::open(&writeable_filepath).unwrap();
-            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(Key::new("Days").string(days.to_string().to_string()));
-            span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
         loop {
-                match f.seek(SeekFrom::Start(0)) {
+            f.seek(SeekFrom::Start(0)).unwrap();
-                    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,
+                    py,
+                    &api,
                     &exec_types,
                     &checkpoint_number,
                     &policy_names,
@@ -268,8 +166,8 @@ pub fn pull_policy_exec_histories(
             if parsed_responses.is_empty() {
                 break;
             }
-                let mut seen: HashMap<(String, String, String), Group> =
+            let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
-                    if writeable_filepath.exists() {
+            {
                 let mut contents = String::new();
                 f.read_to_string(&mut contents).unwrap();
                 let existing_data: ApiResponse =
@@ -328,12 +226,7 @@ pub fn pull_policy_exec_histories(
                 },
             };
             let data_write = serde_json::to_string_pretty(&final_response).unwrap();
-                match fs::write(&writeable_filepath, data_write) {
+            fs::write(&writeable_filepath, data_write).unwrap();
-                    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", ""),
@@ -341,44 +234,49 @@ pub fn pull_policy_exec_histories(
                 )
             {
                 let date_diff = Local::now().naive_local().date() - last_date;
-                    let percentage_diff =
+                let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
-                        (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) {
+    let return_data = fs::read_to_string(&writeable_filepath).unwrap();
-            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()
+    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(headers)
+        .default_headers(header_map)
         .timeout(std::time::Duration::from_secs(300))
         .build()
 }
 
 #[tokio::main]
 async fn history_logging(
-    base_url: &String,
+    py: Python<'_>,
+    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": {},
@@ -436,8 +334,30 @@ 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 = TelemetryConfig::load();
+    let cfg = load_telemetry_config();
     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==5.0.0
+airlock_libs==3.1.2