Bugfix, plus some QoL upgrades from the Async branch
This commit is contained in:
+5
-1
@@ -44,7 +44,11 @@ def generate(api: AirlockAPIWrapper):
|
||||
print(colorText("Please select a duration in minutes: ", "white"))
|
||||
print(colorText("15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", "white"))
|
||||
duration_selected = Selector.select_int(possible_durations)
|
||||
if duration_selected:
|
||||
|
||||
if isinstance(duration_selected, list):
|
||||
duration_selected = duration_selected[0] if duration_selected else None
|
||||
|
||||
if duration_selected is not None:
|
||||
for agent in agents:
|
||||
otp_code = api.otp_generate(agent.agentid, duration_selected, purpose)
|
||||
logger.info(f"Generated OTP for {agent.hostname}: {otp_code}")
|
||||
|
||||
+31
-33
@@ -21,7 +21,7 @@ from typing import List
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from models.execution import ExecutionHistoryRecord, Hash
|
||||
from models.execution import ExecutionHistoryRecord
|
||||
from models.policy import Allowlist, Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.configmanager import get_protected_value, load_env, load_env_json
|
||||
@@ -29,7 +29,6 @@ from utils.selector import Selector
|
||||
from utils.utils import (
|
||||
colorText,
|
||||
formatHTML,
|
||||
import_to_dataframe,
|
||||
regulator,
|
||||
)
|
||||
|
||||
@@ -86,46 +85,45 @@ def sortHashes(
|
||||
if history_days is None:
|
||||
logging.warning("No history range selected. Aborting.")
|
||||
return
|
||||
|
||||
executions = []
|
||||
hashes = []
|
||||
|
||||
# Pull execution histories for each policy
|
||||
|
||||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||||
api, selected_policies, type_=type, history_days=history_days
|
||||
)
|
||||
|
||||
logger.debug(f"Policy_executions is {policy_executions}")
|
||||
|
||||
executions.extend(policy_executions)
|
||||
logger.debug(f"Executions contains {executions}")
|
||||
if executions:
|
||||
hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
|
||||
]
|
||||
logger.debug(f"Executions contains {policy_executions}")
|
||||
|
||||
if hashes:
|
||||
unique_hashes = Hash.deduplicate(hashes)
|
||||
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(api, policy_executions)
|
||||
categorized_executions = ExecutionHistoryRecord.categorize_executions_by_hash_decision(enriched_executions)
|
||||
approved, unapproved, needs_review, unknown = ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
|
||||
|
||||
needs_review, approved, unapproved = Hash.categorize_hashes(
|
||||
hashes=unique_hashes
|
||||
)
|
||||
|
||||
categories = {
|
||||
categories = {
|
||||
"needs_review": needs_review,
|
||||
"approved": approved,
|
||||
"unapproved": unapproved,
|
||||
"unknown" : unknown
|
||||
}
|
||||
|
||||
for label, category in categories.items():
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
|
||||
|
||||
for label, records in categories.items():
|
||||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv"
|
||||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
|
||||
|
||||
# Convert ExecutionHistoryRecord objects to dictionaries
|
||||
df = pd.DataFrame([r.__dict__ for r in records])
|
||||
|
||||
# Optional: flatten hash_obj if needed
|
||||
if not df.empty and 'hash_obj' in df.columns:
|
||||
hash_df = df['hash_obj'].apply(lambda h: h.to_dict() if h else {})
|
||||
df = pd.concat([df.drop(columns=['hash_obj']), hash_df], axis=1)
|
||||
|
||||
# Save to CSV
|
||||
df.to_csv(csv_path, index=False)
|
||||
logger.info(f"Saved {label} executions to {csv_path}")
|
||||
|
||||
# Generate HTML
|
||||
formatHTML(df, html_path)
|
||||
logger.info(f"Generated HTML report at {html_path}")
|
||||
|
||||
ExecutionHistoryRecord.enrich_with_hashes_and_export(
|
||||
executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
|
||||
)
|
||||
df = import_to_dataframe(csv_path)
|
||||
formatHTML(df, html_path)
|
||||
|
||||
def buildPathsandPublishers(split):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
@@ -255,7 +253,7 @@ def splitFilepathsGrouped(df, col="filename"):
|
||||
def clean_split(path):
|
||||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||||
return []
|
||||
parts = os.path.normpath(path).split(os.sep)
|
||||
parts = str(os.path.normpath(path)).split(os.sep)
|
||||
parts = [p for p in parts if p] # Remove empty strings
|
||||
return parts
|
||||
|
||||
@@ -268,9 +266,9 @@ def splitFilepathsGrouped(df, col="filename"):
|
||||
df = df.copy()
|
||||
split_paths = df[col].apply(clean_split)
|
||||
|
||||
# Filter out paths with fewer than `min_files_for_path` components
|
||||
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
if min_files_for_path is not None:
|
||||
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
|
||||
split_paths = split_paths[df.index]
|
||||
|
||||
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
|
||||
grouped = df.groupby("group_key")
|
||||
|
||||
Reference in New Issue
Block a user