mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance

Browse Source 
- Migrated codebase to class-based architecture for better modularity and maintainability
- Introduced system_config.json for centralized configuration (required for runtime)
- Added structured working directories for improved file organization
- Significantly reduced reliance on Parquet; replaced with alternative data handling
- Implemented security improvements across modules
- Several TODOs remain in the main script for future enhancements
- Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
Zarithas
2025-10-05 22:20:42 -04:00
parent b1e6af01c6
commit 89db386ffe
26 changed files with 3530 additions and 2389 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/clientfunctions.py
-430
View File
@@ -1,430 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.hashfunctions as hashf
import utils.utils as ct
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
#Standard Libary Imports:
import datetime
import json
import os
import re
#3rd Party Imports:
import pandas as pd
import requests
def findAgentID(url):
print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
hostname = input(ct.colorText("Enter Device Name: ", "white"))
endpoint = url + '/v1/agent/find'
payload = {
"hostname" : f"{hostname}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["agents"])
return data.loc[0, "agentid"]
def getDestAllowlistFromClientID(url, clientid):
#This is dependant on their being an allowlist in the policy containing the name "localapproval"
endpoint = url + '/v1/agent/find'
payload = {
"agentid" : f"{clientid}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["agents"])
allowlist = policyf.getDestAllowlist(url,data.loc[0, "groupid"])
return allowlist
def getPolicyFromClientID(url, clientid):
endpoint = url + '/v1/agent/find'
payload = {
"agentid" : f"{clientid}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["agents"])
policy_name = getPolicyName(url,data.loc[0, "groupid"])
policy_id = data.loc[0, "groupid"]
return policy_name, policy_id
def getPolicyName(url, groupid):
endpoint = url + '/v1/group/'
payload = {}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["groups"])
# Safely attempt to get the policy name
filtered = data.loc[data['groupid'] == str(groupid), 'name']
if not filtered.empty:
name = filtered.values[0]
else:
name = None # or "Unknown", depending on your preference
return name
def devicehistory(url, outputjson: bool):
endpoint = url + '/v1/getexechistory'
print("\n")
print(ct.colorText("1. Today", "yellow"))
print(ct.colorText("2. Last 24 Hours", "yellow"))
print(ct.colorText("3. Past 7 Days", "yellow"))
print(ct.colorText("4. Past 30 Days", "yellow"))
print(ct.colorText("5. Custom Date Range","yellow"))
choice = input(ct.colorText("\nSelect Date Range: ", "white"))
today = datetime.date.today()
today = today.strftime("%Y-%m-%d")
date_selected = " "
if choice == '1':
date_selected = today
elif choice == '2':
date_selected = datetime.date.today() - datetime.timedelta(days=1)
date_selected = date_selected.strftime('%Y-%m-%d')
elif choice == '3':
date_selected = datetime.date.today() - datetime.timedelta(days=7)
date_selected = date_selected.strftime('%Y-%m-%d')
elif choice == '4':
date_selected = datetime.date.today() - datetime.timedelta(days=30)
date_selected = date_selected.strftime('%Y-%m-%d')
elif choice == "5":
print(ct.colorText("Please Input Dates as YYYY-MM-DD", "cyan"))
date_selected = input(ct.colorText("From: ", "white"))
today = input(ct.colorText("Date To: ", "white"))
print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
device = input(ct.colorText("Enter Device Name: ", "white"))
payload_dict = {
"datefrom": date_selected,
"dateto": today,
"hostname": device
}
payload = json.dumps(payload_dict)
print(ct.colorText(payload, "green"))
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
if outputjson:
return response
parse_text = json.loads(response.text)
# Safely get exechistory
exechistory = parse_text.get('response', {}).get('exechistory')
if isinstance(exechistory, list):
for block in exechistory:
print(ct.colorText(f"Command: {block.get('commandline', 'N/A')}", "green"))
print(ct.colorText(f"Date: {block.get('datetime', 'N/A')}", "green"))
print(ct.colorText(f"Filename: {block.get('filename', 'N/A')}", "green"))
print(ct.colorText(f"Policy Name: {block.get('policyname', 'N/A')}", "green"))
print(ct.colorText(f"Hostname: {block.get('hostname', 'N/A')}", "green"))
print(ct.colorText(f"Hash: {block.get('sha256', 'N/A')}", "green"))
print("\n")
else:
print(ct.colorText("No execution history found or data is not in expected format.", "red"))
def findAllAgents(url):
endpoint = url + '/v1/agent/find'
payload = {}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["agents"])
group_ids = sorted(data['groupid'].unique().tolist())
group_policy_map = {}
for groupid in group_ids:
policy_name = getPolicyName(url, groupid)
group_policy_map[groupid] = policy_name
data['policy_name'] = data['groupid'].map(group_policy_map)
status_map = {
0: 'Offline',
1: 'Online',
2: 'Hidden',
3: 'Safemode'
}
data['status'] = data['status'].map(status_map)
return(data)
def findAgents(url, device_input_str, return_dataframe):
os.makedirs("device_search", exist_ok=True)
df = findAllAgents(url)
#Parse the input string into device names (newline-separated only)
device_names = device_input_str.strip().split('\n')
device_names = [name.strip() for name in device_names if name.strip()]
#Build a regex pattern for case-insensitive matching
pattern = '|'.join([re.escape(name) for name in device_names])
regex = re.compile(pattern, re.IGNORECASE)
#Filter the DataFrame using regex
matched_df = df[df['hostname'].apply(lambda x: bool(regex.search(str(x))))]
if return_dataframe : return matched_df
else:
#Export to CSV with timestamp
timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = f"agentsearch_{timestamp}.csv"
matched_df.to_csv(f"device_search\\{filename}", index=False)
print(ct.colorText(f"\n✅ Matched devices exported to: device_search\\{filename}","green"))
def promptForDevices():
print(ct.colorText("🔍 Device Search", "cyan"))
print(ct.colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
print(ct.colorText("When you're done, press Enter twice.\n", "cyan"))
print(ct.colorText("Example:", "cyan"))
print(ct.colorText("H00000", "cyan"))
print(ct.colorText("UTN00000", "cyan"))
print(ct.colorText("i-hSuperSecretServer", "cyan"))
print(ct.colorText("u-hVenderBroke\n", "cyan"))
# Collect multiline input from user
print(ct.colorText("Paste or type your device names below:", "white"))
device_input_lines = []
while True:
line = input()
if line == "":
break
device_input_lines.append(line)
device_input_str = "\n".join(device_input_lines)
return device_input_str
def findQuietAgents(url):
# Get policy selection and agent list
choice, policynames, policyid = policyf.choosePolicies(url)
policy = policynames[choice]
groupid = policyid[choice]
agents = findGroupAgents(url, groupid)
# Prompt user for history range
while True:
try:
history_days = int(input("Enter how many days of history to pull (1–150): "))
if 1 <= history_days <= 150:
break
else:
print("Invalid input. Please enter a number between 1 and 150.")
except ValueError:
print("Invalid input. Please enter a valid integer.")
while True:
try:
required_quiet = int(input("Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): "))
if 1 <= required_quiet <= 365:
break
else:
print("Invalid input. Please enter a number between 1 and 365.")
except ValueError:
print("Invalid input. Please enter a valid integer.")
# Get policy execution history
policy_exec_history = policyf.getPolicyInfo(url, policy, [1, 2, 6, 7], history_days, False)
# Convert 'datetime' column to timezone-aware datetime objects
policy_exec_history['datetime'] = pd.to_datetime(policy_exec_history['datetime'], format='%Y-%m-%dT%H:%M:%SZ', utc=True)
# Get current UTC time
now = datetime.datetime.now(datetime.timezone.utc)
# Calculate days ago
policy_exec_history['days_ago'] = policy_exec_history['datetime'].apply(lambda dt: (now - dt).days)
# Count total executions per hostname
hostname_counts = policy_exec_history['hostname'].value_counts()
# Map execution counts to agents
agents['execution_count'] = agents['hostname'].map(hostname_counts).fillna(0).astype(int)
# Find most recent execution per hostname
most_recent_exec = policy_exec_history.sort_values(by='days_ago').drop_duplicates(subset='hostname', keep='first')
# Map most recent execution age to agents
agents['days_since'] = agents['hostname'].map(most_recent_exec.set_index('hostname')['days_ago'])
#Check for enforcement readyness
agents['required_quiet'] = required_quiet
agents['enforce_ready'] = agents['days_since'].apply(
lambda x: True if pd.isna(x) or x > required_quiet else False
)
# Sort agents by execution count and hostname
agents = agents.sort_values(by=['execution_count', 'hostname'], ascending=[True, True])
# Save to CSV
print(ct.colorText(f"Saving CSV to {policy}_agents_last_{history_days}_days.csv", "green"))
agents.to_csv(f"{policy}_agents_last_{history_days}_days.csv", index=False)
# Summary statistics
total_agents = len(agents)
ready_agents = agents['enforce_ready'].sum()
not_ready_agents = total_agents - ready_agents
ready_percentage = (ready_agents / total_agents) * 100
# Print results
print(f"Total agents: {total_agents}")
print(f"Agents marked as 'enforce_ready': {ready_agents}")
print(f"Agents not ready: {not_ready_agents}")
print(f"Percentage ready for enforcement: {ready_percentage:.2f}%")
def findGroupAgents(url, groupid):
endpoint = url + '/v1/agent/find'
payload = {
"groupid" : f"{groupid}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["agents"])
group_ids = sorted(data['groupid'].unique().tolist())
group_policy_map = {}
for groupid in group_ids:
policy_name = getPolicyName(url, groupid)
group_policy_map[groupid] = policy_name
data['policy_name'] = data['groupid'].map(group_policy_map)
status_map = {
0: 'Offline',
1: 'Online',
2: 'Hidden',
3: 'Safemode'
}
data['status'] = data['status'].map(status_map)
return(data)
def moveAgentToAudit(url, agentid, policy_relationship_map):
policy_name, policy_id = getPolicyFromClientID(url, agentid)
print(policy_id)
if policy_id in policy_relationship_map:
target_policy = policy_relationship_map[policy_id]
elif policy_id in policy_relationship_map.values():
print(f"Agent {agentid} is already in an audit group. No action needed.")
return
else:
print(f"Error: No corresponding audit policy found for policy: {policy_name} - {policy_id}.")
return
moveAgent(url, agentid, target_policy, "audit")
def moveAgentToEnforcement(url, agentid, policy_relationship_map):
policy_name, policy_id = getPolicyFromClientID(url, agentid)
# Invert the map for audit → enforcement
inverse_map = {v: k for k, v in policy_relationship_map.items()}
if policy_id in inverse_map:
target_policy = inverse_map[policy_id]
elif policy_id in inverse_map.values():
print(f"Agent {agentid} is already in an enforcement group. No action needed.")
return
else:
print(f"Error: No corresponding enforcement policy found for policy: {policy_name} - {policy_id}.")
return
moveAgent(url, agentid, target_policy, "enforcement")
def moveAgent(url, agentid, target_policy, direction):
endpoint = f"{url}/v1/agent/move"
payload = {
"groupid": target_policy,
"agentid": agentid
}
headers = {"X-APIKey": os.getenv('APIKEY')}
response = None # Initialize to avoid unbound errors
try:
response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)
response.raise_for_status() # Raises HTTPError for bad status codes
result = response.json()
# Check if 'error' key exists and if it's not a success message
if "error" in result and result["error"].lower() != "success":
print(f"API returned an error: {result['error']}")
else:
print(f"✅ Agent {agentid} successfully moved to {direction} group {target_policy}.")
except requests.exceptions.HTTPError as http_err:
print(f"HTTP error occurred: {http_err}")
if response is not None:
print("Raw response:", response.text)
except requests.exceptions.RequestException as req_err:
print(f"Request error occurred: {req_err}")
except ValueError:
print("Failed to parse JSON response.")
if response is not None:
print("Raw response:", response.text)
except Exception as e:
print(f"Unexpected error: {e}")
if response is not None:
print("Raw response:", response.text)
utils/hashfunctions.py
-296
View File
@@ -1,296 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.pathfunctions as pathf
import utils.utils as ct
#Standard Libary Imports:
import gc
import json
import os
#3rd Party Imports:
import pandas as pd
import requests
def aggregateHashes(executions_json) -> pd.DataFrame:
"""
Takes the executions, aggregates all the data with sha256 as primary, then returns aggregated dataframe
"""
data = json.loads(executions_json)
df = pd.DataFrame(data["response"]["exechistories"])
if df.empty:
return df
print(df)
# Aggregate by sha256, deduplicate lists, and preserve order
agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()
# Add a column for the number of unique hostnames
agg_df["num_devices"] = agg_df["hostname"].apply(len)
# Sort by num_devices in descending order
agg_df = agg_df.sort_values("num_devices", ascending=False)
return agg_df
def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
"""
Takes output of aggregatedHashes, queries API for those hashes, flattens response while keeping one row per hash,
aggregate applications and baselines into lists, then merges results back into agg_df to create a
"""
if 'sha256' not in agg_df.columns or agg_df.empty:
print("⚠️ 'sha256' column missing or DataFrame is empty. Skipping API query.")
return agg_df.copy() # Return as-is to avoid breaking downstream logic
endpoint = url + '/v1/hash/query'
payload = {
"hashes": agg_df['sha256'].tolist()
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
data = response.json()
results = data.get("response", {}).get("results", [])
rows = []
for res in results:
row = {"sha256": res.get("sha256"), "result": res.get("result")}
if "data" in res:
d = res["data"]
for key in ["filename", "filepath", "description", "filesize", "md5",
"productname", "productversion", "publisher", "createtime", "modtime",
"sha128", "sha384", "sha512", "datetime"]:
row[key] = d.get(key)
row["applications"] = d.get("applications", [])
row["baselines"] = d.get("baselines", [])
reputation = d.get("reputation", {})
for k, v in reputation.items():
row[f"reputation_{k}"] = v
rows.append(row)
df_api = pd.DataFrame(rows)
if 'sha256' not in df_api.columns:
print("⚠️ API response missing 'sha256'. Skipping merge.")
return agg_df.copy()
df = agg_df.merge(df_api, on="sha256", how="left")
# Only include columns that exist to avoid KeyErrors
expected_columns = ['policy','sha256', 'filename_x', 'description', 'productname', 'productversion',
'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username',
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']
available_columns = [col for col in expected_columns if col in df.columns]
aug_df = df[available_columns]
return aug_df
def categorizeHashes(df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
if untrusted_publishers is None: untrusted_publishers = []
if pups is None: pups = []
def reputationtool(row):
val = row["reputation_scannermatch"]
if pd.isna(val) or val == "N/A":
return row["publisher"] == "Not Signed"
try:
return int(val) > threat_tolerance
except (ValueError, TypeError):
return row["publisher"] == "Not Signed"
df["reputation_flag"] = df.apply(reputationtool, axis=1)
mask_needsreview = (
((df["publisher"] == "Not Signed") & df["reputation_flag"]) |
(df["reputation_status"] == "UNKNOWN")
)
mask_approved = (
(
(df["publisher"] != "Not Signed") &
~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
~df["reputation_status"].isna() &
~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
) |
(
(df["publisher"] == "Not Signed") &
~df["reputation_flag"] &
~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
~df["reputation_status"].isna() &
~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
)
)
needsreview_df = df[mask_needsreview]
approved_df = df[mask_approved]
unapproved_df = df[~(mask_needsreview | mask_approved)]
return needsreview_df, approved_df, unapproved_df
def combineHashAndHist(hash_path, condensed_path):
# Load both datasets
condensed_combo = pd.read_parquet(condensed_path)
df = pd.read_parquet(hash_path)
# Merge on sha256
df = pd.merge(condensed_combo, df, on='sha256', how='inner')
# Rename and reorder columns
df = df.rename(columns={'publisher_x': 'publisher'})
df = df.rename(columns={'policy_x': 'policy'})
df = df[['policy','sha256', 'publisher', 'description', 'filename', 'hostname', 'username',
'productname', 'productversion', 'reputation_lastseen', 'reputation_scannermatch',
'reputation_scannercount', 'reputation_status', 'reputation_threatlevel',
'reputation_threatname', 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']]
df = df.sort_values(by='filename')
# Overwrite the original hash file
df.to_parquet(hash_path, index=False)
# Cleanup
del df
del condensed_combo
gc.collect()
def combineHashes(url, parquet_files) -> pd.DataFrame:
combined_hashes = pd.DataFrame()
hashes = []
for file_path in parquet_files:
try:
hash_df = pd.read_parquet(file_path)
pathf.inspect_parquet(file_path)
if not hash_df.empty:
hashes.append(hash_df)
else:
print(f"⚠️ Dataframe is empty: {file_path}")
except Exception as e:
print(f"❌ Error reading Parquet file '{file_path}': {e}")
if hashes:
combined_hashes = pd.concat(hashes, ignore_index=True)
print(f"✅ Combined {len(combined_hashes)} hashes from {len(hashes)} files.")
else:
print("⚠️ No valid dataframes to combine.")
combined_hashes = combined_hashes.drop_duplicates(subset=['sha256'])
augmented_combo = augmentAggregatedHashes(url, combined_hashes)
numeric_reputation_cols = [
'reputation_scannermatch',
'reputation_scannercount',
'reputation_threatlevel'
]
for col in numeric_reputation_cols:
if col in augmented_combo.columns:
augmented_combo[col] = pd.to_numeric(augmented_combo[col].replace('N/A', pd.NA), errors='coerce')
augmented_combo = augmented_combo.rename(columns={'publisher_x': 'publisher'})
augmented_combo = augmented_combo[['sha256', 'publisher', 'description', 'productname', 'productversion',
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
'reputation_timestamp']]
augmented_combo = augmented_combo.sort_values(by=['publisher', 'description', 'productname'])
del combined_hashes
gc.collect()
print(ct.colorText("Hash reputation info added to dataframe", "green"))
return augmented_combo
def condenseExecutions(parquet_paths):
combined_df = pd.DataFrame()
valid_files = []
for file_path in parquet_paths:
try:
df = pd.read_parquet(file_path)
# Optional: pathf.inspect_parquet(file_path)
if not df.empty:
combined_df = pd.concat([combined_df, df], ignore_index=True)
valid_files.append(file_path)
print(f"✅ Loaded {len(df)} rows from {file_path}")
else:
print(f"⚠️ DataFrame from '{file_path}' is empty.")
except Exception as e:
print(f"❌ Error reading Parquet file '{file_path}': {e}")
if not combined_df.empty:
print(f"✅ Combined {len(combined_df)} rows from {len(valid_files)} files.")
else:
print("⚠️ No valid dataframes to combine.")
return combined_df
def divideSortedHashExecutions(unknown_parq, good_parq, bad_parq, condensed_parq, pups) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
# Run combineHashAndHist on each file
combineHashAndHist(unknown_parq, condensed_parq)
combineHashAndHist(good_parq, condensed_parq)
combineHashAndHist(bad_parq, condensed_parq)
# Load data
unknown = pd.read_parquet(unknown_parq)
good = pd.read_parquet(good_parq)
bad = pd.read_parquet(bad_parq)
# Build regex pattern once
pattern = pathf.regulator(pups)
# Move matching rows from unknown and good to bad
bad = pd.concat([
bad,
unknown[unknown["filename"].str.contains(pattern, na=False)],
good[good["filename"].str.contains(pattern, na=False)]
], ignore_index=True)
# Remove matching rows from unknown and good
unknown = unknown[~unknown["filename"].str.contains(pattern, na=False)]
good = good[~good["filename"].str.contains(pattern, na=False)]
return unknown, good, bad
def generatePreflights(hashes, primary_path, secondary_path):
all_hashes = pd.read_parquet(hashes)
primarypathexclusions = ct.tryToReadCSV(primary_path)
secondarypathexclusions = ct.tryToReadCSV(secondary_path)
pathexclusions = pd.concat([primarypathexclusions, secondarypathexclusions], ignore_index=True)
allowbyhash = all_hashes[~all_hashes['sha256'].isin(pathexclusions['sha256'])]
allowbyhash.sort_values(by=["filename"])
return pathexclusions, allowbyhash
def generatePublist(all_hashes, bad_publisher_list):
all_approved_hashes = ct.tryToReadParquet(all_hashes)
#Drop all not signed, only keep unique values
publist = all_approved_hashes[all_approved_hashes['publisher'] != "Not Signed"].drop_duplicates(subset=['publisher'])
#Remove Bad publisher if somehow they made it this far
pattern = pathf.regulator(bad_publisher_list)
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
return publist
utils/localapproval.py
-326
View File
@@ -1,326 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.clientfunctions as clientf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
import utils.utils as ct
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid
#Standard Libary Imports:
import json
import os
import re
import time
import datetime
#3rd Party Imports:
import pandas as pd
import numpy as np
import requests
def getLocalApprovals(url):
endpoint = url + f'/v1/otp/usage'
payload = {
"status" : "0"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
local_approval = pd.DataFrame(result["response"]["otpusage"])
if os.path.exists("Local_Approval\\PARQ\\newest_local_approval.parquet"):
previous_run = pd.read_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet")
previous_run.to_parquet("Local_Approval\\PARQ\\last_local_approval.parquet", index=False)
os.remove("Local_Approval\\PARQ\\newest_local_approval.parquet")
#Only keep rows presumably created by the generate local approval function
local_approval = local_approval[local_approval['purpose'].str.startswith('🎫 Local Approval 🎫')]
local_approval['batchid'] = local_approval['purpose'].apply(lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1))
if not local_approval.empty:
ct.style_dataframe_dark(local_approval, f"Local_Approval\\HTML\\newest_local_approval.html")
local_approval.to_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet", index=False)
return local_approval
def scheduleAddingLAHashes(url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant):
try:
register_function("add_hash", returnFromLocalApproval)
register_function("move_device", clientf.moveAgentToEnforcement)
except Exception as e:
print(f"[ERROR] Failed to register functions: {e}")
return
try:
approvals_df = getNewLocalApprovals(url)
if approvals_df.empty:
print("[INFO] No new local approvals found. Nothing to schedule.")
return
batches = approvals_df.groupby('batchid')
except Exception as e:
print(f"[ERROR] Failed to retrieve or group local approvals: {e}")
return
for batchid, batch_df in batches:
try:
duration_minutes = int(batch_df['duration'].iloc[0])
start_time = datetime.datetime.now()
run_time = start_time + datetime.timedelta(minutes=duration_minutes)
early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))
early_timestamp = early_time.timestamp()
run_timestamp = run_time.timestamp()
# Schedule add_hash job
try:
run_once_job(
f"add_hash_{batchid}",
"add_hash",
early_timestamp,
[url, batch_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant],
None
)
print(f"[INFO] Scheduled add_hash for batch {batchid} at {early_time}")
except Exception as e:
print(f"[ERROR] Failed to schedule add_hash for batch {batchid}: {e}")
# Schedule move_device jobs
devices = batch_df['agentid'].drop_duplicates().tolist()
for device in devices:
try:
run_once_job(
f"move_device_{device}_{batchid}",
"move_device",
run_timestamp,
[url, device, policy_relationship_map],
None
)
print(f"[INFO] Scheduled move_device for device {device} in batch {batchid} at {run_time}")
except Exception as e:
print(f"[ERROR] Failed to schedule move_device for device {device} in batch {batchid}: {e}")
except Exception as e:
print(f"[ERROR] Failed to process batch {batchid}: {e}")
def returnFromLocalApproval(url, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant):
# Get unique policy names from device list
policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
# Create inverse map to go from Audit to Enforcement
inverse_map = {v: k for k, v in policy_relationship_map.items()}
# Fetch all policies
all_policies = policyf.getPolicyDataframe(url)
# Define policy types
policy_types = [1, 2, 6, 7]
# Define output directories
parq_base_dir = "Local_Approval\\PARQ\\"
needappr_base_dir = "Local_Approval\\PARQ\\"
# Ensure output directories exist
os.makedirs(parq_base_dir, exist_ok=True)
os.makedirs(needappr_base_dir, exist_ok=True)
# Build execution history
policyf.buildExecHistory(
url,
policies_in_devicelist,
parq_base_dir,
needappr_base_dir,
policy_types,
threat_tolerance_constant,
bad_publisher_list,
pups,
)
# Read hash data
unknown_hashes = ct.tryToReadCSV(f"{needappr_base_dir}unknown_hashes.csv")
good_hashes = ct.tryToReadCSV(f"{needappr_base_dir}good_hashes.csv")
hashes = pd.concat([unknown_hashes, good_hashes], ignore_index=True)
# Process each policy
for policy in policies_in_devicelist:
# Filter for matching policy
matching_rows = all_policies[all_policies['name'] == policy]
if matching_rows.empty:
print(f"Warning: No group ID found for policy '{policy}'. Skipping.")
continue
# Extract group ID
policy_id = matching_rows['groupid'].values[0]
# Map to destination ID
destination_id = inverse_map.get(policy_id)
if destination_id is None:
print(f"Warning: No corresponding enforcement policy found for group ID '{policy_id}'. Skipping.")
continue
allowlist = policyf.getDestAllowlist(url, destination_id)
policyf.addHash(url, allowlist, hashes[hashes['group'] == policy_id])
def moveToLocalApproval(url, policy_relationship_map):
possible_durations = [15, 60, 360, 1440, 10080]
duration_selected = None
print(ct.colorText("Please select a duration:", "white"))
for i, option in enumerate(possible_durations, start=1):
print(f"{i}. {option}")
try:
choice = int(input("Enter the number of your choice: "))
if 1 <= choice <= len(possible_durations):
duration_selected = possible_durations[choice - 1]
print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
else:
print(ct.colorText("❌ Invalid choice.", "red"))
return
except ValueError:
print(ct.colorText("❌ Invalid input. Please enter a number.", "red"))
return
devicelist = clientf.promptForDevices()
device_df = clientf.findAgents(url, devicelist, True)
batch = int(time.time())
if device_df is None or device_df.empty:
print(ct.colorText("❌ No agents found or error retrieving agents.", "red"))
return
for row in device_df.itertuples(index=False):
try:
addLocalApproval(url, batch, duration_selected, row.agentid)
clientf.moveAgentToAudit(url, row.agentid, policy_relationship_map)
except Exception as e:
print(ct.colorText(f"❌ Error processing agent {row.agentid}: {e}", "red"))
def addLocalApproval(url, batchid, duration_selected, agentid):
purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
endpoint = url + '/v1/otp/retrieve'
payload = {
"duration": str(duration_selected),
"agentid": str(agentid),
"purpose": purpose
}
headers = {"X-APIKey": os.getenv('APIKEY')}
response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)
try:
result = response.json()
otpcode = result["response"]["otpcode"]
print(ct.colorText(f"The OTP code is: {otpcode}", "yellow"))
except Exception as e:
print(ct.colorText(f"An unexpected error occurred for agent {agentid}: {str(e)}", "red"))
def monitorAuditStatus(url: str, policy_relationship_map: dict):
# Simulated current agent list
current_agent_list = clientf.findAllAgents(url)
# Load old agent list
old_agent_path = "Local_Approval\\PARQ\\last_agent_list.parquet"
if os.path.exists(old_agent_path):
old_agent_list = pd.read_parquet(old_agent_path)
else:
old_agent_list = pd.DataFrame(columns=current_agent_list.columns)
# Merge on hostname
merged = pd.merge(
old_agent_list[['hostname', 'groupid']],
current_agent_list[['hostname', 'groupid']],
on='hostname',
how='outer',
suffixes=('_old', '_current'),
indicator=True
)
# Reverse map for enforcement
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())
# 1. Newly added
newly_added = merged[merged['_merge'] == 'right_only']
# 2. Same policy
same_policy = merged[
(merged['_merge'] == 'both') &
(merged['groupid_old'] == merged['groupid_current'])
]
# 3. Moved to audit
moved_to_audit = merged[
(merged['_merge'] == 'both') &
(merged['groupid_old'].isin(policy_relationship_map)) &
(merged['groupid_current'] == merged['groupid_old'].map(policy_relationship_map))
]
# 4. Moved to enforcement
moved_to_enforcement = merged[
(merged['_merge'] == 'both') &
(merged['groupid_old'].isin(reverse_policy_map)) &
(merged['groupid_current'] == merged['groupid_old'].map(reverse_policy_map))
]
# 5. Unusual moves
unusual_move = merged[
(merged['_merge'] == 'both') &
(merged['groupid_old'] != merged['groupid_current']) &
merged.apply(lambda row: (row['groupid_old'], row['groupid_current']) not in known_transitions, axis=1)
]
current_agent_list.to_parquet("Local_Approval\\PARQ\\last_agent_list.parquet", index=False)
# Return all five DataFrames
return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
def getNewLocalApprovals(url):
current_la = getLocalApprovals(url)
# Load old approval list
old_la_path = "Local_Approval\\PARQ\\last_la.parquet"
if os.path.exists(old_la_path):
old_la = pd.read_parquet(old_la_path)
else:
old_la = pd.DataFrame(columns=current_la.columns)
# Create composite keys
current_la['key'] = current_la['clientid'].astype(str) + "_" + current_la['granted'].astype(str)
old_la['key'] = old_la['clientid'].astype(str) + "_" + old_la['granted'].astype(str)
# Find new entries
new_entries = current_la[~current_la['key'].isin(old_la['key'])]
# Convert 'granted' to datetime and filter by last 10 minutes
new_entries['granted'] = pd.to_datetime(new_entries['granted'], utc=True, errors='coerce')
ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
recent_entries = new_entries[new_entries['granted'] > ten_minutes_ago]
# Save current approvals for next run
current_la.drop(columns=['key'], inplace=True)
current_la.to_parquet(old_la_path, index=False)
return recent_entries
utils/menus.py
+258
View File
@@ -0,0 +1,258 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
import os
import re
import dotenv
import pandas as pd
import flows.localApproval as la
import services.policyhandler as policyh
from flows.prepPolicy import (
buildPathsandPublishers,
buildPreflights,
selectAllowlists,
selectPolicies,
sortHashes,
)
from flows.quietAgent import findQuietAgents
from services.agenthandler import findAgents
from services.API import AirlockAPIWrapper
from utils.utils import (
areYouSure,
colorText,
displayIntro,
load_env,
open_directory,
printEnforceChecklist,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def menu_main(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
while True:
displayIntro()
# Add Settings, and give option to change working dir
print(colorText("1. ➡️ - Move Device(s) to local approval", "yellow"))
print(colorText("2. 🎫 - OTP", "yellow"))
print(colorText("3. 🔍 - Device Search", "yellow"))
print(colorText("4. 🔇 - Find Quiet Hosts", "yellow"))
print(colorText("5. 🔒 - Prepare Policy For Enforcement", "yellow"))
print(colorText("6. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
print(colorText("F. 📂 - Open Working Directory", "yellow"))
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("Q. 🔚 - Quit", "yellow"))
choice = input(colorText("\nEnter Menu Item: ", "white"))
if choice == "1":
la.moveToLocalApproval(api)
elif choice == "2":
menu_otp(api)
elif choice == "3":
findAgents(api,False)
elif choice == "4":
findQuietAgents(api)
elif choice == "5":
menu_policy_enforce(api)
elif choice == "6":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
if confirmation.strip().upper() == "I AGREE":
policyh.updateAuditPoliciesFromEnforcementPolices(api)
elif choice == "F":
open_directory(working_dir)
elif choice == "S":
menu_settings()
elif choice == "Q":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_policy_enforce(api: AirlockAPIWrapper):
selected_policies = []
destination_policy = []
destination_allowlist = []
processed_paths = []
processed_hashes = []
processed_publishers = []
tested = False
working_dir = load_env("WORKING_DIR")
while True:
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
choice = input(colorText("\nEnter your choice: ", "white"))
if choice == "1":
selected_policies = selectPolicies(api,True)
elif choice == "2":
print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
destination_policy = selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, False)
elif choice == "3":
sortHashes(
api,
selected_policies,
type=[1, 2, 6, 7],
)
elif choice == "4":
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
buildPathsandPublishers(False)
else:
print("File not found. Please make sure it's saved correctly and try again.")
elif choice == "5":
if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
f"{working_dir}\\Approved\\primary_Paths.csv"
):
buildPreflights()
else:
print("File not found. Please make sure it's saved correctly and try again.")
elif choice == "6":
if (
os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv")
and os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv")
and destination_policy
and destination_allowlist
):
print(colorText("These path exclusions would be added to:", "yellow"))
print(destination_policy)
pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\approved_paths.csv")
hashes = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
# Get unique combinations of longestcfp and file_extension
unique_combinations = pathexclusions[
["longestcfp", "file_extension"]
].drop_duplicates()
# Regex to match a Windows drive letter at the start (e.g., C:\)
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
# Build processed paths like \\path\\**.exe or C:\path\**.jar
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
print(processed_paths)
print(colorText("These publishers would added", "yellow"))
if os.path.exists(f"{working_dir}\\Preflight\\approved_publishers.csv"):
publishers = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
if publishers.empty:
print(colorText("The publishers list is empty.", "red"))
else:
processed_publishers = (
publishers[publishers["publisher_hash"] != "Not Signed"]
["publisher_hash"]
.drop_duplicates()
.tolist()
)
print(processed_publishers)
print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist()
print(processed_hashes)
if processed_paths and processed_hashes:
tested = True
elif choice == "7":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
if (
tested
and destination_policy
and destination_allowlist
and confirmation.strip().upper() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
if processed_publishers:
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
elif choice == "F":
open_directory(working_dir)
elif choice == "S":
menu_settings()
elif choice == "Q":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_otp(api: AirlockAPIWrapper):
while True:
print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
print(colorText("1. Generate OTP", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("Q. Return to Main Menu", "cyan"))
choice = input("Enter your choice: ")
if choice == "1":
# TODO generateOTP(api,findAgents()
break
elif choice == "2":
print("You selected Sub-option B")
elif choice == "Q":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")
def menu_settings():
while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
print(colorText("1. Change Working Dir", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("Q. Return to Main Menu", "cyan"))
choice = input("Enter your choice: ")
if choice == "1":
pass #TODO ADD CHANGE WORKDIR CODE
elif choice == "Q":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")
utils/otpfunctions.py
-226
View File
@@ -1,226 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.clientfunctions as clientf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
import utils.utils as ct
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid
#Standard Libary Imports:
import json
import math
import os
import time
#3rd Party Imports:
import pandas as pd
import requests
def getActiveOTP(url):
endpoint = url + f'/v1/otp/usage'
payload = {
"status" : "1"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
otp = pd.DataFrame(result["response"]["otpusage"])
if os.path.exists("OTP\\PARQ\\newest_active_OTP.parquet"):
previous_run = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet")
previous_run.to_parquet("OTP\\PARQ\\old_active_OTP.parquet", index=False)
os.remove("OTP\\PARQ\\newest_active_OTP.parquet")
otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False)
if not otp.empty:
ct.style_dataframe_dark(otp, f"OTP\\HTML\\newest_active_OTP.html")
def getOTPActivities(url, otpid):
endpoint = url + f'/v1/otp/activities'
payload = {"otpid": f"{otpid}"}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
new_data = pd.DataFrame(result["response"]["otpactivities"])
# Define file path
parquet_path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
# Check if file exists and read it
if os.path.exists(parquet_path):
existing_data = pd.read_parquet(parquet_path)
combined_data = pd.concat([existing_data, new_data], ignore_index=True)
combined_data.drop_duplicates(inplace=True)
else:
combined_data = new_data
# Save combined data
combined_data.to_parquet(parquet_path, index=False)
# Optional: generate styled HTML if there's data
if not combined_data.empty:
ct.style_dataframe_dark(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
def monitorOTP(url, pups):
getActiveOTP(url)
old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet"
new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet"
if os.path.exists(old_otp_path):
old_active_OTP = pd.read_parquet(old_otp_path)
else:
old_active_OTP = pd.DataFrame(columns=['otpid']) # Ensure expected column exists
current_active_OTP = pd.read_parquet(new_otp_path)
if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])]
still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
no_longer_OTP = old_active_OTP[~old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
register_function("addhash", addOTPHashes)
for _, row in newly_added.iterrows():
clientid = row['clientid']
duration = (int(row['duration']) * 60)
hostname = row['hostname']
purpose = row ['purpose']
pid = row['otpid']
early = math.floor(duration * .95)
#If newly added to the list - schedule adding the majority of the executions prior to the expiration of OTP period.
run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, [url, clientid, pid, pups], None)
print(f"Processing: {pid} with other data: {row}")
for _, row in still_in_OTP.iterrows():
pid = row['otpid']
#While still in OTP, continue to update activities list
getOTPActivities(url,pid)
for _, row in no_longer_OTP.iterrows():
clientid = row['clientid']
hostname = row['hostname']
purpose = row ['purpose']
pid = row['otpid']
allowlist = clientf.getDestAllowlistFromClientID(url,clientid)
policy, policyid = clientf.getPolicyFromClientID(url,clientid)
"""
Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
"""
getOTPActivities(url,pid)
find_and_prioritize_jobs_by_pid(pid, 1)
addOTPHashes(url, clientid,pid, pups)
finalhashesadded = pd.read_parquet(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
finalhashesadded['policy'] = policy
finalhashesadded['allowlist'] = allowlist
finalhashesadded['added_at'] = time.localtime()
if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"):
df = pd.DataFrame()
df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
history = pd.concat([history, finalhashesadded], ignore_index=True)
ct.style_dataframe_dark(history, f"localapproval_history.html")
history.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
os.remove(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
del finalhashesadded
def addOTPHashes(url, clientid, otpid, pups):
path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
activities = pd.read_parquet(path)
pattern = pathf.regulator(pups)
allowlist = clientf.getDestAllowlistFromClientID(url, clientid)
# Initialize or preserve 'hash_added' column
if "hash_added" not in activities.columns: activities["hash_added"] = None
# Identify rows that should be added (not matching pattern and not already added)
approve_by_hash = activities[
~activities["filename"].str.contains(pattern, na=False) & (activities["hash_added"] != "added")
]
hashes_to_add = approve_by_hash["sha256"].tolist()
# Add hashes to policy
if hashes_to_add:
policyf.addHash(url, allowlist, hashes_to_add)
# Update 'hash_added' column
activities["hash_added"] = activities.apply(
lambda row: "do not add" if pd.notna(row["filename"]) and pattern in row["filename"]
else ("added" if row["sha256"] in hashes_to_add else row["hash_added"]),
axis=1
)
# Save the updated DataFrame
activities.to_parquet(path)
def generateOTP(url, agentid):
purpose = input(ct.colorText(" Please enter the purpose for the OTP: ", "white"))
possible_durations = [15, 60, 360, 1440, 10080]
duration_selected = " "
print(ct.colorText("Please select a duration:", "white"))
for i, option in enumerate(possible_durations, start=1):
print(f"{i}. {option}")
try:
choice = int(input("Enter the number of your choice: "))
if 1 <= choice <= len(possible_durations):
duration_selected = possible_durations[choice - 1]
print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
else:
print(ct.colorText("Invalid choice.", "red"))
except ValueError:
print(ct.colorText("Invalid input. Please enter a number.", "red"))
endpoint = url + '/v1/otp/retrieve'
payload = {
"duration" : f"{duration_selected}",
"agentid" : f"{agentid}",
"purpose" : f"{purpose}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
otpcode = result["response"]["otpcode"]
print(ct.colorText(f"The OPT code is: {otpcode}", "yellow"))
utils/pathfunctions.py
-213
View File
@@ -1,213 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.hashfunctions as hashf
import utils.pathfunctions as pathf
import utils.utils as ct
#Standard Libary Imports:
import os
import re
#3rd Party Imports:
import pandas as pd
def split_filepaths_grouped(df, col="filename", group_parts=4, min_parts=4):
def clean_split(path):
parts = os.path.normpath(path).split(os.sep)
# Remove leading empty strings caused by UNC paths
parts = [p for p in parts if p]
return parts
df = df.copy()
split_paths = df[col].apply(clean_split)
# Filter out paths with fewer than `min_parts` components
df = df[split_paths.apply(lambda parts: len(parts) >= min_parts)].copy()
split_paths = split_paths[df.index] # Update split_paths to match filtered df
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:group_parts]))
grouped = df.groupby("group_key")
new_rows = []
for _, group_df in grouped:
paths = group_df[col].tolist()
split_parts = [clean_split(p) for p in paths]
def longest_common_prefix(paths):
if not paths:
return []
prefix = paths[0]
for path in paths[1:]:
prefix = [a for a, b in zip(prefix, path) if a == b]
if not prefix:
break
return prefix
common_prefix = longest_common_prefix(split_parts)
prefix_str = os.sep.join(common_prefix)
for i, parts in enumerate(split_parts):
filename = parts[-1]
middle = os.sep.join(parts[len(common_prefix):-1]) if len(parts) > len(common_prefix) + 1 else ""
row = group_df.iloc[i].copy()
row["longestcfp"] = prefix_str
row["middle"] = middle
row["filename_only"] = filename
row["file_extension"] = os.path.splitext(filename)[1].lower()
new_rows.append(row)
return pd.DataFrame(new_rows).drop(columns=["group_key"])
def inspect_parquet(path):
try:
df = pd.read_parquet(path)
print(f"✅ Successfully read: {path}")
print(f"📄 Columns: {df.columns.tolist()}")
print(f"🔢 Rows: {len(df)}")
return df
except Exception as e:
print(f"❌ Error reading {path}: {e}")
return pd.DataFrame()
def regulator(paths, case_insensitive=True):
"""
Build a regex pattern that matches any of the given Windows path fragments.
"""
escaped = [re.escape(p) for p in paths]
pattern = "(?:" + "|".join(escaped) + ")"
if case_insensitive:
pattern = "(?i)" + pattern # Add inline case-insensitive flag
print(f"Regulator is providing: {pattern}")
return pattern
def calculatePath(approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split):
if split : dfs_by_policy = [group for _, group in approved_hashes.groupby('policy')]
else : dfs_by_policy = [approved_hashes]
processed_dfs = []
for df in dfs_by_policy:
haslcp = pathf.split_filepaths_grouped(df, "filename", path_exclusion_constant, min_files_for_path)
haslcp = haslcp.drop_duplicates()
forbidden = pathf.regulator(badpathparts, True)
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
print(ct.colorText("Removing forbidden filepaths for path exceptions", "green"))
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
lcp_not_forbidden_review = lcp_not_forbidden[['policy', 'longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']]
unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index()
unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count']
lcp_not_forbidden_review = lcp_not_forbidden_review.merge(unique_sha_counts, on='longestcfp', how='left')
lcp_not_forbidden_review = lcp_not_forbidden_review[lcp_not_forbidden_review['unique_sha256_count'] >= min_files_for_path]
processed_dfs.append(lcp_not_forbidden_review)
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
return pathExclusions
def generatePathReview(unknown, good, badpathparts, path_exclusion_constant, min_files_for_path, split = False):
df1 = ct.tryToReadCSV(unknown)
df2 = ct.tryToReadCSV(good)
all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename'])
primary_path_exclusions = calculatePath(all_approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split)
remaining_hashes = all_approved_hashes[~all_approved_hashes['sha256'].isin(primary_path_exclusions['sha256'])]
secondary_path_exclusions = calculatePath(remaining_hashes, badpathparts, 3, min_files_for_path, split)
remaining_hashes = remaining_hashes[~remaining_hashes['sha256'].isin(secondary_path_exclusions['sha256'])]
return all_approved_hashes, primary_path_exclusions, secondary_path_exclusions, remaining_hashes
def clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir):
"""
Prompts user to choose whether to delete all .parquet files or preserve execution_history ones.
Then deletes .csv, .html, and .parquet files accordingly from specified folders.
"""
# Prompt user
user_input = input("Do you want to delete *all* .parquet files including execution_history ones? (yes/y or no/n): ").strip().lower()
delete_execution_hist = user_input in ["yes", "y"]
folders = [parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir]
for folder in folders:
folder_path = os.path.abspath(folder)
if not os.path.isdir(folder_path):
print(f"Folder not found: {folder_path}")
continue
for filename in os.listdir(folder_path):
file_path = os.path.join(folder_path, filename)
if not os.path.isfile(file_path):
continue
_, ext = os.path.splitext(filename)
# Delete .csv and .html files
if ext in [".csv", ".html"]:
os.remove(file_path)
print(f"Deleted: {file_path}")
# Delete .parquet files based on user choice
elif ext == ".parquet":
if delete_execution_hist or not filename.startswith("execution_history"):
os.remove(file_path)
def generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path, split= False ):
if os.path.exists(f"{needappr_base_dir}unknown_hashes.csv") and os.path.exists(f"{needappr_base_dir}good_hashes.csv"):
all_hashes, primary_paths, secondary_paths, remaining = pathf.generatePathReview(f"{appr_base_dir}unknown_hashes.csv", f"{appr_base_dir}good_hashes.csv", badpathparts,path_exclusion_constant, min_files_for_path, split)
all_hashes.to_parquet(f"{parq_base_dir}all_hashes.parquet", index=False)
dataframes = {
"all_hashes" : all_hashes,
"primary_Paths": primary_paths,
"secondary_Paths": secondary_paths,
"remaining": remaining
}
for name, df in dataframes.items():
df.to_csv(f"{needappr_base_dir}{name}.csv", index=False)
df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html")
publishers = hashf.generatePublist(f"{parq_base_dir}all_hashes.parquet",bad_publisher_list)
publishers.to_csv(f"{needappr_base_dir}publishers.csv", index=False)
publishers.to_parquet(f"{parq_base_dir}publishers.parquet", index=False)
ct.style_dataframe_dark(publishers, f"{needappr_base_dir}publishers.html")
else:
print(ct.colorText(f"Please manually approve hashes prior to this step","red"))
utils/perstscheduler.py
-308
View File
@@ -1,308 +0,0 @@
#Standard Libary Imports:
import json
import os
import time
from typing import Callable, Any, List, Dict
#3rd Party Imports:
import schedule
# File where all jobs are persisted
JOBS_FILE = "scheduling\\jobs.json"
# Ensure directory exists
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)
# Registry of functions that can be scheduled
FUNCTION_MAP: Dict[str, Callable] = {}
# -------------------------------
# Function Registration
# -------------------------------
def register_function(name: str, func: Callable):
"""
Register a function so it can be called by name later.
Example:
register_function("say_hello", say_hello)
"""
FUNCTION_MAP[name] = func
# -------------------------------
# Persistence Helpers
# -------------------------------
def load_jobs() -> List[Dict[str, Any]]:
"""Load jobs from the JSON file, or return [] if none exist."""
if not os.path.exists(JOBS_FILE):
return []
with open(JOBS_FILE, "r") as f:
return json.load(f)
def _atomic_save(path: str, data: Any):
"""Write JSON atomically to avoid partial writes."""
tmp = f"{path}.tmp"
with open(tmp, "w") as f:
json.dump(data, f, indent=4)
os.replace(tmp, path)
def save_jobs(jobs: List[Dict[str, Any]]):
"""Save jobs to the JSON file (overwrite)."""
_atomic_save(JOBS_FILE, jobs)
# -------------------------------
# Uniqueness Helpers
# -------------------------------
def job_in_store(job_id: str) -> bool:
"""Check if a job id exists in the persisted JSON file."""
return any(j.get("id") == job_id for j in load_jobs())
def job_in_scheduler(job_id: str) -> bool:
"""
Check if a job with this tag exists in the in-memory scheduler.
Uses schedule.get_jobs(tag=...) if available, otherwise scans tags.
"""
try:
jobs = schedule.get_jobs(tag=job_id) # schedule >= 1.2.0
return len(jobs) > 0
except TypeError:
# Fallback for older versions
return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)
def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
"""
Ensure the job_id is unique across persistence and in-memory schedule.
on_conflict:
- "error": raise ValueError if exists.
- "skip" : print and return False.
- "replace": remove existing (in-memory + JSON), then continue.
"""
exists = job_in_store(job_id) or job_in_scheduler(job_id)
if not exists:
return True
if on_conflict == "error":
raise ValueError(f"Job id '{job_id}' already exists.")
elif on_conflict == "skip":
print(f"[INFO] Job '{job_id}' already exists. Skipping creation.")
return False
elif on_conflict == "replace":
# Clear from scheduler
schedule.clear(job_id)
# Remove from persistence
jobs = [j for j in load_jobs() if j.get("id") != job_id]
save_jobs(jobs)
return True
else:
raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")
# -------------------------------
# Internal scheduling (no persistence)
# -------------------------------
def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
args = args or []
kwargs = kwargs or {}
def job_wrapper():
"""Executes the job once, then removes it."""
if func_name not in FUNCTION_MAP:
print(f"[ERROR] Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
# Remove from persistence
jobs = load_jobs()
jobs = [j for j in jobs if j["id"] != job_id]
save_jobs(jobs)
# Clear from in-memory schedule
schedule.clear(job_id)
delay_seconds = run_at_timestamp - time.time()
if delay_seconds <= 0:
print(f"[WARN] Job {job_id} scheduled in the past. Skipping.")
return
# Schedule via schedule library
schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)
def _schedule_recurring(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None):
args = args or []
kwargs = kwargs or {}
def job_wrapper():
if func_name not in FUNCTION_MAP:
print(f"[ERROR] Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
if unit == "seconds":
schedule.every(interval).seconds.do(job_wrapper).tag(job_id)
elif unit == "minutes":
schedule.every(interval).minutes.do(job_wrapper).tag(job_id)
elif unit == "hours":
schedule.every(interval).hours.do(job_wrapper).tag(job_id)
elif unit == "days":
schedule.every(interval).days.do(job_wrapper).tag(job_id)
else:
raise ValueError(f"Unsupported unit: {unit}")
# -------------------------------
# Public APIs (with uniqueness + persistence)
# -------------------------------
def run_once_job(
job_id: str,
func_name: str,
run_at_timestamp: float,
args=None,
kwargs=None,
*,
replace: bool = False,
persist: bool = True,
):
"""
Schedule a job to run once at a specific timestamp.
replace: if True, replace existing job with same id; otherwise print and skip.
persist: if False, do not write to JSON (used by reload_jobs()).
"""
if persist:
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
else:
if job_in_scheduler(job_id):
schedule.clear(job_id)
_schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
jobs.append({
"id": job_id,
"type": "once",
"run_at": run_at_timestamp,
"function": func_name,
"args": args or [],
"kwargs": kwargs or {}
})
save_jobs(jobs)
def recurring_job(
job_id: str,
func_name: str,
interval: int,
unit: str,
args=None,
kwargs=None,
*,
replace: bool = False,
persist: bool = True,
):
"""
Schedule a recurring job.
replace: if True, replace existing job with same id; otherwise print and skip.
persist: if False, do not write to JSON (used by reload_jobs()).
"""
if persist:
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
else:
if job_in_scheduler(job_id):
schedule.clear(job_id)
_schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
jobs.append({
"id": job_id,
"type": "recurring",
"interval": interval,
"unit": unit,
"function": func_name,
"args": args or [],
"kwargs": kwargs or {}
})
save_jobs(jobs)
def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0):
"""
Find all jobs whose ID contains the given PID substring and reschedule them to run sooner.
"""
jobs = load_jobs()
matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]
if not matched_jobs:
print(f"[INFO] No jobs found containing PID substring '{pid_substring}'.")
return
print(f"[INFO] Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
for job in matched_jobs:
job_id = job["id"]
print(f" - Prioritizing job: {job_id}")
# Clear existing job from scheduler
schedule.clear(job_id)
# Reschedule based on job type
if job["type"] == "once":
run_once_job(
job_id,
job["function"],
time.time() + new_delay_seconds,
job.get("args"),
job.get("kwargs"),
replace=True,
persist=True
)
elif job["type"] == "recurring":
recurring_job(
job_id,
job["function"],
job["interval"],
job["unit"],
job.get("args"),
job.get("kwargs"),
replace=True,
persist=True
)
else:
print(f"[WARN] Unknown job type for job '{job_id}'")
# -------------------------------
# Reload Saved Jobs
# -------------------------------
def reload_jobs():
"""Reload jobs from JSON and reschedule them (no re-persist)."""
jobs = load_jobs()
for job in jobs:
if job["type"] == "once":
if job["run_at"] > time.time():
run_once_job(
job["id"], job["function"], job["run_at"],
job.get("args"), job.get("kwargs"),
persist=False
)
elif job["type"] == "recurring":
recurring_job(
job["id"], job["function"], job["interval"], job["unit"],
job.get("args"), job.get("kwargs"),
persist=False
)
# -------------------------------
# Scheduler Loop
# -------------------------------
def start_scheduler():
"""
Start the scheduler loop (blocking).
Call this once in main to begin.
"""
try:
while True:
schedule.run_pending()
time.sleep(0.5)
except KeyboardInterrupt:
print("[INFO] Scheduler stopped.")
utils/policyfunctions.py
-657
View File
@@ -1,657 +0,0 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Local Imports
import utils.utils as ct
import utils.hashfunctions as hashf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
#Standard Libary Imports:
import datetime
import gc
import json
import os
import re
import sys
#3rd Party Imports:
import ijson
import pandas as pd
import requests
import tqdm
from bson import ObjectId
def addHash(url, policy, hash):
print(f"Adding the following: {hash} \n to {policy}:")
for p in hash:
pass
# print(p)
def addPath(url, policy, hash):
print(f"Adding the following Path Exclusions to {policy}:")
for p in hash:
print(p)
def addPub(url, policy, publist):
print(f"Adding the following Publishers to {policy}:")
for p in publist:
print(p)
def addHashReal(url, allowlistID, hashlist):
endpoint = url + '/v1/hash/application/add'
payload = {
"applicationid" : allowlistID,
"hashes" : hashlist
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
payload = json.dumps(payload)
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
response.raise_for_status() # Raise an error for bad status codes
parse_text = json.loads(response.text)
print(parse_text)
def addPathReal(url, grouplistID, pathlist):
endpoint = url + '/v1/group/path/add'
payload = {
"groupid" : grouplistID,
"path" : pathlist
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
print(payload)
payload = json.dumps(payload)
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
print(response.text)
def addPubReal(url, grouplistID, publist):
endpoint = url + '/v1/group/publisher/add'
payload = {
"groupid" : grouplistID,
"publisher" : publist
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
print(payload)
payload = json.dumps(payload)
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
print(response.text)
def getPolicyInfo(url, policy, type, days, parquet=True):
executionhist_policy = pd.DataFrame()
exehist = pullPolicyExechistories(url, policy, type, days, True)
if exehist is not None:
data = json.loads(exehist)
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
if not executionhist_policy.empty:
executionhist_policy = executionhist_policy[['datetime','sha256', 'publisher', 'filename', 'hostname', 'username', 'pprocess', 'gprocess', 'commandline']]
executionhist_policy['policy'] = policy # Add policy column here
executionhist_policy = executionhist_policy.drop_duplicates(subset=['sha256', 'filename', 'hostname'])
executionhist_policy = executionhist_policy.sort_values(by=['sha256', 'filename'])
if parquet:
executionhist_policy.to_parquet(f"prepare_policy\\parquet\\execution_history_{policy}.parquet", index=False)
print(ct.colorText(f"Staging of Execution history for policy: {policy} is complete", "green"))
del data
del exehist
gc.collect()
return executionhist_policy
def sendToPolicy(url, paths, hashes, publishers, destination_name, destination_id, allowlist_name, allowlist_id):
pathexclusions = pd.read_parquet(paths)
allowbyhash = pd.read_parquet(hashes)
publishers = ct.tryToReadCSV(publishers)
ct.areYouSure()
confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white"))
if confirmation.strip().upper() == "I AGREE":
print(ct.colorText("Proceeding with the code...", "yellow"))
print(ct.colorText(f"Adding path exclusions to {destination_name}", "yellow"))
# Get unique combinations of longestcfp and file_extension
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
# Regex to match a Windows drive letter at the start (e.g., C:\)
drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\')
# Build processed paths like \\path\\**.exe or C:\path\**.jar
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
addPathReal(url, destination_id,processed_paths)
print(ct.colorText(f"Adding publishers to {destination_name}", "yellow"))
if publishers.empty:
print(ct.colorText("The publishers list is empty.", "red"))
else:
publisher_list = publishers['publisher'].tolist()
addPubReal(url, destination_id, publisher_list)
print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow"))
allowlist = allowbyhash['sha256'].unique().tolist()
addHash(url, allowlist_id,allowlist)
ct.locked()
exit()
else:
print(ct.colorText("Operation aborted. You MUST EXPLICITLY AGREE to proceed.", "red"))
def pullPolicyExechistories(url, policiesnames, type, days, outputjson: bool):
file_path = 'chunkinator.json'
if not os.path.exists(file_path):
with open(file_path, 'w') as file:
json.dump({'error': 'Success', 'response': {'exechistories': []}}, file)
print(f"File '{file_path}' has been created.")
else:
print(f"File '{file_path}' already exists.")
headers = {"X-APIKey": os.getenv('APIKEY')}
checkpoint = str(skipback(days))
json_output = {'error': 'Success', 'response': {'exechistories': []}}
with tqdm.tqdm(file=sys.stdout, leave=True, total=10000, desc=f"Checkpoint Progess: {checkpoint}", colour="blue", initial=1) as filebar:
with tqdm.tqdm(file=sys.stdout, leave=True, total=100, desc=f"Total of {policiesnames} Complete: ") as pbar:
while True:
item = {}
json_response_data = checkpoint_stomper(checkpoint, url, type, policiesnames, headers)
histories = json_response_data['response']['exechistories']
filebar.total=len(histories)
if not histories:
break
match_found = True
if match_found == True:
for index, item in enumerate(histories):
if index == len(histories) - 1:
checkpoint = item['checkpoint']
filebar.desc = f"Checkpoint Progress: {checkpoint}"
break
else:
if (datetime.date.today() - datetime.timedelta(days=days) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()):
pass
else: json_output['response']['exechistories'].append(item)
filebar.update(1)
filebar.refresh()
seen = {}
if os.path.exists(file_path):
with open(file_path, 'r') as file:
existing_data = json.load(file)
combined = existing_data['response']['exechistories'] + json_output['response']['exechistories']
else:
combined = json_output['response']['exechistories']
for item in combined:
key = (item.get('sha256'), item.get('filename'), item.get('hostname'))
seen[key] = item
deduplicated = list(seen.values())
with open(file_path, 'w') as file:
json.dump({'error': 'Success', 'response': {'exechistories': deduplicated}}, file)
json_output['response']['exechistories'].clear()
date_diff = datetime.date.today() - datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()
percentage_diff = (((days + 10) - date_diff.days) / (days + 10)) * 100
pbar.n = round(percentage_diff)
pbar.set_description_str(f"Total of {policiesnames} Complete: ")
pbar.refresh()
filebar.n = 1
with open(file_path, 'r') as file:
final_output = json.load(file)
os.remove(file_path)
return json.dumps(final_output) if outputjson else None
def checkpoint_stomper(checkpoint, url, type, policy, headers):
json_output = {'error': 'Success', 'response': {'exechistories': []}}
endpoint = url + '/v1/logging/exechistories'
payload_dict = {
"type":[type],
"checkpoint": checkpoint,
"policy": [policy]
}
payload = json.dumps(payload_dict)
with requests.request("POST", endpoint, headers=headers, data=payload, verify=False, stream=True) as response:
parser = ijson.items(response.raw, 'response.exechistories.item')
for item in parser:
key = (item.get('sha256'), item.get('hostname'))
if key not in json_output:
json_output['response']['exechistories'].append(item)
parse_text = json.loads(json.dumps(json_output))
return parse_text
def listPolicies(url):
endpoint = url + '/v1/group'
print(ct.colorText("[+] Grabbing All Policies", "cyan"))
payload = {}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
parse_text = json.loads(response.text)
return parse_text
def choosePolicies(url):
allpolicies = listPolicies(url)
policiesnames = []
policyids = []
for index, list in enumerate(allpolicies['response']['groups'], start=1):
print(ct.colorText(f"{index}. {list['name']}", "yellow"))
policiesnames.append(list['name'])
policyids.append(list['groupid'])
choice = input(ct.colorText("Select Policy Group: ", "white"))
choice = int(choice) - 1
return choice, policiesnames, policyids
def getPolicyDataframe(url) -> pd.DataFrame:
return pd.DataFrame(listPolicies(url)['response']['groups'])
def listATPolicies(url):
endpoint = url + '/v1/group'
print(ct.colorText("[+] Grabbing All Policies", "cyan"))
headers = {
"X-APIKey": os.getenv('APIKEY')
}
try:
response = requests.post(endpoint, headers=headers, json={}, verify=False)
response.raise_for_status()
parse_text = response.json()
at_policies = {}
for index, group in enumerate(parse_text.get('response', {}).get('groups', []), start=1):
name = group.get('name', '')
if "AT" in name:
print(ct.colorText(f"{index}. {name}", "yellow"))
at_policies[name] = group.get('groupid')
return at_policies
except requests.exceptions.RequestException as e:
print(ct.colorText(f"[!] Request failed: {e}", "red"))
return {}
except (KeyError, json.JSONDecodeError) as e:
print(ct.colorText(f"[!] Failed to parse response: {e}", "red"))
return {}
def listAllowlists(url: str) -> tuple[int, list, list]:
endpoint = url + '/v1/application'
print(ct.colorText("[+] Grabbing All Allowlists", "cyan"))
payload = {}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
parse_text = json.loads(response.text)
policiesnames = []
policyids = []
for index, item in enumerate(parse_text['response']['applications'], start=1):
if index >= 38:
print(ct.colorText(f"{index}. {item['name']}", "yellow"))
policiesnames.append(item['name'])
policyids.append(item['applicationid'])
while True:
try:
choice = int(input(ct.colorText("Select allowlist: ", "white")))
if choice < 38 or choice > len(parse_text['response']['applications']):
print(ct.colorText("Please only choose an allowlist designed for this use - '38+'", "red"))
else:
adjusted_choice = choice - 38
return adjusted_choice, policiesnames, policyids
except ValueError:
print(ct.colorText("Invalid input. Please enter a number.", "red"))
def skipback(days):
"""
Generate a MongoDB ObjectId for a given number of days ago from today.
"""
adjusted_days = days
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=adjusted_days)
timestamp = int(date_days_ago.timestamp())
hex_timestamp = format(timestamp, '08x')
objectid_hex = hex_timestamp + '0000000000000000'
return ObjectId(objectid_hex)
def sendToPolicyTest(url, paths, hashes, publishers, destination_name, destination_id, allowlist_id, allowlist_name):
pathexclusions = pd.read_parquet(paths)
allowbyhash = pd.read_parquet(hashes)
publishers = ct.tryToReadCSV(publishers)
print(ct.colorText(f"These path exclusions would be added to {destination_name}", "yellow"))
# Get unique combinations of longestcfp and file_extension
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
# Regex to match a Windows drive letter at the start (e.g., C:\)
drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\')
# Build processed paths like \\path\\**.exe or C:\path\**.jar
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
addPath(url, destination_id,processed_paths)
print(ct.colorText(f"These publishers would added to {destination_name}", "yellow"))
if publishers.empty:
print(ct.colorText("The publishers list is empty.", "red"))
else:
publisher_list = publishers['publisher'].tolist()
addPub(url, destination_id, publisher_list)
print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow"))
allowlist = allowbyhash['sha256'].unique().tolist()
addHash(url, allowlist_id,allowlist)
def updateAuditPoliciesFromEnforcementPolices(url, policy_relationship_map):
for enforcement_policy, audit_policy in policy_relationship_map.items():
assignPoliciesfromGroup(url, enforcement_policy, audit_policy)
turnOnAudit(url, audit_policy)
def assignPoliciesfromGroup(url, source_policy_id, target_policy_id):
endpoint = url + '/v1/group/assign'
payload = {
"groupid" : {source_policy_id},
"targetgroupid" : {target_policy_id}
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
parse_text = json.loads(response.text)
print(parse_text)
def turnOnAudit(url, policyid):
endpoint = url + '/v1/group/settings/auditmode'
payload = {
"groupid" : {policyid},
"auditmode" : "1"
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
parse_text = json.loads(response.text)
print(parse_text)
def agentsInPolicy(url, policyid):
endpoint = url + '/v1/group/agents'
payload = {
"groupid" : {policyid}
}
headers = {
"X-APIKey": os.getenv('APIKEY')
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
parse_text = json.loads(response.text)
print(parse_text)
def prepare_to_enforce(url, bad_publisher_list, pups, badpathparts, threat_tolerance_constant = 4, path_exclusion_constant = 4, min_files_for_path = 4):
destination_name = " "
destination_id = " "
allowlist_name = " "
allowlist_id = " "
policylist = []
type = [1, 2, 6, 7]
parq_base_dir = "prepare_policy\\parquet\\"
appr_base_dir = "prepare_policy\\approved\\"
needappr_base_dir = "prepare_policy\\needs_approved\\"
pflight_base_dir = "prepare_policy\\preflight\\"
#If the directorys where we're going to store our output dont exist, make them.
os.makedirs(parq_base_dir, exist_ok=True)
os.makedirs(needappr_base_dir, exist_ok=True)
os.makedirs(appr_base_dir, exist_ok=True)
os.makedirs(pflight_base_dir, exist_ok=True)
while True:
ct.printEnforceChecklist(parq_base_dir,appr_base_dir, needappr_base_dir, pflight_base_dir, policylist, allowlist_name, destination_name)
choice = input(ct.colorText("\nEnter your choice: ", "white"))
if choice == "1":
policylist, policyids = getMultiplePolicySelections(url)
elif choice == "2":
print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white"))
choice, policynames, policyid = choosePolicies(url)
#print(allowlist_parent_tuple)
destination_name = policynames[choice]
destination_id = policyid[choice]
print(ct.colorText(f"Please choose Allowlist for Hashes","white"))
choice, allowlists,allowid = listAllowlists(url)
#print(allowlist_parent_tuple)
allowlist_name = allowlists[choice]
allowlist_id = allowid[choice]
print(destination_name, allowlist_name)
elif choice == "3":
policyf.buildExecHistory(url,
policylist,
parq_base_dir,
needappr_base_dir,
type,
threat_tolerance_constant,
bad_publisher_list,
pups,
)
csvs = [f"{needappr_base_dir}unknown_hashes.csv", f"{needappr_base_dir}good_hashes.csv"]
#Since we want to build paths as if they were all in the same policy to begin with, lets group them that way
for csv in csvs:
df = ct.tryToReadCSV(csv)
df['policy'] = destination_name
df.to_csv(csv)
elif choice == "4":
pathf.generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path)
elif choice == "5":
savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir)
elif choice == "6":
if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ":
sendToPolicyTest(
url,
f"{parq_base_dir}final_path_exclusions.parquet",
f"{parq_base_dir}final_hash_approvals.parquet",
f"{appr_base_dir}publishers.parquet",
destination_name,
destination_id,
allowlist_name,
allowlist_id
)
elif choice == "7":
if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ":
sendToPolicy(
url,
f"{parq_base_dir}final_path_exclusions.parquet",
f"{parq_base_dir}final_hash_approvals.parquet",
f"{appr_base_dir}publishers.parquet",
destination_name,
destination_id,
allowlist_name,
allowlist_id,
)
elif choice == "R":
pathf.clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir)
elif choice == "Q":
break
else:
print(ct.colorText("Invalid choice. Please try again.", "red"))
def buildExecHistory(url,
policylist,
parq_base_dir,
needappr_base_dir,
type,
threat_tolerance_constant,
bad_publisher_list,
pups
):
exe_hist_parq_list = []
while True:
try:
history_days = int(input("Enter the how many days in of history do you want to pull - select a number between 1 and 150: "))
if 1 <= history_days <= 150:
break
else:
print("Invalid input. Please enter a number between 1 and 150.")
except ValueError:
print("Invalid input. Please enter a valid integer.")
for policy in policylist:
policy_exec_history = policyf.getPolicyInfo(url, policy, type, history_days)
policy_exec_history['policy'] = policy
policy_exec_history.to_parquet(f"{parq_base_dir}Exec_Hist_{policy}.parquet")
exe_hist_parq_list.append(f"{parq_base_dir}Exec_Hist_{policy}.parquet")
augmented_hashlist = hashf.combineHashes(url, exe_hist_parq_list)
augmented_hashlist.to_parquet(f"{parq_base_dir}augmentedHashlist.parquet",index=False)
needsreview_df, approved_df, unapproved_df = hashf.categorizeHashes(augmented_hashlist, threat_tolerance_constant, bad_publisher_list, pups)
needsreview_df.to_parquet(f"{parq_base_dir}needsreview.parquet",index=False)
approved_df.to_parquet(f"{parq_base_dir}approved.parquet",index=False)
unapproved_df.to_parquet(f"{parq_base_dir}unapproved.parquet",index=False)
condensed_executions = hashf.condenseExecutions(exe_hist_parq_list)
condensed_executions.to_parquet(f"{parq_base_dir}condensed_executions.parquet", index=False)
unknown, good, bad = hashf.divideSortedHashExecutions(
f"{parq_base_dir}needsreview.parquet",
f"{parq_base_dir}approved.parquet",
f"{parq_base_dir}unapproved.parquet",
f"{parq_base_dir}condensed_executions.parquet",
pups
)
dataframes = {
"unknown_hashes" : unknown,
"good_hashes": good,
"bad_hashes": bad
}
for name, df in dataframes.items():
df.to_csv(f"{needappr_base_dir}{name}.csv", index=False)
df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html")
def savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir):
if os.path.exists(f"{appr_base_dir}primary_Paths.csv"):
if not os.path.exists(f"{parq_base_dir}final_hash_approvals.parquet") and not os.path.exists(f"{parq_base_dir}final_path_exclusions.parquet"):
pathexclusions, allowbyhash = hashf.generatePreflights(
f"{parq_base_dir}all_hashes.parquet",
f"{appr_base_dir}primary_Paths.csv",
f"{appr_base_dir}secondary_Paths.csv")
dataframes = {
"final_path_exclusions" : pathexclusions,
"final_hash_approvals": allowbyhash
}
for name, df in dataframes.items():
df.to_csv(f"{pflight_base_dir}{name}.csv", index=False)
df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
ct.style_dataframe_dark(df, f"{pflight_base_dir}{name}.html")
def getMultiplePolicySelections(url):
policynameslist = []
policyidlist = []
while True:
choice, policynames, policyid = listPolicies(url)
selected_policy = policynames[choice]
selected_policyid = policyid[choice]
if selected_policy not in policynameslist:
policynameslist.append(selected_policy)
if selected_policyid not in policyidlist:
policyidlist.append(selected_policyid)
while True:
answer = input(ct.colorText("Do you want to load another policy? (yes/no): ", "white")).strip().lower()
if answer in ("no", "n"):
break # Exit the inner loop and then the outer loop
elif answer in ("yes", "y"):
break # Exit the inner loop and continue the outer loop
else:
print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))
if answer in ("no", "n"):
break
return policynameslist, policyidlist
def getDestAllowlist(url, groupid):
allowlists = getPolicyAllowlists(url, groupid)
matches = allowlists[
allowlists['name'].str.contains('local', case=False, na=False) &
allowlists['name'].str.contains('approval', case=False, na=False),
'applicationid'
].values
app_id = matches[0] if len(matches) > 0 else None
return app_id
def getPolicyAllowlists(url, groupid):
endpoint = url + '/v1/group/policies'
payload = {
"groupid" : f"{groupid}"
}
headers = {"X-APIKey": os.getenv('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
data = pd.DataFrame(result["response"]["applications"])
return data
utils/utils.py
+606 -249
View File
@@ -13,11 +13,597 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#Standard Libary Imports:
import json
import logging
import os
import platform
import re
import subprocess
import tempfile
import tkinter as tk
from tkinter import filedialog, messagebox, simpledialog
from typing import Callable, Optional, TypeVar
import pandas as pd
def colorText(text: str, color: str) -> str:
logger = logging.getLogger(__name__)
T = TypeVar("T")
def load_env_json(key: str, default: str):
raw = os.getenv(key, default)
try:
return json.loads(raw)
except json.JSONDecodeError:
try:
escaped = raw.encode('unicode_escape').decode('utf-8')
return json.loads(escaped)
except Exception as e:
logging.error(f"Failed to parse {key}: {e}")
return json.loads(default)
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
"""
Safely retrieves an environment variable and casts it to the desired type.
Parameters:
key (str): The name of the environment variable.
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
default (Optional[T], optional): Default value if the variable is not set or invalid.
Returns:
Optional[T]: The casted value or the default.
"""
value = os.getenv(key)
if value is None:
logger.warning(f"Environment variable '{key}' not set.")
return default
try:
value = value.strip("'\"") # Strip surrounding quotes
return cast_type(value)
except (ValueError, TypeError):
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
return default
def import_to_dataframe(file_path: str) -> pd.DataFrame:
df = pd.DataFrame()
try:
if not os.path.exists(file_path):
print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
return df
ext = os.path.splitext(file_path)[1].lower()
if ext == ".csv":
df = pd.read_csv(file_path)
elif ext == ".parquet":
df = pd.read_parquet(file_path)
else:
print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
return df
if df.empty:
print(colorText("Error: File has headers but no data rows.", "red"))
else:
print(colorText(f"Data loaded successfully from {file_path}", "green"))
return df
except pd.errors.EmptyDataError:
print(
colorText(
"Notice: CSV file is completely empty, falling back to empty frame",
"white",
)
)
return pd.DataFrame()
except Exception as e:
print(colorText(f"Error reading file: {e}", "red"))
return pd.DataFrame()
def choose_directory():
root = tk.Tk()
root.withdraw() # Hide the main window
directory = filedialog.askdirectory(title="Select a Directory")
print("Selected directory:", directory)
return directory
def choose_file(initial_directory=None, required_substring=None):
"""Open a file dialog and ensure the selected file contains a required substring."""
while True:
root = tk.Tk()
root.withdraw() # Hide the main window
file_path = filedialog.askopenfilename(initialdir=initial_directory)
if not file_path:
print("No file selected.")
return None
if required_substring and required_substring not in file_path:
print(
f"The selected file must contain '{required_substring}' in its path or name. Please try again."
)
else:
return file_path
def choose_save_location():
root = tk.Tk()
root.withdraw()
save_path = filedialog.asksaveasfilename(defaultextension=".txt")
return save_path
def ask_user_input(message):
root = tk.Tk()
root.withdraw()
user_input = simpledialog.askstring("Input", "{message}]:")
return user_input
def show_info_message_cli(title, message):
root = tk.Tk()
root.withdraw()
root.after(100, lambda: messagebox.showinfo(title, message))
root.mainloop()
def show_confirm_question():
root = tk.Tk()
root.withdraw()
response = messagebox.askquestion("Confirm", "Do you want to continue?")
print("User response:", response)
def regulator(paths, case_insensitive=True):
"""
Build a regex pattern that matches any of the given Windows path fragments.
"""
escaped = [re.escape(p) for p in paths]
pattern = "(?:" + "|".join(escaped) + ")"
if case_insensitive:
pattern = "(?i)" + pattern # Add inline case-insensitive flag
print(f"Regulator is providing: {pattern}")
return pattern
def displayIntro():
print(
colorText(
r"""
███
████ ░████████
█████████████ ███████████████
█████████████████████ █████████████████████
███████████████████ ██████████████████████▓
███████████████████ ██████████████████████
█████████████████████ ███████████████████████
████████████████████████████████████████████████████████
█████████ ██ ██ █████████
█████████ ██ ███ █ █████████
█████████ ██ ████ █████ █████████████
█████████ ██ ██████ █████████████
████████ ██ ███████ ████████████░
███████ ██ ██▓ ██████ ████████████
██████ ██ ████ █████ ███████████
█████████████████████████████████████████████████
▒████████████████████ ██████████████████
███████████████████ ███████████████▒
███████████████ █████████████
██████████ ███████████
████████
████
""",
"yellow",
)
)
print(
colorText(
r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
print(
colorText(
"======================== Welcome to the Airlock API Tool ========================",
"cyan",
)
)
print(
colorText(
"=================================================================================",
"cyan",
)
)
def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
working_dir = load_env("WORKING_DIR")
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
"cyan",
)
)
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
"\nSequentually follow these steps to prepare a policy for enforcement:",
"white",
)
)
print(
colorText(
"\n1. Choose which originating policy or policies to move to enforcement",
"cyan",
)
)
if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red"))
else:
print(colorText("The following policies have been choosen:", "green"))
for policy in selected_policies:
print(colorText(f" [✓] {policy.name}", "green"))
print(colorText("2. Choose the destination policy and allowlist", "cyan"))
if not destination_policy:
print(colorText(" [✗] No destination policy has been chosen", "red"))
elif destination_policy:
print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green"))
if not destination_allowlist:
print(colorText(" [✗] No allowlist has been chosen", "red"))
elif destination_allowlist:
print(
colorText(
f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
"green",
)
)
print(
colorText(
"3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
"cyan",
)
)
if not selected_policies:
print(colorText(" [✗] No policies have been chosen", "red"))
else:
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
print(colorText(" [✓] Data has been fetched", "green"))
else:
print(colorText(" [✗] Data has not been fetched", "red"))
print(colorText("4. Manually review the files:", "cyan"))
print(
colorText(
" 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n",
"cyan",
)
)
print(
colorText(
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
"cyan",
)
)
print(
colorText(
" If metarules need to be created, please make note of them, and remove the row from the csv.",
"cyan",
)
)
print(
colorText(
" When complete, save both csv files to the directory 'approved' and choose this option.",
"cyan",
)
)
print(
colorText(
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
"cyan",
)
)
if os.path.exists(f"{working_dir}\\Approved\\approved_executions.csv"):
print(colorText(" [✓] Reviewed hashes have been loaded", "green"))
else:
print(colorText(" [✗] Reviewed hashes have not been loaded", "red"))
if os.path.exists(
f"{working_dir}\\Needs_Review\\Review_Second\\primary_Paths.csv",
):
print(colorText(" [✓] Path review list created", "green"))
else:
print(colorText(" [✗] Path review list has not been created", "red"))
print(
colorText(
"5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful",
"cyan",
)
)
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(
colorText(
" Do the same process with the list of publishers forthe same directories",
"cyan",
)
)
print(colorText(" Preflight Lists will be generated", "cyan"))
if os.path.exists(
f"{working_dir}\\Approved\\primary_Paths.csv",
):
print(colorText(" [✓] Reviewed path list detected", "green"))
else:
print(colorText(" [✗] Path review list has not been detected", "red"))
if os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") and os.path.exists(
f"{working_dir}\\Preflight\\approved_hashes.csv"
):
print(colorText(" [✓] Preflight Path Exclusion List has been generated", "green"))
else:
print(colorText(" [✗] Preflight Path Exclusion List has not been generated", "red"))
print(colorText("6. Test ------------------------------------------------------", "cyan"))
print(colorText(" Print rather than apply selected data.", "cyan"))
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
print(
colorText(
" Apply path exclusions according to allowed and approved paths",
"cyan",
)
)
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(
colorText(
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
"cyan",
)
)
print(colorText("F. 📂 - Open Working Directory", "cyan"))
print(colorText("Q. 🔚 - Quit", "cyan"))
def areYouSure():
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
"yellow",
)
)
print(
colorText(
"🛑========================================================================================================================================🛑",
"red",
)
)
print(
colorText(
"⚠️=========================================================================================================================================⚠️",
"yellow",
)
)
print(
colorText(
"🛑****************************************************************************************************************************************🛑",
"red",
)
)
def locked():
print(
colorText(
r"""
████████████████████████████████████████████████████████████████
███ ██
██ ██████ ███
██ ████████████ ███
██ ████ ███ ███
██ ███ ███ ███
██ ███ ███ ███
██ ▒████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ███
███ ███
████████████████████████████████████████████████████████████████████
▒██████████████████████████████████████████████████████████████████▒
▒████
▒████
▓██████████████████████████████████████████
█████████████████████████████████████████████░
""",
"yellow",
)
)
def printDeviceEnforceChecklist():
print(
colorText(
"\n --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
"cyan",
)
)
print(
colorText(
" --------------------------------------------------------------------",
"cyan",
)
)
print(
colorText(
"\nSequentually follow these steps to prepare a policy for enforcement:",
"white",
)
)
print(
colorText(
"\n1. Choose which originating policy or policies to move to enforcement",
"cyan",
)
)
print(
colorText(
"2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
"cyan",
)
)
print(colorText("3. Manually review the files:", "cyan"))
print(
colorText(
" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
"cyan",
)
)
print(
colorText(
" If metarules need to be created, please make note of them, and remove the row from the csv.",
"cyan",
)
)
print(
colorText(
" When complete, save both csv files to the directory 'approved' and choose this option.",
"cyan",
)
)
print(
colorText(
" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
"cyan",
)
)
print(
colorText(
"4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
"cyan",
)
)
print(
colorText(
" Remove the rows containing path exclusions you do not approve of",
"cyan",
)
)
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(
colorText(
" Do the same process with the list of publishers forthe same directories",
"cyan",
)
)
print(colorText(" Preflight Lists will be generated", "cyan"))
print(colorText("5. Choose the destination policy and parent and child allow list", "cyan"))
print(colorText("6. Test ------------------------------------------------------", "cyan"))
print(colorText(" Print rather than apply selected data.", "cyan"))
print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
print(
colorText(
" Apply path exclusions according to allowed and approved paths",
"cyan",
)
)
print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(
colorText(
"R. Remove/Reset Generated data - will prompt to allow keeping execution history",
"cyan",
)
)
print(colorText("Q. Quit", "cyan"))
def colorText(text, color):
colors = {
"red": "\033[91m",
"green": "\033[92m",
@@ -26,17 +612,17 @@ def colorText(text: str, color: str) -> str:
"magenta": "\033[95m",
"cyan": "\033[96m",
"white": "\033[97m",
"reset": "\033[0m"
"reset": "\033[0m",
}
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
def style_dataframe_dark(df, output_html_path=None, overwrite=True):
def formatHTML(df, output_html_path=None, overwrite=True):
from datetime import datetime
# Get current date and filename for subtitle
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
filename = output_html_path.replace('.html', '') if output_html_path else "Report"
filename = output_html_path.replace(".html", "") if output_html_path else "Report"
dark_css = """
<style>
@@ -147,255 +733,26 @@ def style_dataframe_dark(df, output_html_path=None, overwrite=True):
f.write(styled_html)
print(f"✅ Styled table saved to '{output_html_path}'")
elif overwrite:
import tempfile
temp_path = tempfile.mktemp(suffix=".html")
with open(temp_path, "w", encoding="utf-8") as f:
with tempfile.NamedTemporaryFile(
suffix=".html", delete=False, mode="w", encoding="utf-8"
) as f:
f.write(styled_html)
print(f"✅ Styled table saved to temporary file: {temp_path}")
temp_path = f.name
print(f"✅ Styled table saved to temporary file: {temp_path}")
else:
return styled_html
def displayIntro():
def open_directory(path):
system = platform.system()
print(colorText(r"""
███
████ ░████████
█████████████ ███████████████
█████████████████████ █████████████████████
███████████████████ ██████████████████████▓
███████████████████ ██████████████████████
█████████████████████ ███████████████████████
████████████████████████████████████████████████████████
█████████ ██ ██ █████████
█████████ ██ ███ █ █████████
█████████ ██ ████ █████ █████████████
█████████ ██ ██████ █████████████
████████ ██ ███████ ████████████░
███████ ██ ██▓ ██████ ████████████
██████ ██ ████ █████ ███████████
█████████████████████████████████████████████████
▒████████████████████ ██████████████████
███████████████████ ███████████████▒
███████████████ █████████████
██████████ ███████████
████████
████
""", "yellow"))
print(colorText(r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""", "cyan"))
print(colorText("=================================================================================", "cyan"))
print(colorText("======================== Welcome to the Airlock API Tool ========================", "cyan"))
print(colorText("=================================================================================", "cyan"))
def printEnforceChecklist(parq_base_dir,appr_base_dir, needappr_base_dir, pflight_base_dir, policy_list, allowlist_name, destination_name):
print(colorText("\n --------------------------------------------------------------------", "cyan"))
print(colorText(" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------", "cyan"))
print(colorText(" --------------------------------------------------------------------", "cyan"))
print(colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
print(colorText("\n1. Choose which originating policy or policies to move to enforcement", "cyan"))
if not policy_list:
print(colorText(f" [✗] No policies have been chosen","red"))
else:
print(colorText(f"The following policies have been choosen:", "green"))
for policy in policy_list:
print(colorText(f" [✓] {policy}","green"))
print(colorText(f"2. Choose the destination policy and allowlist", "cyan"))
if allowlist_name == " ":
print(colorText(f" [✗] No allowlists have been chosen","red"))
elif allowlist_name != " " and allowlist_name != " " and allowlist_name is not allowlist_name:
print(colorText(f" [✓] {allowlist_name} has been selected as allowlist","green"))
if destination_name == " ":
print(colorText(f" [✗] No destination policy has been chosen","red"))
if system == "Windows":
os.startfile(path)
elif system == "Linux":
subprocess.run(["xdg-open", path])
else:
print(colorText(f" [✓] destination policy is {destination_name}","green"))
print(colorText("3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", "cyan"))
if not policy_list:
print(colorText(f" [✗] No policies have been chosen","red"))
else:
for policy in policy_list:
if os.path.exists(f"{parq_base_dir}Exec_Hist_{policy}.parquet"): print(colorText(f" [✓] Data for {policy} has been fetched","green"))
else: print(colorText(f" [✗] Data for {policy} has not been fetched","red"))
print(colorText(f"4. Manually review the files:","cyan"))
print(colorText(" 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n", "cyan"))
print(colorText(" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", "cyan"))
print(colorText(" If metarules need to be created, please make note of them, and remove the row from the csv.", "cyan"))
print(colorText(" When complete, save both csv files to the directory 'approved' and choose this option.","cyan"))
print(colorText(" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", "cyan"))
if os.path.exists(f"{appr_base_dir}good_hashes.csv") and os.path.exists(f"{appr_base_dir}unknown_hashes.csv"):
print(colorText(" [✓] Reviewed hashes have been loaded","green"))
else:
print(colorText(" [✗] Reviewed hashes have not been loaded","red"))
if os.path.exists(f"{parq_base_dir}all_hashes.parquet"):
print(colorText(" [✓] The combined approved hashes list has been generated","green"))
else:
print(colorText(" [✗] The combined approved hashes list has not been generated","red"))
if os.path.exists(f"{needappr_base_dir}primary_Paths.csv"):
print(colorText(" [✓] Path review list created","green"))
else:
print(colorText(" [✗] Path review list has not been created","red"))
raise OSError(f"Unsupported operating system: {system}")
print(colorText(f"5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'", "cyan"))
print(colorText(" Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful" , "cyan"))
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(colorText(" Do the same process with the list of publishers forthe same directories", "cyan"))
print(colorText(" Preflight Lists will be generated", "cyan"))
if os.path.exists(f"{appr_base_dir}primary_Paths.csv"):
print(colorText(" [✓] Reviewed path list detected","green"))
else:
print(colorText(" [✗] Path review list has not been detected","red"))
if os.path.exists(f"{pflight_base_dir}final_path_exclusions.csv"):
print(colorText(" [✓] Preflight Path Exclusion List has been generated","green"))
else:
print(colorText(" [✗] Preflight Path Exclusion List has not been generated","red"))
if os.path.exists(f"{pflight_base_dir}final_hash_approvals.csv"):
print(colorText(" [✓] Preflight hash approval list has been generated","green"))
else:
print(colorText(" [✗] Preflight hash approval list has not been generated","red"))
print(colorText(f"6. Test ------------------------------------------------------", "cyan"))
print(colorText(f" Print rather than apply selected data.", "cyan"))
print(colorText(f"7. Liftoff ------------------------------------------------------", "cyan"))
print(colorText(f" Apply path exclusions according to allowed and approved paths", "cyan"))
print(colorText(f" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(f" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(colorText("R. Remove/Reset Generated data - will prompt to allow keeping execution history", "cyan"))
print(colorText("Q. Quit", "cyan"))
def areYouSure():
print(colorText(f"🛑****************************************************************************************************************************************🛑","red"))
print(colorText(f"⚠️=========================================================================================================================================⚠️","yellow"))
print(colorText(f"🛑========================================================================================================================================🛑","red"))
print(colorText(f"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️", "yellow"))
print(colorText(f"🛑========================================================================================================================================🛑","red"))
print(colorText(f"⚠️=========================================================================================================================================⚠️","yellow"))
print(colorText(f"🛑****************************************************************************************************************************************🛑","red"))
def locked():
print(colorText(r"""
████████████████████████████████████████████████████████████████
███ ██
██ ██████ ███
██ ████████████ ███
██ ████ ███ ███
██ ███ ███ ███
██ ███ ███ ███
██ ▒████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ██████████████████████ ███
██ ███
███ ███
████████████████████████████████████████████████████████████████████
▒██████████████████████████████████████████████████████████████████▒
▒████
▒████
▓██████████████████████████████████████████
█████████████████████████████████████████████░
""", "yellow"))
def printDeviceEnforceChecklist():
print(colorText("\n --------------------------------------------------------------------", "cyan"))
print(colorText(" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------", "cyan"))
print(colorText(" --------------------------------------------------------------------", "cyan"))
print(colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
print(colorText("\n1. Choose which originating policy or policies to move to enforcement", "cyan"))
print(colorText("2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", "cyan"))
print(colorText(f"3. Manually review the files:","cyan"))
print(colorText(" 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'", "cyan"))
print(colorText(" Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", "cyan"))
print(colorText(" If metarules need to be created, please make note of them, and remove the row from the csv.", "cyan"))
print(colorText(" When complete, save both csv files to the directory 'approved' and choose this option.","cyan"))
print(colorText(" This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", "cyan"))
print(colorText(f"4. Manually review the file 'needs_approved\\paths_needing_review.csv'", "cyan"))
print(colorText(" Remove the rows containing path exclusions you do not approve of" , "cyan"))
print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
print(colorText(" Do the same process with the list of publishers forthe same directories", "cyan"))
print(colorText(" Preflight Lists will be generated", "cyan"))
print(colorText(f"5. Choose the destination policy and parent and child allow list", "cyan"))
print(colorText(f"6. Test ------------------------------------------------------", "cyan"))
print(colorText(f" Print rather than apply selected data.", "cyan"))
print(colorText(f"7. Liftoff ------------------------------------------------------", "cyan"))
print(colorText(f" Apply path exclusions according to allowed and approved paths", "cyan"))
print(colorText(f" Apply signed or attested hashes to Parent Allow List", "cyan"))
print(colorText(f" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
print(colorText("R. Remove/Reset Generated data - will prompt to allow keeping execution history", "cyan"))
print(colorText("Q. Quit", "cyan"))
def apivalidation():
match os.getenv('APIKEY'):
case '':
print(colorText("Please add your API Key to the .env file", "red"))
def tryToReadCSV(csv):
try:
if not os.path.exists(csv):
print(colorText(f"Error: File '{csv}' does not exist.", "red"))
return pd.DataFrame() # Return empty DataFrame if file doesn't exist
df = pd.read_csv(csv)
if df.empty:
print(colorText("Error: CSV file has headers but no data rows.", "red"))
else:
print(colorText(f"Data loaded successfully from {csv}", "green"))
except pd.errors.EmptyDataError:
print(colorText("Notice: CSV file is completely empty (no headers, no data), falling back to empty frame", "white"))
df = pd.DataFrame() # Create an empty DataFrame as fallback
return df
def tryToReadParquet(parquet):
try:
df = pd.read_parquet(parquet)
if df.empty:
print(colorText("Error: Parquet file has headers but no data rows.", "red"))
else:
print(colorText(f"Data loaded successfully from {parquet}", "green"))
except pd.errors.EmptyDataError:
print(colorText("Notice : Parquet file is completely empty (no headers, no data), falling back to empty frame", "white"))
df = pd.DataFrame() # Create an empty DataFrame as fallback
return df
def deduplicate_list(lst):
seen = set()
return [x for x in lst if not (x in seen or seen.add(x))]