use chrono::{Duration, Local, NaiveDate}; use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; use mongodb::bson::oid::ObjectId; use opentelemetry::global::shutdown_tracer_provider; use opentelemetry::sdk::Resource; use opentelemetry::trace::noop::NoopTracerProvider; use opentelemetry::trace::{Status, TraceContextExt, TraceError}; use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer}; use opentelemetry::{Key, global}; use opentelemetry_otlp::WithExportConfig; use pyo3::{prelude::*, types::PyString}; use reqwest::{ Client, header::{HeaderMap, HeaderName, HeaderValue}, }; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::{ collections::HashMap, env, fmt::Write, fs::{self, File}, io::{Read, Seek, SeekFrom}, path::PathBuf, str::FromStr, }; #[allow(non_snake_case)] #[derive(Deserialize, Debug)] struct TelemetryConfig { TELEMETRY: bool, TELEM_URL: Option, } 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::(&contents).unwrap_or(Self { TELEMETRY: false, TELEM_URL: None, }), Err(_) => Self { TELEMETRY: false, TELEM_URL: None, }, } } } #[derive(Debug, Deserialize, Serialize)] struct ApiResponse { error: String, response: ExecHistories, } #[derive(Debug, Deserialize, Serialize)] struct ExecHistories { exechistories: Vec, } #[derive(Debug, Deserialize, Serialize, Clone)] struct Group { checkpoint: String, #[serde(rename = "type")] exectype: u8, username: String, hostname: String, netdomain: String, filename: String, ppolicy: String, policyname: String, policyver: String, commandline: String, publisher: String, pprocess: String, gprocess: String, sha256: String, datetime: String, md5: String, sha128: String, sha384: String, sha512: String, ip: String, localip: String, } #[pyfunction] pub fn pull_policy_exec_histories( py: Python<'_>, py_self: Py, policy_names: String, exec_types: String, days: i64, ) -> Py { 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() { 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 { error: "Success".to_string(), response: ExecHistories { exechistories: vec![], }, }; 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 } 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 = 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().to_string())); 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( 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 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(), }, }; 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 = fs::read_to_string(&writeable_filepath).unwrap(); shutdown_tracer_provider(); PyString::new(py, &return_data).into() } fn build_client(py: Python<'_>, py_self: &Py) -> Result { 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) .timeout(std::time::Duration::from_secs(300)) .build() } #[tokio::main] async fn history_logging( py: Python<'_>, py_self: &Py, 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": {}, "checkpoint": "{}", "policy": ["{}"] }}"#, exec_types, checkpoint_number, policy_names ); let res = 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"); return first_response; } Err(_res) => { let failed_response: ApiResponse = ApiResponse { error: "Failed".to_string(), response: ExecHistories { exechistories: vec![], }, }; return failed_response; } } } 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") } _ => home.join(".local").join("share").join("Loxide"), } } fn skipback(days: i64) -> ObjectId { let date_days_ago = Local::now() - Duration::days(days); let timestamp = date_days_ago.timestamp() as u32; let mut hex_timestamp = String::new(); write!(&mut hex_timestamp, "{:08x}", timestamp).unwrap(); let objectid_hex = format!("{}0000000000000000", hex_timestamp); ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex") } fn init_tracer() -> Result, TraceError> { let cfg = TelemetryConfig::load(); if !cfg.TELEMETRY { global::set_tracer_provider(NoopTracerProvider::new()); return Ok(None); } let endpoint = cfg.TELEM_URL.unwrap_or_default(); let tracer = opentelemetry_otlp::new_pipeline() .tracing() .with_exporter( opentelemetry_otlp::new_exporter() .tonic() .with_endpoint(endpoint), ) .with_trace_config(sdktrace::config().with_resource(Resource::new(vec![ KeyValue::new("service.name", "LoxideLibs"), ]))) .install_simple() .unwrap(); Ok(Some(tracer)) }