RustImplementation #49

Merged
mysticmomba merged 33 commits from RustImplementation into master 2025-12-17 14:19:20 -05:00
5 changed files with 784 additions and 513 deletions
Showing only changes of commit f3c1d97d28 - Show all commits
+536 -297
View File
File diff suppressed because it is too large Load Diff
+2 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "airlock_libs" name = "airlock_libs"
version = "4.0.3" version = "5.0.0"
edition = "2024" edition = "2024"
[lib] [lib]
@@ -24,6 +24,7 @@ tonic = { version = "0.8.2", features = ["tls-roots"] }
tracing = "0.1.41" tracing = "0.1.41"
tracing-subscriber = "0.3.20" tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0" tracing-opentelemetry = "0.32.0"
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
[package.metadata.maturin] [package.metadata.maturin]
generate-abi-stubs = true generate-abi-stubs = true
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project] [project]
name = "airlock_libs" name = "airlock_libs"
version = "4.0.3" version = "5.0.0"
description = "Airlock Digital API Wrapper" description = "Airlock Digital API Wrapper"
readme = "README.md" readme = "README.md"
license = { text = "AGPL-3.0-only" } license = { text = "AGPL-3.0-only" }
+57 -26
View File
@@ -9,6 +9,7 @@ use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
use opentelemetry::{Key, global}; use opentelemetry::{Key, global};
use opentelemetry_otlp::WithExportConfig; use opentelemetry_otlp::WithExportConfig;
use pyo3::{prelude::*, types::PyString}; use pyo3::{prelude::*, types::PyString};
use pyo3_async_runtimes::async_std;
use reqwest::{ use reqwest::{
Client, Client,
header::{HeaderMap, HeaderName, HeaderValue}, header::{HeaderMap, HeaderName, HeaderValue},
@@ -88,6 +89,40 @@ struct Group {
localip: String, 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] #[pyfunction]
pub fn pull_policy_exec_histories( pub fn pull_policy_exec_histories(
py: Python<'_>, py: Python<'_>,
@@ -96,6 +131,15 @@ pub fn pull_policy_exec_histories(
exec_types: String, exec_types: String,
days: i64, days: i64,
) -> Py<PyString> { ) -> Py<PyString> {
let headers: HeaderMap = match PyData::convert(py, &py_self, true) {
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() { let rt = match tokio::runtime::Runtime::new() {
Ok(rt) => rt, Ok(rt) => rt,
Err(e) => { Err(e) => {
@@ -159,7 +203,7 @@ pub fn pull_policy_exec_histories(
); );
progress_bar.enable_steady_tick(std::time::Duration::from_millis(100)); progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
let client = tracer.in_span("Building HTTP Client", |cx| { let client = tracer.in_span("Building HTTP Client", |cx| {
let client_result = build_client(py, &py_self); let client_result = build_client(headers);
match client_result { match client_result {
Ok(client_result) => { Ok(client_result) => {
cx.span().add_event( cx.span().add_event(
@@ -186,7 +230,6 @@ pub fn pull_policy_exec_histories(
} }
} }
}); });
let api: Py<PyAny> = py_self;
let cutoff = Local::now().naive_local() - Duration::days(days); let cutoff = Local::now().naive_local() - Duration::days(days);
let mut f = match File::open(&writeable_filepath) { let mut f = match File::open(&writeable_filepath) {
Ok(f) => f, Ok(f) => f,
@@ -209,8 +252,7 @@ pub fn pull_policy_exec_histories(
} }
let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| { let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
let results: ApiResponse = history_logging( let results: ApiResponse = history_logging(
py, &base_url,
&api,
&exec_types, &exec_types,
&checkpoint_number, &checkpoint_number,
&policy_names, &policy_names,
@@ -226,8 +268,8 @@ pub fn pull_policy_exec_histories(
if parsed_responses.is_empty() { if parsed_responses.is_empty() {
break; break;
} }
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() let mut seen: HashMap<(String, String, String), Group> =
{ if writeable_filepath.exists() {
let mut contents = String::new(); let mut contents = String::new();
f.read_to_string(&mut contents).unwrap(); f.read_to_string(&mut contents).unwrap();
let existing_data: ApiResponse = let existing_data: ApiResponse =
@@ -299,7 +341,8 @@ pub fn pull_policy_exec_histories(
) )
{ {
let date_diff = Local::now().naive_local().date() - last_date; 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; 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_position(percentage_diff.round() as u64);
progress_bar.set_message("Total Percent Complete"); progress_bar.set_message("Total Percent Complete");
} }
@@ -314,40 +357,28 @@ pub fn pull_policy_exec_histories(
} }
}; };
shutdown_tracer_provider(); shutdown_tracer_provider();
PyString::new(py, &return_data).into() return_data.to_string()
});
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> { 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() Client::builder()
.danger_accept_invalid_certs(true) .danger_accept_invalid_certs(true)
.default_headers(header_map) .default_headers(headers)
.timeout(std::time::Duration::from_secs(300)) .timeout(std::time::Duration::from_secs(300))
.build() .build()
} }
#[tokio::main] #[tokio::main]
async fn history_logging( async fn history_logging(
py: Python<'_>, base_url: &String,
py_self: &Py<PyAny>,
exec_types: &String, exec_types: &String,
checkpoint_number: &String, checkpoint_number: &String,
policy_names: &String, policy_names: &String,
client: &Client, client: &Client,
) -> ApiResponse { ) -> ApiResponse {
let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
let payload = format!( let payload = format!(
r#"{{ r#"{{
"type": {}, "type": {},
+1 -1
View File
@@ -11,4 +11,4 @@ urllib3==2.5.0
pyperclip==1.11.0 pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/ --extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==4.0.3 airlock_libs==5.0.0