use crate::modules::datatypes::*; use crate::prelude::*; use opentelemetry::trace::SpanContext; #[pyfunction] pub fn pull_policy_exec_histories( py: Python<'_>, py_self: Py, policy_names: Option, exec_types: String, days: i64, ) -> Py { println!(); let data: PyData = PyData::extract_data(py, &py_self); let headers: HeaderMap = data.headers; let base_url: String = data.base_url; let handle: thread::JoinHandle = 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); std::process::abort(); } }; let tracer_provider = rt.block_on(async { TelemetryConfig::init_tracer() }); global::set_tracer_provider(tracer_provider.clone()); let tracer: global::BoxedTracer = global::tracer("tracer"); let _cx: Context = Context::new(); let file_path: PathBuf = format!( "{}\\cache\\chunkinator.json", get_base_directory().display() ) .into(); 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 data: ApiResponse = ApiResponse { error: "Success".to_string(), response: ExecHistories { exechistories: vec![], }, }; 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) => { println!("Failed to write to: {:?}: {}", &writeable_filepath, e); std::process::abort(); } } let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string(); let progress_bar = Arc::new(Mutex::new(ProgressBar::new(100))); progress_bar .lock() .unwrap() .set_draw_target(ProgressDrawTarget::stderr()); progress_bar.lock().unwrap().set_style( ProgressStyle::default_bar() .template("Total - Policy Name: {msg}: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len}") .unwrap().progress_chars("⣿⣦⣀") ); let client: Client = tracer.in_span("Building HTTP Client", |cx| { let client_result: 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 } 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: chrono::NaiveDateTime = Local::now().naive_local() - chrono::Duration::days(days); let (tx, rx) = unbounded::<(SpanContext, Vec)>(); let pb_clone = progress_bar.clone(); thread::spawn(move || { let tracer = global::tracer("loxide"); 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((parent_spancontext, parsed_responses)) = rx.recv() { let parent_ctx = Context::new().with_remote_span_context(parent_spancontext); let span = tracer.build_with_context( tracer .span_builder("Deduplicate and Write") .with_kind(trace::SpanKind::Consumer), &parent_ctx, ); let cx = Context::current_with_span(span); cx.span().add_event( "Received Data from Producer", vec![KeyValue::new( "Items to Process", parsed_responses.len().to_string(), )], ); 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(_) => { cx.span().add_event( "Writing Data to File", vec![KeyValue::new("Success", "Ok".to_string())], ); } Err(e) => { cx.span().add_event( "Writing Data to File", vec![KeyValue::new("Failed", e.to_string())], ); cx.span() .set_status(Status::error("Failed to Write to File")); } } cx.span().add_event( "Finished Deduplicating Data", vec![KeyValue::new( "Items Successfully Processed", seen.len().to_string(), )], ); } }); let mut first_date: Option = None; tracer.in_span("Airlock Data Retreival", |cx| { pb_clone .lock() .unwrap() .enable_steady_tick(std::time::Duration::from_millis(100)); pb_clone .lock() .unwrap() .set_message(policy_names.clone().unwrap_or("Statistics".to_string())); let span: opentelemetry::trace::SpanRef<'_> = cx.span(); span.set_attribute(KeyValue::new("Days", days.to_string())); span.set_attribute(KeyValue::new( "Policy Name", policy_names .clone() .unwrap_or("Statistics Monitoring".to_string()), )); loop { let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| { cx.span().add_event( "Retrieving Responses from API", vec![KeyValue::new( "Checkpoint Number", checkpoint_number.to_string(), )], ); let results: ApiResponse = rt.block_on(history_logging( &base_url, &exec_types, &checkpoint_number, &policy_names, &client, )); cx.span().add_event( "Got Responses from API", vec![KeyValue::new( "Items in Response", results.response.exechistories.len().to_string(), )], ); cx.span().set_status(Status::Ok); cx.span().set_attribute(KeyValue::new( "items_in_response", results.response.exechistories.len().to_string(), )); results }); let parsed_responses: Vec = execution_histories.response.exechistories; if parsed_responses.is_empty() { break; } match tx.send((cx.span().span_context().clone(), parsed_responses.clone())) { Ok(_) => {} Err(e) => { cx.span().add_event( "Failed to Send Items to Processor", vec![KeyValue::new("Response from Processor", e.to_string())], ); cx.span().set_status(Status::error("Processor Failed")) } } 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", ) { 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; pb_clone.lock().unwrap().set_position(percentage); } } } }); progress_bar .lock() .unwrap() .finish_with_message("All Checkpoints Complete"); 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: {:?}: {}", &file_path, e); std::process::abort(); } }; tracer_provider .shutdown() .expect("Failed to Shutdown Tracer Provdier"); drop(tx); return_data.to_string() }); let gil_value: String = handle.join().unwrap(); Python::attach(|py: Python<'_>| PyString::new(py, &gil_value).into()) } fn build_client(headers: HeaderMap) -> Result { Client::builder() .danger_accept_invalid_certs(true) .default_headers(headers) .timeout(std::time::Duration::from_secs(300)) .build() } #[tracing::instrument(name = "history_logging")] async fn history_logging( base_url: &String, exec_types: &String, checkpoint_number: &String, policy_names: &Option, client: &Client, ) -> ApiResponse { let policy_json = match policy_names { Some(name) => format!(r#"[ "{}" ]"#, name), // JSON array with one element None => "[]".to_string(), // Empty JSON array }; let payload = format!( r#"{{ "type": {}, "checkpoint": "{}", "policy": {} }}"#, exec_types, checkpoint_number, policy_json ); let res: Result = client .post(format!("{}/v1/logging/exechistories", base_url)) .body(payload) .send() .await; match res { Ok(res) => { let first_response: ApiResponse = serde_json::from_str(&res.text().await.unwrap()) .expect("Failed to retrieve response from API"); first_response } Err(_res) => { let failed_response: ApiResponse = ApiResponse { error: "Failed".to_string(), response: ExecHistories { exechistories: vec![], }, }; failed_response } } } pub fn get_base_directory() -> PathBuf { let home = env::var_os("HOME") .map(PathBuf::from) .or_else(|| env::var_os("USERPROFILE").map(PathBuf::from)) .expect("Could not find Home Directory"); let os = std::env::consts::OS; match os { "windows" => { let appdata = env::var_os("APPDATA") .map(PathBuf::from) .unwrap_or_else(|| home.join("AppData").join("Roaming")); appdata.join("Loxide") } "linux" => home.join(".local").join("share").join("Loxide"), _ => { println!("{} is currently not compatible with LoxideLibs", os); std::process::abort(); } } }