diff --git a/AirlockTools.py b/AirlockTools.py index b944afd..00c5d69 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -19,9 +19,11 @@ import utils.getdeviceevents import utils.allowlist import utils.hashfunctions import utils.pathfunctions +import utils.colortext as ct import urllib3 import pandas as pd + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) dotenv.load_dotenv() @@ -64,62 +66,235 @@ def menu_main(): choice = input(ct.colorText("\nEnter Menu Item: ", "white")) if choice == '1': utils.getdeviceevents.devicehistory(url,False) - if choice == '2': - utils.allowlist.allowlistexechistories(url,False) - if choice == '3': - executionhist = utils.allowlist.allowlistexechistories(url,True) - print(executionhist) - aggregated = utils.hashfunctions.aggregateHashes(executionhist) - print(aggregated) - augmented = utils.hashfunctions.augmentAggregatedHashes(url,aggregated) - print(augmented) + elif choice == "2": + menu_local_approve() + elif choice == "3": + menu_feature2() + elif choice == "4": + menu_prepare_to_enforce() + elif choice == "Q": + break + else: + print(ct.colorText("Invalid choice. Please try again.","red")) - augmented.to_html("augmentedlist.html", index=False) +def menu_local_approve(): + while True: + print("\n--- Submenu ---") + print("1. Sub-option A") + print("2. Sub-option B") + print("3. Return to Main Menu") + choice = input("Enter your choice: ") - badpublisherlist = [] - categorized = utils.hashfunctions.categorizeHashes(augmented, 5, badpublisherlist) - categorized[0].to_html("needsreview.html", index=False) - categorized[1].to_html("approved.html", index=False) - categorized[2].to_html("remaining.html", index=False) + if choice == "1": + print("You selected Sub-option A") + elif choice == "2": + print("You selected Sub-option B") + elif choice == "3": + print("Returning to Main Menu...") + break + else: + print("Invalid choice. Please try again.") - if choice == '4': - html_file = "augmentedlist.html" - augmented_df = pd.read_html(html_file) - print(augmented_df) - combined_df = pd.concat(augmented_df, ignore_index=True) +def menu_feature2(): + while True: + print("\n--- Submenu ---") + print("1. Sub-option A") + print("2. Sub-option B") + print("3. Return to Main Menu") + choice = input("Enter your choice: ") - path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(combined_df) - path_eligible.to_html("EligblePaths.html", index=False) - path_ineligible.to_html("IneligiblePaths.html",index=False) + if choice == "1": + print("You selected Sub-option A") + elif choice == "2": + print("You selected Sub-option B") + elif choice == "3": + print("Returning to Main Menu...") + break + else: + print("Invalid choice. Please try again.") + +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") + if not os.path.exists("approvals"): os.makedirs("approvals") + + df_aggregated_combo = pd.DataFrame() + while True: + print(ct.colorText("\n --------------------------------------------------------------------", "cyan")) + print(ct.colorText(" -------------------- Prepare to Enforce Policy ---------------------", "cyan")) + print(ct.colorText(" --------------------------------------------------------------------", "cyan")) + 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 choice == '5': - - executionhist = utils.allowlist.allowlistexechistories(url,True) - print(executionhist) + if first_policy == " " and second_policy == " ": + print(ct.colorText(f" [✗] No policies have been chosen","red")) + elif first_policy != " " and second_policy is first_policy: + print(ct.colorText(f" [✓] {first_policy} has been selected,", "green")) + elif first_policy != " " and second_policy != " ": + 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")) - aggregated = utils.hashfunctions.aggregateHashes(executionhist) - print(aggregated) - - augmented = utils.hashfunctions.augmentAggregatedHashes(url,aggregated) - print(augmented) - - augmented.to_html("augmentedlist.html", index=False) - html_file = "augmentedlist.html" - augmented_df = pd.read_html(html_file) - combined_df = pd.concat(augmented_df, ignore_index=True) - - path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(combined_df) - path_eligible.to_html("EligblePaths.html", index=False) - path_ineligible.to_html("IneligiblePaths.html",index=False) + print(ct.colorText("2. Pull and stage event history", "cyan")) - badpublisherlist = [] - categorized = utils.hashfunctions.categorizeHashes(augmented, 5, badpublisherlist) - categorized[0].to_html("needsreview.html", index=False) - categorized[1].to_html("approved.html", index=False) - categorized[2].to_html("remaining.html", index=False) - + if os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == True: + 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(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(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(ct.colorText(f" [✓] This has not been completed for {second_policy}","red")) + + 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(ct.colorText(" [✓] This step has been completed","green")) + else: + print(ct.colorText(" [✗] This step has not been completed","red")) + + 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(ct.colorText(" [✓] This step has been completed","green")) + else: + print(ct.colorText(" [✗] This step has not been completed","red")) + + 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(ct.colorText(" [✓] This step has been completed","green")) + else: + print(ct.colorText(" [✗] This step has not been completed", "red")) + + print(ct.colorText("6. Categorize your hashes ", "cyan")) + + if os.path.isfile(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv") and os.path.isfile(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv") and os.path.isfile(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv"): + print(ct.colorText(" [✓] This step has been completed","green")) + else: + print(ct.colorText(" [✗] This step has not been completed","red")) + 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(ct.colorText(" [✓] This step has been completed","green")) + else: + print(ct.colorText(" [✗] This step has not been completed","red")) + + + print(ct.colorText("Q. Quit", "cyan")) + + choice = input(ct.colorText("\nEnter your choice: ", "white")) + if choice == "1": + first_policy_tuple = utils.allowlist.listPolicies(url) + first_policy = first_policy_tuple[1][first_policy_tuple[0]] + while True: + answer = input(ct.colorText(f"{"Do you want to load a second policy?"} (yes/no): ", "white").strip().lower()) + if answer in ("yes", "y"): + second_policy_tuple = utils.allowlist.listPolicies(url) + second_policy = second_policy_tuple[1][second_policy_tuple[0]] + break + elif answer in ("no", "n"): + second_policy_tuple = first_policy_tuple + second_policy = first_policy + break + else: + 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"): + 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(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(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(f"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(f"dataframe_csv\\df_aggregated_{first_policy}.csv") and os.path.exists(f"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) + else: + print(ct.colorText(f"Please stage your data before attempting this step","red")) + + elif choice == "4": + if os.path.exists(f"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) + print(ct.colorText(f"Hash reputation info added to dataframe","green")) + else: + 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"): + 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) + print(ct.colorText(f"Eligible paths determined","green")) + else: + 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_augmented_combo_{first_policy}_{second_policy}.csv"): + categorized = utils.hashfunctions.categorizeHashes(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"), threat_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) + categorized[1].to_html(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html", index=False) + 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_unapproved_hashes__{first_policy}_{second_policy}.html", index=False) + categorized[2].to_csv(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv", index=False) + print(ct.colorText(f"Hashes have been categorized","green")) + else: + print(ct.colorText(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red")) + + elif choice == "7": + if os.path.exists(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv") and os.path.exists(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv") and os.path.exists(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv"): + 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) + print(ct.colorText(f"Allowable paths determined","green")) + else: + print(ct.colorText(f"Please complete step 6 prior to attempting this step","red")) + + elif choice == "Q": + break + + else: + print(ct.colorText("Invalid choice. Please try again.", "red")) + +def tryToReadCSV(csv): + try: + df =pd.read_csv(csv) + if df.empty: + print(ct.colorText("Error: CSV file has headers but no data rows.", "red")) + else: + print(ct.colorText("Data loaded successfully.", "green")) + except pd.errors.EmptyDataError: + 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 + if __name__ == "__main__":