diff --git a/AirlockTools.py b/AirlockTools.py index 3478647..d2aa6ea 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -92,8 +92,11 @@ def menu_feature2(): print("Invalid choice. Please try again.") def menu_prepare_to_enforce(): - first_policy = [""] - second_policy = [""] + first_policy = ["",""] + second_policy = ["",""] + policylist = first_policy[1] + history_pol1_staged = False + history_pol2_staged = False history_staged = False hashes_threat_pulled = False hashes_categorized = False @@ -104,9 +107,12 @@ def menu_prepare_to_enforce(): 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: " + str([first_policy[0]])) - print("2. Choose second policy, If an enforcement policy of that type exists, include it here - Currently selected second policy is: " + str(second_policy[0])) - print("3. Pull and stage event history for the selected policies - This step has been done - " + str(history_staged)) + #print("1. Choose first policy, Typically the audit version of the policy - Currently selected first policy is: " + str([[first_policy[0]])) + #print("2. Choose second policy, If an enforcement policy of that type exists, include it here - Currently selected second policy is: " + str(second_policy[0])) + 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("4. Pull Hash info - This step has been done - " + str(hashes_threat_pulled)) print("5. Determine if path exclusions are possible: - This step has been done - " + str(path_exclusions_calculated)) print("6. Categorize your hashes - This step has been done - " + str(hashes_categorized)) @@ -121,32 +127,45 @@ def menu_prepare_to_enforce(): second_policy = utils.allowlist.listPolicies(url) elif choice == "3": - executionhist_policy1 = utils.allowlist.pullPolicyExechistories(url,first_policy[0], first_policy[1],True) - executionhist_policy2 = utils.allowlist.pullPolicyExechistories(url,second_policy[0], second_policy[1],True) + executionhist_policy1 = utils.allowlist.pullPolicyExechistories(url,first_policy, first_policy,True) df_aggregated_policy1 = utils.hashfunctions.aggregateHashes(executionhist_policy1) - df_aggregated_policy2 = utils.hashfunctions.aggregateHashes(executionhist_policy2) - df_aggregated_combo= pd.concat([df_aggregated_policy1, df_aggregated_policy2], ignore_index=True) - df_aggregated_combo.to_html(f"df_aggregated_combo_{first_policy}_{second_policy}.html", index=False) - history_staged = True + df_aggregated_policy1.to_html(f"df_aggregated_{first_policy}.html") + history_pol1_staged =True elif choice == "4": + executionhist_policy2 = utils.allowlist.pullPolicyExechistories(url,second_policy, second_policy,True) + df_aggregated_policy2 = utils.hashfunctions.aggregateHashes(executionhist_policy2) + df_aggregated_policy2.to_html(f"df_aggregated_{second_policy}.html") + history_pol2_staged =True + + elif choice == "5": + if history_pol1_staged == True & history_pol2_staged == True: + df1 = pd.read_html(f"df_aggregated_{first_policy}.html")[0] + df2 = pd.read_html(f"df_aggregated_{second_policy}.html")[0] + df_aggregated_combo= pd.concat([df1 , df2], ignore_index=True) + df_aggregated_combo.to_html(f"df_aggregated_combo_{first_policy}_{second_policy}.html", index=False) + history_staged = True + else: + print("Both Policies have to be staged to combine them") + + elif choice == "6": if history_staged == True: df_augmented = utils.hashfunctions.augmentAggregatedHashes(url, pd.read_html(f"df_aggregated_combo_{first_policy}_{second_policy}.html")[0]) df_augmented.to_html(f"df_augmented_combo_{first_policy}_{second_policy}.html", index=False) hashes_threat_pulled = True else: - print("Data not yet staged, please complete steps 1-3") + print("Data not yet staged, please complete earlier steps") - elif choice == "5": + elif choice == "7": if hashes_threat_pulled == True: path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(pd.read_html(f"df_augmented_combo_{first_policy}_{second_policy}.html")[0]) path_eligible.to_html(f"df_path_eligible_{first_policy}_{second_policy}.html", index=False) path_ineligible.to_html(f"df_path_ineligible_{first_policy}_{second_policy}.html", index=False) path_exclusions_calculated = True else: - print("Step 4 not complete, Please complete step 4") + print("Step 6 not complete, Please complete step 6") - elif choice == "6": + elif choice == "8": if path_exclusions_calculated == True: categorized = utils.hashfunctions.categorizeHashes(pd.read_html(f"df_augmented_combo_{first_policy}_{second_policy}.html")[0], treat_tolerance_constant, badpublisherlist) categorized[0].to_html(f"df_hashes_needing_approval_{first_policy}_{second_policy}.html", index=False) @@ -154,15 +173,15 @@ def menu_prepare_to_enforce(): categorized[2].to_html(f"df_remaining_hashes__{first_policy}_{second_policy}.html", index=False) hashes_categorized = True else: - print("Step 5 not complete, Please complete step 5") + print("Step 7 not complete, Please complete step 7") - elif choice == "7": + elif choice == "9": if hashes_categorized == True: allowpaths = utils.allowfunctions.filter_and_drop(pd.read_html(f"df_automatically_approved_hashes_{first_policy}_{second_policy}.html")[0], pd.read_html(f"df_path_eligible_{first_policy}_{second_policy}.html")[0], path_exclusion_constant) allowpaths.to_html(f"df_allowed_paths_{first_policy}_{second_policy}.html", index=False) allowed_paths_determined = True else: - print("Step 6 not complete, Please complete step 6") + print("Step 8 not complete, Please complete step 8") elif choice == "11": break else: