diff --git a/AirlockTools.py b/AirlockTools.py index 26de05c..9a10dc9 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -90,16 +90,21 @@ def apivalidation(): 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(f"Data loaded successfully from {csv}", "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 + try: + if not os.path.exists(csv): + print(ct.colorText(f"Error: File '{csv}' does not exist.", "red")) + return pd.DataFrame() # Return empty DataFrame if file doesn't exist + + 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(f"Data loaded successfully from {csv}", "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 + def tryToReadParquet(parquet): try: @@ -479,6 +484,9 @@ def menu_prepare_to_enforce(): allowlist_child_name, allowlist_child_id ) + elif choice == "R": + + utils.pathfunctions.clean_folders_enforcement_prep() elif choice == "Q": break diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py index 2c45e32..0a69c21 100644 --- a/utils/hashfunctions.py +++ b/utils/hashfunctions.py @@ -364,8 +364,10 @@ def divideSortedHashExecutions(first_policy,second_policy, pups): def generatePreflights(first_policy, second_policy): allhashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet") - pathexclusions = tryToReadCSV(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv") - pathexclusions.to_parquet(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet", index=False) + primarypathexclusions = tryToReadCSV(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv") + secondarypathexclusions = tryToReadCSV(f"approved\\secondary_paths_{first_policy}_{second_policy}.csv") + + pathexclusions = pd.concat([primarypathexclusions, secondarypathexclusions], ignore_index=True) publishers = tryToReadCSV(f"approved\\publishers_{first_policy}_{second_policy}.csv") publishers.to_parquet(f"parquet\\publishers_{first_policy}_{second_policy}.parquet", index=False) diff --git a/utils/pathfunctions.py b/utils/pathfunctions.py index 658b870..9c5e4cd 100644 --- a/utils/pathfunctions.py +++ b/utils/pathfunctions.py @@ -196,7 +196,48 @@ def generatePathReview(first_policy, second_policy, badpathparts, path_exclusion del unique_sha_counts del lcp_not_forbidden_review + if not os.path.exists(f"parquet\\secondary_paths_{first_policy}_{second_policy}.parquet"): + all_approved_hashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet") + recommended_paths = pd.read_parquet(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet") + remaining_hashes = all_approved_hashes[~all_approved_hashes['sha256'].isin(recommended_paths['sha256'])] + + print(ct.colorText(f"Beginning calculating longest common filepaths for path exceptions","green")) + + haslcp = pathf.split_filepaths_grouped(remaining_hashes,"filename", 3, min_files_for_path) + haslcp.drop_duplicates() + + forbidden = pathf.regulator(badpathparts, True) + forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False) + + + print(ct.colorText("Removing forbidden filepaths for path exceptions", "green")) + + # Make a real DataFrame copy before modifying + lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() + + #For the review, drop down to only the columns we care, and then group by the commmon file path, consolidating and dropping dupes + lcp_not_forbidden_review = lcp_not_forbidden[['longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']] + + # Count unique sha256 per longestcfp + unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index() + unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count'] + + # Merge the count back into the original DataFrame + lcp_not_forbidden_review = lcp_not_forbidden_review.merge(unique_sha_counts, on='longestcfp', how='left') + lcp_not_forbidden_review = lcp_not_forbidden_review[lcp_not_forbidden_review['unique_sha256_count'] >= min_files_for_path] + + lcp_not_forbidden_review.to_parquet(f"parquet\\secondary_paths_{first_policy}_{second_policy}.parquet",index=False) + lcp_not_forbidden_review.to_csv(f"needs_approved\\secondary_paths_{first_policy}_{second_policy}.csv",index=False) + ct.style_dataframe_dark(lcp_not_forbidden_review,f"needs_approved\\secondary_paths_{first_policy}_{second_policy}.html", True) + + remaining_hashes = remaining_hashes[~remaining_hashes['sha256'].isin(lcp_not_forbidden_review['sha256'])] + + remaining_hashes.to_csv(f"needs_approved\\not_covered_by_path_exclusion_{first_policy}_{second_policy}.csv",index=False) + + del lcp_not_forbidden + del unique_sha_counts + del lcp_not_forbidden_review def allATpaths(url, pups, untrusted_publishers, badpathparts, threat_tolerance, path_exclusion_constant, min_files_for_path): import pandas as pd @@ -364,4 +405,41 @@ def wildcardRegex(pattern): pattern = pattern.replace("*", "[^\\\\]*") pattern = pattern.replace("?", ".") pattern = pattern.replace("___RECURSIVE___", ".*") - return re.compile(f"^{pattern}$", re.IGNORECASE) \ No newline at end of file + return re.compile(f"^{pattern}$", re.IGNORECASE) + +def clean_folders_enforcement_prep(): + """ + Prompts user to choose whether to delete all .parquet files or preserve execution_history ones. + Then deletes .csv, .html, and .parquet files accordingly from specified folders. + """ + # Prompt user + user_input = input("Do you want to delete *all* .parquet files including execution_history ones? (yes/y or no/n): ").strip().lower() + delete_execution_hist = user_input in ["yes", "y"] + + folders = ["approved", "exclusions", "needs_approved", "parquet", "preflight"] + + for folder in folders: + folder_path = os.path.abspath(folder) + + if not os.path.isdir(folder_path): + print(f"Folder not found: {folder_path}") + continue + + for filename in os.listdir(folder_path): + file_path = os.path.join(folder_path, filename) + + if not os.path.isfile(file_path): + continue + + _, ext = os.path.splitext(filename) + + # Delete .csv and .html files + if ext in [".csv", ".html"]: + os.remove(file_path) + print(f"Deleted: {file_path}") + + # Delete .parquet files based on user choice + elif ext == ".parquet": + if delete_execution_hist or not filename.startswith("execution_history"): + os.remove(file_path) + diff --git a/utils/pretty.py b/utils/pretty.py index 8f6907a..3a19d2b 100644 --- a/utils/pretty.py +++ b/utils/pretty.py @@ -304,6 +304,8 @@ def printEnforceChecklist(first_policy, second_policy, allowlist_child_name, all print(colorText(f" Apply signed or attested hashes to Parent Allow List", "cyan")) print(colorText(f" Apply approved, but unsigned hashes to the Child Allow List", "cyan")) + print(colorText("R. Remove/Reset Generated data - will prompt to allow keeping execution history", "cyan")) + print(colorText("Q. Quit", "cyan"))