Compare commits

...

6 Commits

Author SHA1 Message Date
brotoskyj 0ac3b54d89 Refactored Progress Bar
Build Library / Build Library (push) Successful in 5m54s
Refactored Progress Bar to remove multiprogress bar and only draw one instance. #36 is still open and not fixed with this push, but I believe this is the way to fix the issue.
Also implemented an Arc Mutex on the progress bar so it can be controlled via different threads.
2025-12-05 17:19:40 -05:00
Zarithas 154a7efcc8 Bug fix for Revoke OTP resolved, no longer crashes when select all is chosen when there are no active sessions, Swapped Revoke OTP and Quiet Hosts locations on menus 2025-12-05 15:30:24 -05:00
Zarithas ab5f00d8e7 Merge branch 'RustImplementation' of https://git.racooncity.org/brotoskyj/Airlocktools into RustImplementation 2025-12-05 15:06:06 -05:00
Zarithas 3ab803c12e Quiet Agent UI improvements 2025-12-05 15:05:46 -05:00
Zarithas 98cb23e5ea Bugfix for Issue 39.
Fixes:
brotoskyj/AirlockTools#39
2025-12-05 14:48:23 -05:00
brotoskyj 0aabbfd36e Style Change
Build Library / Build Library (push) Successful in 6m25s
Cleaned up services.rs file and removed whitespaces
Specified Data Types were needed instead of allowing the compiler to select types
2025-12-05 11:12:29 -05:00
8 changed files with 290 additions and 191 deletions
+30 -31
View File
@@ -58,9 +58,14 @@ class OTPRevokeWidget(Static):
}
#button_container {
height: auto;
width: 100%;
padding: 1;
align: center middle;
}
#button_container Button {
min-width: 16;
margin: 0 1;
}
#result_container {
height: auto;
max-height: 10;
@@ -87,29 +92,10 @@ class OTPRevokeWidget(Static):
# Action buttons
with Horizontal(id="button_container"):
self.refresh_button = Button("🔄 Refresh", id="refresh_btn")
self.refresh_button.styles.width = "15%"
self.refresh_button.styles.margin = (1, 1, 1, 1)
yield self.refresh_button
self.select_all_button = Button("☑️ Select All", id="select_all_btn")
self.select_all_button.styles.width = "15%"
self.select_all_button.styles.margin = (1, 1, 1, 1)
yield self.select_all_button
self.select_none_button = Button(
"❌ Clear Selection", id="select_none_btn"
)
self.select_none_button.styles.width = "20%"
self.select_none_button.styles.margin = (1, 1, 1, 1)
yield self.select_none_button
self.revoke_button = Button(
"🛑 Revoke Selected", id="revoke_btn", variant="error"
)
self.revoke_button.styles.width = "20%"
self.revoke_button.styles.margin = (1, 1, 1, 1)
yield self.revoke_button
yield Button("Refresh", id="refresh_btn")
yield Button("Select All", id="select_all_btn")
yield Button("Clear Selection", id="select_none_btn")
yield Button("Revoke Selected", id="revoke_btn", variant="error")
# Results display
with Vertical(id="result_container"):
@@ -122,7 +108,7 @@ class OTPRevokeWidget(Static):
# Configure sessions table
self.sessions_table.clear()
self.sessions_table.add_columns(
"", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
"", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
)
# Enable row selection with checkbox column
@@ -213,7 +199,11 @@ class OTPRevokeWidget(Static):
elif btn.id == "select_all_btn":
# Select all visible rows
if self._filtered_df is not None:
if (
self._filtered_df is not None
and not self._filtered_df.empty
and "otpid" in self._filtered_df.columns
):
self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"])
await self._refresh_table()
@@ -235,7 +225,12 @@ class OTPRevokeWidget(Static):
# Get the row index from the cursor row
row_index = self.sessions_table.cursor_row
if self._filtered_df is not None and row_index < len(self._filtered_df):
if (
self._filtered_df is not None
and not self._filtered_df.empty
and "otpid" in self._filtered_df.columns
and row_index < len(self._filtered_df)
):
# Get the OTP ID for this row
otpid = str(self._filtered_df.iloc[row_index]["otpid"])
@@ -257,12 +252,12 @@ class OTPRevokeWidget(Static):
async def _revoke_selected(self) -> None:
"""Revoke the selected OTP sessions."""
if not self._selected_otpids:
self.results_display.update("No sessions selected for revocation")
self.results_display.update("No sessions selected for revocation")
return
api = getattr(self.app, "api", None)
if not api:
self.results_display.update("API not available")
self.results_display.update("API not available")
return
# Collect results
@@ -303,13 +298,13 @@ class OTPRevokeWidget(Static):
else "No response"
)
results.append(
f"Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
f"Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
)
logger.error(f"Failed to revoke OTP {otpid}: {error_msg}")
except Exception as e:
failure_count += 1
results.append(f"Error revoking OTP {otpid}: {str(e)}")
results.append(f"Error revoking OTP {otpid}: {str(e)}")
logger.exception(f"Exception revoking OTP {otpid}: {e}")
# Update results display
@@ -368,7 +363,11 @@ class OTPRevokeScreen(Screen):
async def action_select_all(self) -> None:
"""Select all visible sessions."""
if self.widget._filtered_df is not None:
if (
self.widget._filtered_df is not None
and not self.widget._filtered_df.empty
and "otpid" in self.widget._filtered_df.columns
):
self.widget._selected_otpids = set(
str(x) for x in self.widget._filtered_df["otpid"]
)
+183 -110
View File
@@ -34,7 +34,7 @@ from textual.app import ComposeResult
from textual.containers import Horizontal, Vertical
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import Button, DataTable, Footer, Header, Static
from textual.widgets import Button, DataTable, Footer, Header, Input, Static
from models.policy import Policy
from services.API import AirlockAPIWrapper
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
This screen provides a multi-step workflow:
1. Select initial policy to analyze
2. View categorized agents (enforce ready vs. non-enforce ready)
3. Select target policies for each category
4. Execute agent migrations
2. Configure analysis parameters (history period and quiet time period)
3. View categorized agents (enforce ready vs. non-enforce ready)
4. Select target policies for each category
5. Execute agent migrations
Attributes:
api (AirlockAPIWrapper): API wrapper for Airlock operations
policies (List[Policy]): List of all available policies
selected_policy (Optional[Policy]): The initially selected policy to analyze
history_days (int): Number of days of history to pull (default: 150)
quiet_days (int): Number of days without execution to be considered quiet (default: 45)
history_days (int): Number of days of history to pull (default: 150, range: 1-365)
quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
@@ -86,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
self.api = api
self.policies = policies
self.selected_policy: Optional[Policy] = None
self.history_days = 150 # Fixed as per requirements
self.history_days = 150 # Default value, user-selectable
self.quiet_days = 45 # Default value
self.agents_df: Optional[pd.DataFrame] = None
self.enforce_ready_df: Optional[pd.DataFrame] = None
@@ -130,7 +131,7 @@ class QuietAgentWorkflowScreen(Screen):
stage_messages = {
"select_policy": "Step 1: Select Policy to Analyze",
"select_quiet_days": "Step 2: Select Quiet Time Period",
"select_history_days": "Step 2: Configure Analysis Parameters",
"analyzing": "Analyzing agent activity...",
"view_results": "Step 3: Review Categorized Agents",
"select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
@@ -161,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
# Initial policy selection for analysis
self.selected_policy = message.policy
logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
self._show_quiet_days_selection()
self._show_history_days_selection()
elif self.workflow_stage == "select_enforce_target":
# Target policy selection for enforce ready agents
self.enforce_ready_target_policy = message.policy
@@ -177,48 +178,167 @@ class QuietAgentWorkflowScreen(Screen):
)
self._show_migration_confirmation()
def _show_quiet_days_selection(self) -> None:
"""Show the quiet days selection screen."""
self.workflow_stage = "select_quiet_days"
def _show_history_days_selection(self) -> None:
"""Show the history days and quiet days selection screen."""
self.workflow_stage = "select_history_days"
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Create info text
info_widget = Static(
f"Policy Selected: {self.selected_policy.name}\n\n"
f"History Period: {self.history_days} days\n\n"
"Select quiet time period (days without untrusted execution):",
id="quiet_days_info",
"Configure Analysis Parameters:",
id="analysis_params_info",
)
info_widget.styles.margin = (0, 0, 2, 0)
content.mount(info_widget)
# Create button container and mount it first
button_container = Vertical(id="quiet_days_buttons")
button_container.styles.height = "auto"
content.mount(button_container)
# Create input container
input_container = Vertical(id="analysis_params_input_container")
input_container.styles.height = "auto"
content.mount(input_container)
# Now add buttons to the mounted container
for days in [15, 30, 45, 60]:
btn = Button(
f"{days} days {'(Default)' if days == 45 else ''}",
id=f"quiet_days_{days}",
classes="quiet_day_btn",
# History days label
history_label = Static("History Period (days of execution history to pull):")
history_label.styles.margin = (0, 0, 1, 0)
input_container.mount(history_label)
# Add history days input field
history_input = Input(
placeholder="Enter days (1-365, default: 150)",
value="150",
id="history_days_input",
)
history_input.styles.width = "50"
history_input.styles.margin = (0, 0, 2, 0)
input_container.mount(history_input)
# Quiet days label
quiet_label = Static(
"Quiet Time Period (days without execution to be considered quiet):"
)
quiet_label.styles.margin = (0, 0, 1, 0)
input_container.mount(quiet_label)
# Add quiet days input field
quiet_input = Input(
placeholder="Enter days (1-365, default: 45)",
value="45",
id="quiet_days_input",
)
quiet_input.styles.width = "50"
quiet_input.styles.margin = (0, 0, 2, 0)
input_container.mount(quiet_input)
# Add submit button
submit_btn = Button(
"Continue",
id="analysis_params_submit",
variant="primary",
)
submit_btn.styles.width = "50"
submit_btn.styles.margin = (1, 0, 0, 0)
input_container.mount(submit_btn)
# Focus the first input field
history_input.focus()
def _validate_and_submit_history_days(self) -> None:
"""Validate and submit the history days and quiet days inputs."""
try:
history_input = self.query_one("#history_days_input", Input)
quiet_input = self.query_one("#quiet_days_input", Input)
history_value = history_input.value.strip()
quiet_value = quiet_input.value.strip()
# Validate history days
if not history_value:
self.app.notify(
"Please enter a history period value", severity="error", timeout=3
)
history_input.focus()
return
try:
history_days = int(history_value)
except ValueError:
self.app.notify(
"Please enter a valid number for history period",
severity="error",
timeout=3,
)
history_input.focus()
return
if history_days < 1 or history_days > 365:
self.app.notify(
"History period must be between 1 and 365 days",
severity="error",
timeout=3,
)
history_input.focus()
return
# Validate quiet days
if not quiet_value:
self.app.notify(
"Please enter a quiet time period value",
severity="error",
timeout=3,
)
quiet_input.focus()
return
try:
quiet_days = int(quiet_value)
except ValueError:
self.app.notify(
"Please enter a valid number for quiet time period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
if quiet_days < 1 or quiet_days > 365:
self.app.notify(
"Quiet time period must be between 1 and 365 days",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# Check that quiet days doesn't exceed history days
if quiet_days > history_days:
self.app.notify(
"Quiet time period cannot exceed history period",
severity="error",
timeout=3,
)
quiet_input.focus()
return
# All validation passed
self.history_days = history_days
self.quiet_days = quiet_days
logger.info(
f"Selected history days: {history_days}, quiet days: {quiet_days}"
)
btn.styles.width = "100%"
btn.styles.margin = (0, 0, 1, 0)
button_container.mount(btn)
self._start_analysis()
except Exception as e:
logger.error(f"Error validating analysis parameters: {e}")
self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button press events."""
button_id = event.button.id
# Quiet days selection buttons
if button_id and button_id.startswith("quiet_days_"):
days = int(button_id.split("_")[-1])
self.quiet_days = days
logger.info(f"Selected quiet days: {days}")
self._start_analysis()
# Analysis parameters submit button
if button_id == "analysis_params_submit":
self._validate_and_submit_history_days()
return
# Navigation buttons
@@ -258,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection()
return
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle input submission (Enter key pressed)."""
if event.input.id in ["history_days_input", "quiet_days_input"]:
self._validate_and_submit_history_days()
def _start_analysis(self) -> None:
"""Start the agent activity analysis."""
self.workflow_stage = "analyzing"
content = self.query_one("#content_area", Vertical)
content.remove_children()
# Show analyzing message with detailed steps
analyzing_msg = Static(
f"Analyzing Agent Activity\n"
f"{'=' * 50}\n\n"
f"Policy: {self.selected_policy.name}\n"
f"History Period: {self.history_days} days\n"
f"Quiet Threshold: {self.quiet_days} days\n\n"
f"Progress:\n"
f"Step 1/4: Fetching agents from policy...\n"
f"Step 2/4: Pulling execution history (this may take a moment)...\n"
f"Step 3/4: Analyzing activity patterns...\n"
f"Step 4/4: Categorizing agents...\n\n"
f"Please wait - this operation cannot be cancelled.",
id="analyzing_message",
)
analyzing_msg.styles.margin = (2, 1)
content.mount(analyzing_msg)
# Show notification
# Show notification that analysis is starting
self.app.notify(
"Starting analysis - this may take several minutes for large policies",
severity="information",
timeout=5,
)
# Perform the analysis asynchronously
self.call_later(self._perform_analysis)
# Clear the screen to provide a blank canvas for Rust progress output
# (Rust output displays over the TUI, so we clear everything except header/footer)
try:
# Clear title
title_widget = self.query_one("#workflow_title", Static)
title_widget.update("")
def _perform_analysis(self) -> None:
# Clear status
status_widget = self.query_one("#workflow_status", Static)
status_widget.update("")
# Clear content area
content = self.query_one("#content_area", Vertical)
content.remove_children()
except Exception as e:
logger.debug(f"Could not clear screen for analysis: {e}")
# Delay the analysis start to ensure UI refresh completes first
# This prevents Rust output from starting before the screen is cleared
self.set_timer(0.5, self._perform_analysis_worker)
def _perform_analysis_worker(self) -> None:
"""Perform the actual agent activity analysis."""
try:
# Update status: Fetching agents
self._update_analysis_status("Step 1/4: Fetching agents from policy...")
# Get agents in the selected policy
agents = self.api.agents_find_by_group(self.selected_policy.groupid)
@@ -310,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
self._show_policy_selection()
return
agent_count = len(agents)
self.app.notify(
f"Found {agent_count} agents - fetching execution history...",
severity="information",
timeout=3,
)
# Update status: Pulling execution history
self._update_analysis_status(
f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
f"(This may take several minutes - progress shown in terminal)"
)
# Get execution history (this shows progress bars in terminal via airlock_libs)
policy_exec_history = getPolicyInfo(
self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
)
# Update status: Analyzing patterns
self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
self.app.notify(
"History retrieved - analyzing patterns...",
severity="information",
timeout=2,
)
if policy_exec_history.empty:
logger.info(
"No execution history found for the selected policy and time range."
@@ -381,9 +478,6 @@ class QuietAgentWorkflowScreen(Screen):
lambda x: True if pd.isna(x) or x > self.quiet_days else False
)
# Update status: Categorizing
self._update_analysis_status("Step 4/4: Categorizing agents...")
# Sort agents
agents = agents.sort_values(
by=["execution_count", "hostname"], ascending=[True, True]
@@ -394,7 +488,7 @@ class QuietAgentWorkflowScreen(Screen):
# Categorize agents into DataFrames
self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
logger.info(
f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
@@ -416,27 +510,6 @@ class QuietAgentWorkflowScreen(Screen):
self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
self._show_policy_selection()
def _update_analysis_status(self, status_text: str) -> None:
"""Update the analysis status message."""
try:
analyzing_msg = self.query_one("#analyzing_message", Static)
# Build updated message
updated_text = (
f"Analyzing Agent Activity\n"
f"{'=' * 50}\n\n"
f"Policy: {self.selected_policy.name}\n"
f"History Period: {self.history_days} days\n"
f"Quiet Threshold: {self.quiet_days} days\n\n"
f"Progress:\n"
f"{status_text}\n\n"
f"Please wait - this operation cannot be cancelled."
)
analyzing_msg.update(updated_text)
except Exception as e:
logger.debug(f"Could not update analysis status: {e}")
def _show_results(self) -> None:
"""Show the categorized results."""
self.workflow_stage = "view_results"
@@ -821,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
# Depending on stage, go back to previous stage or exit
if self.workflow_stage in ["select_policy", "view_results", "complete"]:
self.app.pop_screen()
elif self.workflow_stage == "select_quiet_days":
elif self.workflow_stage == "select_history_days":
self._show_policy_selection()
elif self.workflow_stage == "select_enforce_target":
self._show_results()
+4 -4
View File
@@ -92,15 +92,15 @@ class MainMenuScreen(Screen):
BUTTON_DEFS = {
"agent_actions": [
(
"🖥️ - Find, Move, or Generate OTP for Agents",
"🖥️ - Find agent, Move agent, or Generate One Time Pass",
"move_agent_workflow_button",
),
("🎫 - Review and appove OTP Activities", "otp_activities_button"),
("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
("🎫 - Review and approve OTP Activities", "otp_activities_button"),
("🛑 - Revoke Active OTP Session", "otp_revoke_button"),
],
"policy": [
("⚖️ - Prepare Policy For Enforcement", "policy_prep_button"),
("🛑 - Revoke OTPs", "otp_revoke_button"),
("🔕 - Find and Move Quiet Hosts to Enforcement", "find_quiet_button"),
],
}
+18 -3
View File
@@ -26,11 +26,13 @@ dependencies = [
[[package]]
name = "airlock_libs"
version = "5.1.1"
version = "5.2.0"
dependencies = [
"chrono",
"crossbeam",
"flexi_logger",
"indicatif",
"log",
"mongodb",
"opentelemetry 0.18.0",
"opentelemetry-otlp",
@@ -799,6 +801,19 @@ version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
[[package]]
name = "flexi_logger"
version = "0.31.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31e5335674a3a259527f97e9176a3767dcc9b220b8e29d643daeb2d6c72caf8b"
dependencies = [
"chrono",
"log",
"nu-ansi-term",
"regex",
"thiserror 2.0.17",
]
[[package]]
name = "fnv"
version = "1.0.7"
@@ -1590,9 +1605,9 @@ dependencies = [
[[package]]
name = "log"
version = "0.4.28"
version = "0.4.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432"
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
dependencies = [
"value-bag",
]
+3 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "airlock_libs"
version = "5.1.1"
version = "5.2.0"
edition = "2024"
[lib]
@@ -26,6 +26,8 @@ tracing-subscriber = "0.3.20"
tracing-opentelemetry = "0.32.0"
pyo3-async-runtimes = { version = "0.27.0", features = ["async-std", "tokio"] }
crossbeam = "0.8.4"
log = "0.4.29"
flexi_logger = "0.31.7"
[package.metadata.maturin]
generate-abi-stubs = true
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "maturin"
[project]
name = "airlock_libs"
version = "5.1.1"
version = "5.2.0"
description = "Airlock Digital API Wrapper"
readme = "README.md"
license = { text = "AGPL-3.0-only" }
+50 -40
View File
@@ -1,9 +1,8 @@
use std::thread;
use crossbeam::channel::unbounded;
use crate::modules::datatypes::*;
use crate::prelude::*;
use crossbeam::channel::unbounded;
use std::sync::{Arc, Mutex};
use std::thread;
#[pyfunction]
pub fn pull_policy_exec_histories(
py: Python<'_>,
@@ -16,12 +15,12 @@ pub fn pull_policy_exec_histories(
ExtractedValues::Headers(h) => h,
ExtractedValues::BaseUrl(_) => std::process::abort(),
};
let base_url = match PyData::convert(py, &py_self, false) {
let base_url: String = match PyData::convert(py, &py_self, false) {
ExtractedValues::Headers(_) => std::process::abort(),
ExtractedValues::BaseUrl(b) => b,
};
let handle = std::thread::spawn(move || {
let rt = match tokio::runtime::Runtime::new() {
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);
@@ -31,8 +30,8 @@ pub fn pull_policy_exec_histories(
rt.block_on(async {
let _ = init_tracer();
});
let tracer = global::tracer("global_tracer");
let _cx = Context::new();
let tracer: global::BoxedTracer = global::tracer("global_tracer");
let _cx: Context = Context::new();
let file_path: PathBuf = format!(
"{}\\cache\\chunkinator.json",
get_base_directory().display()
@@ -58,14 +57,14 @@ pub fn pull_policy_exec_histories(
}
}
}
let data = ApiResponse {
let data: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: vec![],
},
};
let writeable_filepath = file_path.clone();
let data_write = serde_json::to_string_pretty(&data).expect("Failed to serialize");
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) => {
@@ -74,17 +73,18 @@ pub fn pull_policy_exec_histories(
}
}
let mut checkpoint_number: String = SkipBack::find_checkpoint(days).to_string();
let multi_progress = MultiProgress::new();
multi_progress.set_draw_target(ProgressDrawTarget::stderr());
let progress_bar = multi_progress.add(ProgressBar::new(100));
progress_bar.set_style(
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 Completion: {spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {pos}/{len} {message}")
.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(headers);
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(
@@ -111,12 +111,13 @@ pub fn pull_policy_exec_histories(
}
}
});
let cutoff = Local::now().naive_local() - Duration::days(days);
let cutoff: chrono::NaiveDateTime = Local::now().naive_local() - Duration::days(days);
let (tx, rx) = unbounded::<Vec<Group>>();
let pb_clone = progress_bar.clone();
thread::spawn(move || {
let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
{
let contents = fs::read_to_string(&writeable_filepath).unwrap_or_default();
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(),
@@ -128,7 +129,7 @@ pub fn pull_policy_exec_histories(
.response
.exechistories
.into_iter()
.map(|entry| {
.map(|entry: Group| {
(
(
entry.sha256.clone(),
@@ -147,7 +148,7 @@ pub fn pull_policy_exec_histories(
if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
continue;
}
let history_date = match NaiveDate::parse_from_str(
let history_date: NaiveDate = match NaiveDate::parse_from_str(
&executions.datetime.replace(" +0000 UTC", ""),
"%Y-%m-%dT%H:%M:%SZ",
) {
@@ -155,7 +156,7 @@ pub fn pull_policy_exec_histories(
Err(_) => continue,
};
if history_date >= cutoff.into() {
let key = (
let key: (String, String, String) = (
executions.sha256.clone(),
executions.filename.clone(),
executions.hostname.clone(),
@@ -163,13 +164,13 @@ pub fn pull_policy_exec_histories(
seen.entry(key).or_insert(executions.clone());
}
}
let final_response = ApiResponse {
let final_response: ApiResponse = ApiResponse {
error: "Success".to_string(),
response: ExecHistories {
exechistories: seen.values().cloned().collect(),
},
};
let data_write = serde_json::to_string_pretty(&final_response).unwrap();
let data_write: String = serde_json::to_string_pretty(&final_response).unwrap();
match fs::write(&writeable_filepath, data_write) {
Ok(_) => {}
Err(e) => {
@@ -180,7 +181,11 @@ pub fn pull_policy_exec_histories(
});
let mut first_date: Option<NaiveDate> = None;
tracer.in_span("Airlock Data Retreival", |cx| {
let span = cx.span();
pb_clone
.lock()
.unwrap()
.enable_steady_tick(std::time::Duration::from_millis(100));
let span: opentelemetry::trace::SpanRef<'_> = cx.span();
span.set_attribute(Key::new("Days").string(days.to_string()));
span.set_attribute(KeyValue::new("Policy Name", policy_names.clone()));
loop {
@@ -198,7 +203,7 @@ pub fn pull_policy_exec_histories(
));
results
});
let parsed_responses = execution_histories.response.exechistories;
let parsed_responses: Vec<Group> = execution_histories.response.exechistories;
if parsed_responses.is_empty() {
break;
}
@@ -214,18 +219,23 @@ pub fn pull_policy_exec_histories(
first_date = Some(last_date);
}
if let Some(base_date) = first_date {
let date_diff = last_date - base_date;
let total_span = (Local::now().naive_local().date() - base_date).num_days();
let percentage = ((date_diff.num_days() as f64 / total_span as f64) * 100.0)
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;
progress_bar.set_position(percentage);
pb_clone.lock().unwrap().set_position(percentage);
}
}
}
});
progress_bar.finish_with_message("All Checkpoints Complete");
let return_data = match fs::read_to_string(file_path.clone()) {
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);
@@ -236,8 +246,8 @@ pub fn pull_policy_exec_histories(
shutdown_tracer_provider();
return_data.to_string()
});
let gil_value = handle.join().unwrap();
Python::attach(|py| PyString::new(py, &gil_value).into())
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> {
@@ -264,7 +274,7 @@ async fn history_logging(
}}"#,
exec_types, checkpoint_number, policy_names
);
let res = client
let res: Result<reqwest::Response, reqwest::Error> = client
.post(format!("{}/v1/logging/exechistories", base_url))
.body(payload)
.send()
@@ -305,13 +315,13 @@ pub fn get_base_directory() -> PathBuf {
}
fn init_tracer() -> Result<Option<sdktrace::Tracer>, TraceError> {
let cfg = TelemetryConfig::load();
let cfg: TelemetryConfig = TelemetryConfig::load();
if !cfg.TELEMETRY {
global::set_tracer_provider(NoopTracerProvider::new());
return Ok(None);
}
let endpoint = cfg.TELEM_URL.unwrap_or_default();
let tracer =
let endpoint: String = cfg.TELEM_URL.unwrap_or_default();
let tracer: sdktrace::Tracer =
opentelemetry_otlp::new_pipeline()
.tracing()
.with_exporter(
+1 -1
View File
@@ -11,4 +11,4 @@ urllib3==2.5.0
pyperclip==1.11.0
--extra-index-url https://git.racooncity.org/api/packages/brotoskyj/pypi/simple/
airlock_libs==5.1.1
airlock_libs==5.2.0