AirlockTools/airlock_libs/src/services.rs

use crate::modules::datatypes::*;
use crate::prelude::*;
use opentelemetry::trace::SpanContext;

#[pyfunction]
pub fn pull_policy_exec_histories(
    py: Python<'_>,
    py_self: Py<PyAny>,
    policy_names: Option<String>,
    exec_types: String,
    days: i64,
) -> Py<PyString> {
    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<String> = 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<Client, reqwest::Error> = 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<Group>)>();
        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<NaiveDate> = 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<Group> = 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<reqwest::Client, reqwest::Error> {
    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<String>,
    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<reqwest::Response, reqwest::Error> = 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();
        }
    }
}