AirlockTools/airlock_libs/src/services.rs

use chrono::{Duration, Local, NaiveDate};
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
use mongodb::bson::oid::ObjectId;
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,
    path::PathBuf,
    str::FromStr, sync::mpsc::{Sender, Receiver, channel},
};

#[derive(Debug, Deserialize, Serialize)]
struct ApiResponse {
    error: String,
    response: ExecHistories,
}
#[derive(Debug, Deserialize, Serialize)]
struct ExecHistories {
    exechistories: Vec<Group>,
}
#[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<PyAny>,
    policy_names: String,
    exec_types: String,
    days: i64,
) -> Py<PyString> {
    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()
        {
            fs::create_dir_all(parent_dir).unwrap();
        }
        fs::File::create(file_path).unwrap();
    }
    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 = build_client(py, &py_self);
    let api: Py<PyAny> = py_self;
    let cutoff = Local::now().naive_local() - Duration::days(days);
    let (tx, rx):(Sender<Group>, Receiver<Group>) = channel();
    loop {
        let thread_tx = tx.clone();
        let execution_histories = history_logging(
            py,
            &api,
            &exec_types,
            &checkpoint_number,
            &policy_names,
            &client,
        );
        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 f = File::open(&writeable_filepath).unwrap();
            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();
    PyString::new(py, &return_data).into()
}

fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
    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(30))
        .build()
        .unwrap()
}

#[tokio::main]
async fn history_logging(
    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": {},
        "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")
}