RustImplementation #23

Merged
mysticmomba merged 118 commits from RustImplementation into master 2025-11-04 18:13:24 -05:00
10 changed files with 3627 additions and 4 deletions
Showing only changes of commit a04ce62d83 - Show all commits
+1 -1
View File
@@ -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)
---
+3
View File
@@ -0,0 +1,3 @@
/target
build.sh
pythontest.py
+3016
View File
File diff suppressed because it is too large Load Diff
+59
View File
@@ -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/)
+87
View File
@@ -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'
"""
...
+170
View File
@@ -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();
}
+7
View File
@@ -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(())
}
+282
View File
@@ -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")
}
+2 -3
View File
@@ -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(
api, policy, type_, history_days, True
)
execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days)
if execs:
data = json.loads(execs)
exechistories = data.get("response", {}).get("exechistories", [])