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

@@ -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
 \____|__  /__||__|  |____/\____/ \___  >__|_ \   |____| \____/ \____/|____/____  >
         \/                           \/     \/                                 \/ 
-""")
+""", "green"))
-	print("=== Welcome to the Airlock API Tool ===")
+	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(color_text("\n--- Main Menu ---", "yellow") + color_text("", "bold"))
-		print("1. Get All Events for Single Device")
+		print(color_text("1. Get All Events for Single Device", "yellow"))
-		print("2. Placeholder for Local Approval")
+		print(color_text("2. Placeholder for Local Approval", "yellow"))
-		print("3. Placeholder for Another Tool" )
+		print(color_text("3. Placeholder for Another Tool", "yellow"))
-		print("4. Prepare Policy For Enforcement")
+		print(color_text("4. Prepare Policy For Enforcement", "yellow"))
-		print("11. Exit")
+		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"
+	second_policy = " "
-	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
 
 	df_aggregated_combo = pd.DataFrame()
 	while True:
-		print("\n--- Prepare to Enforce Policy ---")
+		print(color_text("\n---------Prepare to Enforce Policy ---------------------", "white") + color_text("", "bold"))
-		print("Sequentually follow steps to prepare for policy enforcement")
+		print(color_text("Sequentually follow steps to prepare for policy enforcement", "white") + color_text("", "underline"))
-		print("1. Choose first policy, Typically the audit version of the policy - Currently selected first policy is: " + first_policy)
+		print(color_text("1. Choose which policy or policies to work with - : ", "cyan"))
-		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")
 		
-		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]]
+			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)
+			if  not os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
-			second_policy = second_policy_tuple[1][second_policy_tuple[0]]
+				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)
+			if  second_policy is first_policy and os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
-			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:
 				df1 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{first_policy}.csv")
-				df2 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{second_policy}.csv")
+				df_aggregated_combo = df1
-				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)
+			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":
+		elif choice == "4":
-			if history_staged == True:
+			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":
+		elif choice == "5":
-			if hashes_threat_pulled == True:
+			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":
+		elif choice == "6":
-			if path_exclusions_calculated == True:
+			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":
+		elif choice == "7":
-			if hashes_categorized == True:
+			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()