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

Continued Colorization, pushed it to utils and colorized other helper scripts

Browse Source
This commit is contained in:
Driven-Element
2025-08-24 13:46:20 -04:00
parent 615fb77891
commit bdf3d4b9ec
5 changed files with 105 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AirlockTools.py
+65 -68
View File
@@ -24,6 +24,7 @@ import urllib3
import pandas as pd
import sqlite3
import pathlib
import utils.colortext as ct
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
dotenv.load_dotenv()
@@ -33,24 +34,10 @@ badpublisherlist = ["Brave Software, Inc.", "Zoom Video Communications, Inc."]
path_exclusion_constant = 3
treat_tolerance_constant = 4
def color_text(text: str, color: str) -> str:
colors = {
"red": "\033[91m",
"green": "\033[92m",
"yellow": "\033[93m",
"blue": "\033[94m",
"magenta": "\033[95m",
"cyan": "\033[96m",
"white": "\033[97m",
"reset": "\033[0m",
"bold": "\033[1m",
"underline": "\033[4m",
}
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
def apivalidation():
print(color_text(r"""
print(ct.colorText(r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
@@ -58,23 +45,27 @@ def apivalidation():
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""", "green"))
print(color_text("============ Welcome to the Airlock API Tool ============", "magenta"))
print(ct.colorText("=================================================================================", "magenta"))
print(ct.colorText("======================== Welcome to the Airlock API Tool ========================", "magenta"))
print(ct.colorText("=================================================================================", "magenta"))
match os.getenv('APIKEY'):
case '':
print("Please add your API Key to the .env file")
print(ct.colorText("Please add your API Key to the .env file", "red"))
case _:
menu_main()
def menu_main():
while True:
print(color_text("\n--- Main Menu ---", "yellow") + color_text("", "bold"))
print(color_text("1. Get All Events for Single Device", "yellow"))
print(color_text("2. Placeholder for Local Approval", "yellow"))
print(color_text("3. Placeholder for Another Tool", "yellow"))
print(color_text("4. Prepare Policy For Enforcement", "yellow"))
print(color_text("11. Exit", "yellow"))
print(ct.colorText("\n-----------------------------------", "yellow"))
print(ct.colorText("------------ Main Menu ------------", "yellow"))
print(ct.colorText("-----------------------------------", "yellow"))
print(ct.colorText("1. Get All Events for Single Device", "yellow"))
print(ct.colorText("2. Placeholder for Local Approval", "yellow"))
print(ct.colorText("3. Placeholder for Another Tool", "yellow"))
print(ct.colorText("4. Prepare Policy For Enforcement", "yellow"))
print(ct.colorText("11. Exit", "yellow"))
choice = input("Enter Menu Item: ")
choice = input(ct.colorText("Enter Menu Item: ", "white"))
if choice == '1':
utils.getdeviceevents.devicehistory(url,False)
elif choice == "2":
@@ -86,7 +77,7 @@ def menu_main():
elif choice == "11":
break
else:
print(color_text("Invalid choice. Please try again.","red"))
print(ct.colorText("Invalid choice. Please try again.","red"))
def menu_local_approve():
while True:
@@ -129,75 +120,81 @@ def menu_prepare_to_enforce():
first_policy = " "
second_policy = " "
#If the directorys where we're going to store our output dont exist, make them.
if not os.path.exists("dataframe_html"): os.makedirs("dataframe_html")
if not os.path.exists("dataframe_csv"): os.makedirs("dataframe_csv")
df_aggregated_combo = pd.DataFrame()
while True:
print(color_text("\n---------Prepare to Enforce Policy ---------------------", "white") + color_text("", "bold"))
print(color_text("Sequentually follow steps to prepare for policy enforcement", "white") + color_text("", "underline"))
print(color_text("1. Choose which policy or policies to work with - : ", "cyan"))
print(ct.colorText("\n --------------------------------------------------------------------", "magenta"))
print(ct.colorText(" -------------------- Prepare to Enforce Policy ---------------------", "magenta"))
print(ct.colorText(" --------------------------------------------------------------------", "magenta"))
print(ct.colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
print(ct.colorText("\n1. Choose which policy or policies to work with - : ", "cyan"))
if first_policy == " " and second_policy == " ":
print(color_text(f" No policies have been chosen","red"))
print(ct.colorText(f" No policies have been chosen","red"))
elif first_policy != " " and second_policy is first_policy:
print(color_text(f" {first_policy} has been selected,", "green"))
print(ct.colorText(f" {first_policy} has been selected,", "green"))
elif first_policy != " " and second_policy != " ":
print(color_text(f" {first_policy} has been selected as Policy 1","green"))
print(color_text(f" {second_policy} has been selected as Policy 2","green"))
print(ct.colorText(f" {first_policy} has been selected as Policy 1","green"))
print(ct.colorText(f" {second_policy} has been selected as Policy 2","green"))
print(color_text("2. Pull and stage event history.", "cyan"))
print(ct.colorText("2. Pull and stage event history", "cyan"))
if os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == True:
print(color_text(f" This has been completed for {first_policy}","green"))
print(ct.colorText(f" This has been completed for {first_policy}","green"))
elif os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == False:
print(color_text(f" This has not been completed for {first_policy}","red"))
print(ct.colorText(f" This step has not been completed","red"))
elif second_policy is not first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv") == True:
print(color_text(f" This has been completed for {second_policy}","green"))
print(ct.colorText(f" This has been completed for {second_policy}","green"))
elif second_policy is not first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv") == False:
print(color_text(f" This has not been completed for {second_policy}","red"))
print(ct.colorText(f" This has not been completed for {second_policy}","red"))
print(color_text("3. Combine Staged policies", "cyan"))
print(ct.colorText("3. Combine Staged policies", "cyan"))
if os.path.exists(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv") == True:
print(color_text(" This step has been completed","green"))
print(ct.colorText(" This step has been completed","green"))
else:
print(color_text(" This step has not been completed","red"))
print(ct.colorText(" This step has not been completed","red"))
print(color_text("4. Add hash threat information to list of executions", "cyan"))
print(ct.colorText("4. Add hash threat information to list of executions", "cyan"))
if os.path.exists(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv") == True:
print(color_text(" This step has been completed","green"))
print(ct.colorText(" This step has been completed","green"))
else:
print(color_text(" This step has not been completed","red"))
print(ct.colorText(" This step has not been completed","red"))
print(color_text("5. Determine if path exclusions are possible", "cyan"))
print(ct.colorText("5. Determine if path exclusions are possible", "cyan"))
if os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True:
print(color_text(" This step has been completed","green"))
print(ct.colorText(" This step has been completed","green"))
else:
print(color_text(" This step has not been completed", "red"))
print(ct.colorText(" This step has not been completed", "red"))
print(color_text("6. Categorize your hashes ", "cyan"))
print(ct.colorText("6. Categorize your hashes ", "cyan"))
if os.path.exists(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_remaining_hashes_{first_policy}_{second_policy}.html") == True:
print(color_text(" This step has been completed","green"))
print(ct.colorText(" This step has been completed","green"))
else:
print(color_text(" This step has not been completed","red"))
print(ct.colorText(" This step has not been completed","red"))
print(color_text("7. Compare potential path exclusions with allowed hashes", "cyan"))
print(ct.colorText("7. Compare potential path exclusions with allowed hashes", "cyan"))
if os.path.exists(f"dataframe_csv\\df_allowed_paths_{first_policy}_{second_policy}.csv") == True:
print(color_text(" This step has been completed","green"))
print(ct.colorText(" This step has been completed","green"))
else:
print(color_text(" This step has not been completed","red"))
print(ct.colorText(" This step has not been completed","red"))
print(color_text("11. Exit", "cyan"))
print(ct.colorText("11. Exit", "cyan"))
choice = input(color_text("Enter your choice: ", "yellow"))
choice = input(ct.colorText("Enter your choice: ", "yellow"))
if choice == "1":
first_policy_tuple = utils.allowlist.listPolicies(url)
first_policy = first_policy_tuple[1][first_policy_tuple[0]]
while True:
answer = input(color_text(f"{"Do you want to load a second policy?"} (yes/no): ", "magenta").strip().lower())
answer = input(ct.colorText(f"{"Do you want to load a second policy?"} (yes/no): ", "magenta").strip().lower())
if answer in ("yes", "y"):
second_policy_tuple = utils.allowlist.listPolicies(url)
second_policy = second_policy_tuple[1][second_policy_tuple[0]]
@@ -207,7 +204,7 @@ def menu_prepare_to_enforce():
second_policy = first_policy
return False
else:
print(color_text("Please answer with 'yes' or 'no'.", "red"))
print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))
elif choice == "2":
if not os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
@@ -215,14 +212,14 @@ def menu_prepare_to_enforce():
df_aggregated_policy1 = utils.hashfunctions.aggregateHashes(executionhist_policy1)
df_aggregated_policy1.to_html(f"dataframe_html\\df_aggregated_{first_policy}.html", index=False)
df_aggregated_policy1.to_csv(f"dataframe_csv\\df_aggregated_{first_policy}.csv", index=False)
print(color_text(f"Staging of Exection history for policy: {first_policy} is complete","green"))
print(ct.colorText(f"Staging of Exection history for policy: {first_policy} is complete","green"))
if not os.path.exists("dataframe_csv\\df_aggregated_{second_policy}.csv"):
executionhist_policy2 = utils.allowlist.pullPolicyExechistories(url,second_policy_tuple[0], second_policy_tuple[1],True)
df_aggregated_policy2 = utils.hashfunctions.aggregateHashes(executionhist_policy2)
df_aggregated_policy2.to_html(f"dataframe_html\\df_aggregated_{second_policy}.html", index=False)
df_aggregated_policy2.to_csv(f"dataframe_csv\\df_aggregated_{second_policy}.csv", index=False)
print(color_text(f"Staging of Exection history for policy: {second_policy} is complete","green"))
print(ct.colorText(f"Staging of Exection history for policy: {second_policy} is complete","green"))
elif choice == "3":
if second_policy is first_policy and os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
@@ -237,7 +234,7 @@ def menu_prepare_to_enforce():
df_aggregated_combo.to_html(f"dataframe_html\\df_aggregated_combo_{first_policy}_{second_policy}.html", index=False)
df_aggregated_combo.to_csv(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv", index=False)
else:
print(color_text(f"Please stage your data before attempting this step","red"))
print(ct.colorText(f"Please stage your data before attempting this step","red"))
elif choice == "4":
if os.path.exists("dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv"):
@@ -245,7 +242,7 @@ def menu_prepare_to_enforce():
df_augmented.to_html(f"dataframe_html\\df_augmented_combo_{first_policy}_{second_policy}.html", index=False)
df_augmented.to_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv", index=False)
else:
print(color_text(f"Please combine your data with step 3 prior to attempting this step","red"))
print(ct.colorText(f"Please combine your data with step 3 prior to attempting this step","red"))
elif choice == "5":
if os.path.exists(f"dataframe_html\\df_augmented_combo_{first_policy}_{second_policy}.html"):
@@ -255,7 +252,7 @@ def menu_prepare_to_enforce():
path_ineligible.to_html(f"dataframe_html\\df_path_ineligible_{first_policy}_{second_policy}.html", index=False)
path_ineligible.to_csv(f"dataframe_csv\\df_path_ineligible_{first_policy}_{second_policy}.csv", index=False)
else:
print(color_text(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red"))
print(ct.colorText(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red"))
elif choice == "6":
if os.path.exists(f"dataframe_csv\\df_path_ineligible_{first_policy}_{second_policy}.csv") == True and os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True:
@@ -267,7 +264,7 @@ def menu_prepare_to_enforce():
categorized[2].to_html(f"dataframe_html\\df_remaining_hashes__{first_policy}_{second_policy}.html", index=False)
categorized[2].to_csv(f"dataframe_csv\\df_remaining_hashes__{first_policy}_{second_policy}.csv", index=False)
else:
print(color_text(f"Please complete step 5 prior to attempting this step","red"))
print(ct.colorText(f"Please complete step 5 prior to attempting this step","red"))
elif choice == "7":
if os.path.exists(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_remaining_hashes_{first_policy}_{second_policy}.html") == True:
@@ -276,23 +273,23 @@ def menu_prepare_to_enforce():
allowpaths.to_csv(f"dataframe_csv\\df_allowed_paths_{first_policy}_{second_policy}.csv", index=False)
allowed_paths_determined = True
else:
print(color_text(f"Please complete step 6 prior to attempting this step","red"))
print(ct.colorText(f"Please complete step 6 prior to attempting this step","red"))
elif choice == "11":
break
else:
print("Invalid choice. Please try again.")
print(ct.colorText("Invalid choice. Please try again.", "red"))
def tryToReadCSV(csv):
try:
df =pd.read_csv(csv)
if df.empty:
print("CSV file has headers but no data rows.")
print(ct.colorText("Error: CSV file has headers but no data rows.", "red"))
else:
print("Data loaded successfully.")
print(ct.colorText("Data loaded successfully.", "green"))
except pd.errors.EmptyDataError:
print("CSV file is completely empty (no headers, no data).")
print(ct.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
utils/allowlist.py
+5 -4
View File
@@ -17,6 +17,7 @@ import requests
import json
import os
import time
import utils.colortext as ct
def pullPolicyExechistories(url, choice, policiesnames, outputjson: bool):
@@ -32,7 +33,7 @@ def pullPolicyExechistories(url, choice, policiesnames, outputjson: bool):
for index, item in enumerate(json_response_data['response']['exechistories']):
if index == len(json_response_data['response']['exechistories']) -1:
checkpoint = item['checkpoint']
print(f"Date Greater than 30 Days, Stepping to new Checkpoint. {item['checkpoint']}")
print(ct.colorText(f"Date Greater than 30 Days, Stepping to new Checkpoint. {item['checkpoint']}", "blue"))
else:
if (datetime.date.today() - datetime.timedelta(days=10) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()):
pass
@@ -57,7 +58,7 @@ def checkpoint_stomper(checkpoint, url, policy, headers):
def listPolicies(url):
endpoint = url + '/v1/group'
print("[+] Grabbing All Policies")
print(ct.colorText("[+] Grabbing All Policies", "magenta"))
payload = {}
headers = {
"X-APIKey": os.getenv('APIKEY')
@@ -67,9 +68,9 @@ def listPolicies(url):
policiesnames = []
policyids = []
for index, list in enumerate(parse_text['response']['groups'], start=1):
print(f"{index}. {list['name']}")
print(ct.colorText(f"{index}. {list['name']}", "yellow"))
policiesnames.append(list['name'])
policyids.append(list['groupid'])
choice = input("Select Policy Group: ")
choice = input(ct.colorText("Select Policy Group: ", "yellow"))
choice = int(choice) - 1
return choice, policiesnames
utils/colortext.py
+14
View File
@@ -0,0 +1,14 @@
def colorText(text: str, color: str) -> str:
colors = {
"red": "\033[91m",
"green": "\033[92m",
"yellow": "\033[93m",
"blue": "\033[94m",
"magenta": "\033[95m",
"cyan": "\033[96m",
"white": "\033[97m",
"reset": "\033[0m"
}
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
utils/getdeviceevents.py
+19 -18
View File
@@ -16,16 +16,17 @@ import datetime
import requests
import json
import os
import utils.colortext as ct
def devicehistory(url, outputjson: bool):
endpoint = url + '/v1/getexechistory'
print("\n")
print("1. Today")
print("2. Last 24 Hours")
print("3. Past 7 Days")
print("4. Past 30 Days")
print("5. Custom Date Range")
choice = input("\nSelect Date Range: ")
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")
if choice == '1':
@@ -40,18 +41,18 @@ def devicehistory(url, outputjson: bool):
date_selected = datetime.date.today() - datetime.timedelta(days=30)
date_selected = date_selected.strftime('%Y-%m-%d')
elif choice == "5":
print("Please Input Dates as YYYY-MM-DD")
date_selected = input("From: ")
today = input("Date To: ")
print("WARNING: Device Name is Case Sensitive")
device = input("Enter Device Name: ")
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(payload)
print(ct.colorText(payload, "green"))
headers = {
"X-APIKey": os.getenv('APIKEY')
}
@@ -64,10 +65,10 @@ def devicehistory(url, outputjson: bool):
parse_text = json.loads(response.text)
for block in parse_text['response']['exechistory']:
print(f"Command: {block['commandline']}")
print(f"Date: {block['datetime']}")
print(f"Filename: {block['filename']}")
print(f"Policy Name: {block['policyname']}")
print(f"Hostname: {block['hostname']}")
print(f"Hash: {block['sha256']}")
print(ct.colorText(f"Command: {block['commandline']}","green"))
print(ct.colorText(f"Date: {block['datetime']}","green"))
print(ct.colorText(f"Filename: {block['filename']}","green"))
print(ct.colorText(f"Policy Name: {block['policyname']}","green"))
print(ct.colorText(f"Hostname: {block['hostname']}","green"))
print(ct.colorText(f"Hash: {block['sha256']}","green"))
print("\n")
utils/hashfunctions.py
+2 -1
View File
@@ -16,6 +16,7 @@ import pandas as pd
import requests
import os
import json
import utils.colortext as ct
def aggregateHashes(executions_json) -> pd.DataFrame:
@@ -27,7 +28,7 @@ def aggregateHashes(executions_json) -> pd.DataFrame:
if df.empty:
return df
print(df)
print(color_text(df,"green"))
# Aggregate by sha256, deduplicate lists, and preserve order
agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()