Merge branch 'RustImplementation' of https://git.racooncity.org/brotoskyj/Airlocktools into RustImplementation

README.md

@@ -45,7 +45,7 @@ Python toolkit for secure, auditable, and automated airlock agent and policy man
 
 ## 🧑‍💻 Requirements
 
-TBD
+[airlock_libs](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/0.1.1)
 
 ---
 

airlock_libs/.gitignore

@@ -0,0 +1,3 @@
+/target
+build.sh
+pythontest.py

airlock_libs/README.md

@@ -0,0 +1,59 @@
+
+# Airlock Libs
+
+A Rust implementation of common Airlock API integrations for use in [AirlockTools](https://git.racooncity.org/brotoskyj/AirlockTools)
+
+## Deployment
+
+To install and use this library in a standalone Python script
+### Install Using pip
+Follow the instructions [here](https://git.racooncity.org/brotoskyj/-/packages/pypi/airlock-libs/)
+### Install Wheel
+#### Linux
+
+```bash
+  cd target/wheels/
+  python3 -m pip install airlock_libs-<versionNumber>-cp313-manylinux_2_34_x86_64.whl
+```
+
+#### Windows
+```powershell
+  cd target/wheels
+  python3 -m pip install airlock_libs-<versionNumber>-cp313-win_amd64.whl
+```
+
+or
+
+### Import DLL/SO
+#### Linux
+
+```bash
+cd /target/x86_64-pc-windows-gnu/release/
+Copy libairlock_libs.so to current project directory
+```
+
+#### Windows
+```powershell
+cd /target/x86_64-unknown-linux-gnu/release/
+Copy airlock_libs.dll to current project directory
+```
+## Usage/Examples
+
+```python
+import airlock_libs
+
+def pullExecHistories() {
+  airlock_libs.pull_policy_exec_histories(api, execution_types, checkpoint_number, policy_names)
+}
+```
+
+
+## Features
+
+- Pull Policy Execution Histories
+
+
+## License
+
+[AGPLv3](https://choosealicense.com/licenses/agpl-3.0/)
+

airlock_libs/airlock_libs.pyi

@@ -0,0 +1,87 @@
+from typing import Dict, List, Optional
+def pull_policy_exec_histories(self, type: List[str], checkpoint: str, policy: List[str]) -> str:
+        """Retrieve execution history logs."""
+
+def api(AirlockAPIWrapper):
+    """
+    An implementation of the python AirlockAPIWrapper class to pass Python data into Rust
+
+    Parameters
+    ----------
+    base_url : str
+        (Required) Base URL of the Airlock API, this should be in your .env file.
+    api_key : str
+        (Required) API Key for your profile in airlock, this should be in your credential manager.
+    headers : {"X-APIKey": self.api_key}
+
+    ```def __init__(self, base_url: str, api_key: str):
+            self.base_url = base_ur.rstrip("/")
+            self.api_key = api_key
+            self.headers = {"X-APIKey": self.api_key}
+    ```
+    """
+
+def history_logging(
+        api,
+        exec_types: str,
+        checkpoint_number: str,
+        policy_names: str,
+    ) -> List[Dict[str, Any]]:
+        """
+        Query execution history logs from the Airlock API.
+
+        Parameters
+        ----------
+        exec_types : str
+            A JSON-style string list of execution types to retrieve.
+            Example: "[3,5,8]"
+            - 0 = Trusted Execution
+            - 1 = Blocked Execution
+            - 2 = Untrusted Execution [Audit]
+            - 3 = Untrusted Execution [OTP]
+            - 5 = Trusted Publisher Execution
+            - 8 = Trusted Process Execution
+            (etc.)
+
+        checkpoint_number : str
+            The checkpoint ID. Used to fetch results after a certain event.
+            Example: "601d275487bacb01e3470713"
+
+        policy_names : str
+            A comma-separated or JSON-style list of policy group names.
+            Example: "Apple Mac" or "["Apple Mac", "Servers London"]"
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            A list of dictionaries, where each dictionary represents an
+            execution history record. Each record can include fields like:
+
+            - checkpoint: str
+            - type: int
+            - username: str
+            - hostname: str
+            - filename: str
+            - ppolicy: str
+            - policyname: str
+            - policyver: str
+            - commandline: str
+            - publisher: str
+            - pprocess: str
+            - gprocess: str
+            - sha256: str
+            - datetime: str
+            - ip: str
+            - localip: str
+        Raises
+        ------
+        RuntimeError
+            If the request fails or the response cannot be parsed.
+
+        Example
+        -------
+        >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac")
+        >>> print(histories[0]["filename"])
+        'chrome.exe'
+        """
+        ...

airlock_libs/rust_rewrite.txt

@@ -0,0 +1,170 @@
+use chrono::{Datelike, Duration, NaiveDate, Utc};
+use indicatif::{ProgressBar, ProgressStyle};
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+use std::collections::HashMap;
+use std::fs::{self, File};
+use std::io::{Read, Write};
+use std::path::Path;
+
+// Example API response type
+#[derive(Debug, Serialize, Deserialize, Clone)]
+struct HistoryItem {
+    checkpoint: Option<String>,
+    datetime: String,
+    sha256: Option<String>,
+    filename: Option<String>,
+    hostname: Option<String>,
+    // other fields...
+}
+
+// JSON file structure
+#[derive(Debug, Serialize, Deserialize)]
+struct JsonFile {
+    error: String,
+    response: ResponseData,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct ResponseData {
+    exechistories: Vec<HistoryItem>,
+}
+
+// Mock API function
+fn history_logging(_type: &str, checkpoint: &str, _policy: &[&str]) -> Vec<HistoryItem> {
+    // Replace with actual API call
+    vec![]
+}
+
+fn main() {
+    let file_path = Path::new("data/example.json");
+    let mut checkpoint = "initial_checkpoint".to_string();
+    let policy_name = "policy1".to_string();
+    let days = 7;
+
+    // Initialize JSON output
+    let mut json_output = JsonFile {
+        error: "Success".to_string(),
+        response: ResponseData {
+            exechistories: vec![],
+        },
+    };
+
+    // Outer progress bar (filebar)
+    let filebar = ProgressBar::new(10_000);
+    filebar.set_style(
+        ProgressStyle::default_bar()
+            .template("Checkpoint Progress: {msg} [{bar:40.cyan/blue}] {pos}/{len}")
+            .unwrap(),
+    );
+    filebar.set_message(&checkpoint);
+
+    // Inner progress bar (total progress)
+    let pbar = ProgressBar::new(100);
+    pbar.set_style(
+        ProgressStyle::default_bar()
+            .template("Total of {msg} Complete: [{bar:40.cyan/blue}] {pos}/{len}")
+            .unwrap(),
+    );
+    pbar.set_message(&policy_name);
+
+    loop {
+        let histories = history_logging("type", &checkpoint, &[&policy_name]);
+
+        if !histories.iter().all(|h| h.datetime.len() > 0) {
+            eprintln!("Unexpected response format from API.");
+            break;
+        }
+
+        filebar.set_length(histories.len() as u64);
+
+        if histories.is_empty() {
+            break;
+        }
+
+        for (index, history_item) in histories.iter().enumerate() {
+            if history_item.checkpoint.is_none() || history_item.datetime.is_empty() {
+                continue;
+            }
+
+            if index == histories.len() - 1 {
+                checkpoint = history_item.checkpoint.clone().unwrap();
+                filebar.set_message(&checkpoint);
+                break;
+            }
+
+            // Parse date
+            let history_date = match NaiveDate::parse_from_str(
+                &history_item.datetime.replace(" +0000 UTC", ""),
+                "%Y-%m-%dT%H:%M:%SZ",
+            ) {
+                Ok(date) => date,
+                Err(_) => continue,
+            };
+
+            let cutoff = Utc::today().naive_utc() - Duration::days(days);
+            if history_date >= cutoff {
+                json_output.response.exechistories.push(history_item.clone());
+            }
+
+            filebar.inc(1);
+            filebar.tick();
+        }
+
+        // Deduplicate
+        let mut seen: HashMap<(Option<String>, Option<String>, Option<String>), HistoryItem> =
+            HashMap::new();
+
+        let combined = if file_path.exists() {
+            let mut f = File::open(file_path).unwrap();
+            let mut contents = String::new();
+            f.read_to_string(&mut contents).unwrap();
+            let existing_data: JsonFile = serde_json::from_str(&contents).unwrap_or(JsonFile {
+                error: "Success".to_string(),
+                response: ResponseData {
+                    exechistories: vec![],
+                },
+            });
+            [existing_data.response.exechistories, json_output.response.exechistories.clone()]
+                .concat()
+        } else {
+            json_output.response.exechistories.clone()
+        };
+
+        for entry in combined {
+            let key = (entry.sha256.clone(), entry.filename.clone(), entry.hostname.clone());
+            seen.insert(key, entry);
+        }
+
+        let deduplicated: Vec<HistoryItem> = seen.into_values().collect();
+
+        // Write to file
+        let output_file = File::create(file_path).unwrap();
+        serde_json::to_writer_pretty(&output_file, &json!({
+            "error": "Success",
+            "response": { "exechistories": deduplicated }
+        }))
+        .unwrap();
+
+        json_output.response.exechistories.clear();
+
+        // Update inner progress bar (percentage based on last valid item)
+        if let Some(last_item) = histories.last() {
+            if let Ok(last_date) = NaiveDate::parse_from_str(
+                &last_item.datetime.replace(" +0000 UTC", ""),
+                "%Y-%m-%dT%H:%M:%SZ",
+            ) {
+                let date_diff = Utc::today().naive_utc() - last_date;
+                let percentage_diff = ((days + 10) - date_diff.num_days()) as f64 / (days + 10) as f64 * 100.0;
+                pbar.set_position(percentage_diff.round() as u64);
+                pbar.set_message(&policy_name);
+                pbar.tick();
+            }
+        }
+
+        filebar.set_position(1);
+    }
+
+    filebar.finish();
+    pbar.finish();
+}

airlock_libs/src/lib.rs

@@ -0,0 +1,7 @@
+use pyo3::prelude::*;
+mod services;
+#[pymodule]
+fn airlock_libs(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(services::pull_policy_exec_histories, py)?)?;
+    Ok(())
+}

airlock_libs/src/services.rs

@@ -0,0 +1,282 @@
+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,
+};
+
+#[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,
+    //output_json: bool,
+) -> 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;
+    loop {
+        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,
+            };
+            let cutoff = Local::now().naive_local() - Duration::days(days);
+            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 + 10) - date_diff.num_days()) as f64 / (days + 10) 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();
+    //let json_data = serde_json::from_str(&return_data).unwrap();
+    PyString::new(py, &return_data).into()
+    //let py_any: Py<String> = serde_pyobject::to_pyobject(py, &json_data).unwrap().into();
+    //py_any
+}
+
+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("AirlockTools")
+        }
+        _ => home.join(".local").join("share").join("AirlockTools"),
+    }
+}
+
+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")
+}

models/execution.py

@@ -26,6 +26,7 @@ from typing import List, Optional, Tuple
 import dotenv
 import pandas as pd
 
+import airlock_libs
 from services.API import AirlockAPIWrapper
 from services.policyhandler import pullPolicyExechistories
 from utils.configmanager import get_protected_value, load_env_json
@@ -263,9 +264,7 @@ class ExecutionHistoryRecord:
     ) -> List["ExecutionHistoryRecord"]:
         executions = []
         for policy in selected_policies:
-            execs = pullPolicyExechistories(
+            execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days)
-                api, policy, type_, history_days, True
-            )
             if execs:
                 data = json.loads(execs)
                 exechistories = data.get("response", {}).get("exechistories", [])