AirlockTools/airlock_libs/src/services.rs

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<String>,
}

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::<Self>(&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<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 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<PyAny> = 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<PyAny>) -> 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()
        .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<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")
}

fn init_tracer() -> Result<Option<sdktrace::Tracer>, 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))
}