From 615fb778913b40bbad355e5eeb9f68ed98e9d999 Mon Sep 17 00:00:00 2001 From: Driven-Element <129630760+Driven-Element@users.noreply.github.com> Date: Sun, 24 Aug 2025 01:43:12 -0400 Subject: [PATCH] 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 --- AirlockTools.py | 209 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 139 insertions(+), 70 deletions(-) diff --git a/AirlockTools.py b/AirlockTools.py index 6721512..9e70a80 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -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")) + + 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")) - choice = input("Enter your choice: ") + 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]] + 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 == "2": - second_policy_tuple = utils.allowlist.listPolicies(url) - second_policy = second_policy_tuple[1][second_policy_tuple[0]] + 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")) + + 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")) - elif choice == "3": - 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) - - history_pol1_staged =True - - - elif choice == "4": - 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 - - 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") - df2 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{second_policy}.csv") - df_aggregated_combo= pd.concat([df1 , df2], ignore_index=True) + 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()