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

Refactored and Colorized Menu. Do not have a chance to verify if api calls are still working, will do some more looking at that tomorrow

Browse Source
This commit is contained in:
Driven-Element
2025-08-24 01:43:12 -04:00
parent 6a6057e612
commit 615fb77891
1 changed files with 139 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AirlockTools.py
+127 -58
View File
@@ -33,16 +33,32 @@ 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(r"""
print(color_text(r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ 4
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""")
print("=== Welcome to the Airlock API Tool ===")
""", "green"))
print(color_text("============ Welcome to the Airlock API Tool ============", "magenta"))
match os.getenv('APIKEY'):
case '':
print("Please add your API Key to the .env file")
@@ -51,12 +67,12 @@ def apivalidation():
def menu_main():
while True:
print("\n--- Main Menu ---")
print("1. Get All Events for Single Device")
print("2. Placeholder for Local Approval")
print("3. Placeholder for Another Tool" )
print("4. Prepare Policy For Enforcement")
print("11. Exit")
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"))
choice = input("Enter Menu Item: ")
if choice == '1':
@@ -70,7 +86,7 @@ def menu_main():
elif choice == "11":
break
else:
print("Invalid choice. Please try again.")
print(color_text("Invalid choice. Please try again.","red"))
def menu_local_approve():
while True:
@@ -111,88 +127,138 @@ def menu_feature2():
def menu_prepare_to_enforce():
first_policy = " "
second_policy = " 4"
history_pol1_staged = True
history_pol2_staged = True
history_staged = True
hashes_threat_pulled = True
hashes_categorized = True
path_exclusions_calculated = True
allowed_paths_determined = True
second_policy = " "
df_aggregated_combo = pd.DataFrame()
while True:
print("\n--- Prepare to Enforce Policy ---")
print("Sequentually follow steps to prepare for policy enforcement")
print("1. Choose first policy, Typically the audit version of the policy - Currently selected first policy is: " + first_policy)
print("2. Choose second policy, If an enforcement policy of that type exists, include it here - Currently selected second policy is: " + second_policy)
print("3. Pull and stage event history for the first policy - This step has been done - " + str(history_pol1_staged))
print("4. Pull and stage event history for the second policy - This step has been done - " + str(history_pol2_staged))
print("5. Combine aggregated policies -")
print("6. Pull Hash info - This step has been done - " + str(hashes_threat_pulled))
print("7. Determine if path exclusions are possible: - This step has been done - " + str(path_exclusions_calculated))
print("8. Categorize your hashes - This step has been done - " + str(hashes_categorized))
print("9. Combine Potential path exclusions and allowed hashes - This step has been done" + str(allowed_paths_determined))
print("11. Exit")
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"))
choice = input("Enter your choice: ")
if first_policy == " " and second_policy == " ":
print(color_text(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"))
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(color_text("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"))
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"))
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"))
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(color_text("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"))
else:
print(color_text(" This step has not been completed","red"))
print(color_text("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"))
else:
print(color_text(" This step has not been completed","red"))
print(color_text("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"))
else:
print(color_text(" This step has not been completed", "red"))
print(color_text("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"))
else:
print(color_text(" This step has not been completed","red"))
print(color_text("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"))
else:
print(color_text(" This step has not been completed","red"))
print(color_text("11. Exit", "cyan"))
choice = input(color_text("Enter your choice: ", "yellow"))
if choice == "1":
first_policy_tuple = utils.allowlist.listPolicies(url)
first_policy = first_policy_tuple[1][first_policy_tuple[0]]
elif choice == "2":
while True:
answer = input(color_text(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]]
return False
elif answer in ("no", "n"):
second_policy_tuple = first_policy_tuple
second_policy = first_policy
return False
else:
print(color_text("Please answer with 'yes' or 'no'.", "red"))
elif choice == "3":
elif choice == "2":
if not os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
executionhist_policy1 = utils.allowlist.pullPolicyExechistories(url,first_policy_tuple[0], first_policy_tuple[1],True)
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"))
history_pol1_staged =True
elif choice == "4":
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)
history_pol2_staged =True
print(color_text(f"Staging of Exection history for policy: {second_policy} is complete","green"))
elif choice == "5":
if history_pol1_staged == True & history_pol2_staged == True:
elif choice == "3":
if second_policy is first_policy and os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
df1 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{first_policy}.csv")
df_aggregated_combo = df1
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)
elif os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv") and os.path.exists("dataframe_csv\\df_aggregated_{second_policy}.csv"):
df1 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{first_policy}.csv")
df2 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{second_policy}.csv")
df_aggregated_combo = pd.concat([df1 , df2], ignore_index=True)
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)
history_staged = True
else:
print("Both Policies have to be staged to combine them")
print(color_text(f"Please stage your data before attempting this step","red"))
elif choice == "6":
if history_staged == True:
elif choice == "4":
if os.path.exists("dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv"):
df_augmented = utils.hashfunctions.augmentAggregatedHashes(url,tryToReadCSV(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv"))
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)
hashes_threat_pulled = True
else:
print("Data not yet staged, please complete earlier steps")
print(color_text(f"Please combine your data with step 3 prior to attempting this step","red"))
elif choice == "7":
if hashes_threat_pulled == True:
elif choice == "5":
if os.path.exists(f"dataframe_html\\df_augmented_combo_{first_policy}_{second_policy}.html"):
path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"))
path_eligible.to_html(f"dataframe_html\\df_path_eligible_{first_policy}_{second_policy}.html", index=False)
path_eligible.to_csv(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv", index=False)
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)
path_exclusions_calculated = True
else:
print("Step 6 not complete, Please complete step 6")
print(color_text(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red"))
elif choice == "8":
if path_exclusions_calculated == True:
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:
categorized = utils.hashfunctions.categorizeHashes(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"), treat_tolerance_constant, badpublisherlist)
categorized[0].to_html(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html", index=False)
categorized[0].to_csv(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv", index=False)
@@ -200,20 +266,21 @@ def menu_prepare_to_enforce():
categorized[1].to_csv(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv", index=False)
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)
hashes_categorized = True
else:
print("Step 7 not complete, Please complete step 7")
print(color_text(f"Please complete step 5 prior to attempting this step","red"))
elif choice == "9":
if hashes_categorized == True:
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:
allowpaths = utils.allowfunctions.filter_and_drop(pd.read_csv(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv"),tryToReadCSV(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv"), path_exclusion_constant)
allowpaths.to_html(f"dataframe_html\\df_allowed_paths_{first_policy}_{second_policy}.html", index=False)
allowpaths.to_csv(f"dataframe_csv\\df_allowed_paths_{first_policy}_{second_policy}.csv", index=False)
allowed_paths_determined = True
else:
print("Step 8 not complete, Please complete step 8")
print(color_text(f"Please complete step 6 prior to attempting this step","red"))
elif choice == "11":
break
else:
print("Invalid choice. Please try again.")
@@ -229,6 +296,8 @@ def tryToReadCSV(csv):
df = pd.DataFrame() # Create an empty DataFrame as fallback
return df
if __name__ == "__main__":
apivalidation()