215 lines
8.6 KiB
Python
215 lines
8.6 KiB
Python
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
#Local Imports
|
|
import utils.hashfunctions as hashf
|
|
import utils.pathfunctions as pathf
|
|
import utils.utils as ct
|
|
|
|
#Standard Libary Imports:
|
|
import ast
|
|
import gc
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
#3rd Party Imports:
|
|
import pandas as pd
|
|
import requests
|
|
|
|
def split_filepaths_grouped(df, col="filename", group_parts=4, min_parts=4):
|
|
def clean_split(path):
|
|
parts = os.path.normpath(path).split(os.sep)
|
|
# Remove leading empty strings caused by UNC paths
|
|
parts = [p for p in parts if p]
|
|
return parts
|
|
|
|
df = df.copy()
|
|
split_paths = df[col].apply(clean_split)
|
|
|
|
# Filter out paths with fewer than `min_parts` components
|
|
df = df[split_paths.apply(lambda parts: len(parts) >= min_parts)].copy()
|
|
split_paths = split_paths[df.index] # Update split_paths to match filtered df
|
|
|
|
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:group_parts]))
|
|
grouped = df.groupby("group_key")
|
|
new_rows = []
|
|
|
|
for _, group_df in grouped:
|
|
paths = group_df[col].tolist()
|
|
split_parts = [clean_split(p) for p in paths]
|
|
|
|
def longest_common_prefix(paths):
|
|
if not paths:
|
|
return []
|
|
prefix = paths[0]
|
|
for path in paths[1:]:
|
|
prefix = [a for a, b in zip(prefix, path) if a == b]
|
|
if not prefix:
|
|
break
|
|
return prefix
|
|
|
|
common_prefix = longest_common_prefix(split_parts)
|
|
prefix_str = os.sep.join(common_prefix)
|
|
|
|
for i, parts in enumerate(split_parts):
|
|
filename = parts[-1]
|
|
middle = os.sep.join(parts[len(common_prefix):-1]) if len(parts) > len(common_prefix) + 1 else ""
|
|
row = group_df.iloc[i].copy()
|
|
row["longestcfp"] = prefix_str
|
|
row["middle"] = middle
|
|
row["filename_only"] = filename
|
|
row["file_extension"] = os.path.splitext(filename)[1].lower()
|
|
new_rows.append(row)
|
|
|
|
return pd.DataFrame(new_rows).drop(columns=["group_key"])
|
|
|
|
def inspect_parquet(path):
|
|
try:
|
|
df = pd.read_parquet(path)
|
|
print(f"✅ Successfully read: {path}")
|
|
print(f"📄 Columns: {df.columns.tolist()}")
|
|
print(f"🔢 Rows: {len(df)}")
|
|
return df
|
|
except Exception as e:
|
|
print(f"❌ Error reading {path}: {e}")
|
|
return pd.DataFrame()
|
|
|
|
def regulator(paths, case_insensitive=True):
|
|
"""
|
|
Build a regex pattern that matches any of the given Windows path fragments.
|
|
"""
|
|
escaped = [re.escape(p) for p in paths]
|
|
pattern = "(?:" + "|".join(escaped) + ")"
|
|
if case_insensitive:
|
|
pattern = "(?i)" + pattern # Add inline case-insensitive flag
|
|
print(f"Regulator is providing: {pattern}")
|
|
return pattern
|
|
|
|
def calculatePath(approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split):
|
|
|
|
if split : dfs_by_policy = [group for _, group in approved_hashes.groupby('policy')]
|
|
else : dfs_by_policy = [approved_hashes]
|
|
|
|
processed_dfs = []
|
|
|
|
for df in dfs_by_policy:
|
|
haslcp = pathf.split_filepaths_grouped(df, "filename", path_exclusion_constant, min_files_for_path)
|
|
haslcp = 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"))
|
|
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
|
|
|
lcp_not_forbidden_review = lcp_not_forbidden[['policy', 'longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']]
|
|
|
|
unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index()
|
|
unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count']
|
|
|
|
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]
|
|
processed_dfs.append(lcp_not_forbidden_review)
|
|
|
|
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
|
|
|
return pathExclusions
|
|
|
|
def generatePathReview(unknown, good, badpathparts, path_exclusion_constant, min_files_for_path, split = False):
|
|
|
|
df1 = ct.tryToReadCSV(unknown)
|
|
df2 = ct.tryToReadCSV(good)
|
|
|
|
all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename'])
|
|
|
|
primary_path_exclusions = calculatePath(all_approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split)
|
|
|
|
remaining_hashes = all_approved_hashes[~all_approved_hashes['sha256'].isin(primary_path_exclusions['sha256'])]
|
|
|
|
secondary_path_exclusions = calculatePath(remaining_hashes, badpathparts, 3, min_files_for_path, split)
|
|
|
|
remaining_hashes = remaining_hashes[~remaining_hashes['sha256'].isin(secondary_path_exclusions['sha256'])]
|
|
|
|
return all_approved_hashes, primary_path_exclusions, secondary_path_exclusions, remaining_hashes
|
|
|
|
def clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir):
|
|
"""
|
|
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 = [parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir]
|
|
|
|
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)
|
|
|
|
def generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path, split= False ):
|
|
|
|
if os.path.exists(f"{needappr_base_dir}unknown_hashes.csv") and os.path.exists(f"{needappr_base_dir}good_hashes.csv"):
|
|
|
|
all_hashes, primary_paths, secondary_paths, remaining = pathf.generatePathReview(f"{appr_base_dir}unknown_hashes.csv", f"{appr_base_dir}good_hashes.csv", badpathparts,path_exclusion_constant, min_files_for_path, split)
|
|
|
|
all_hashes.to_parquet(f"{parq_base_dir}all_hashes.parquet", index=False)
|
|
|
|
dataframes = {
|
|
"all_hashes" : all_hashes,
|
|
"primary_Paths": primary_paths,
|
|
"secondary_Paths": secondary_paths,
|
|
"remaining": remaining
|
|
|
|
}
|
|
|
|
for name, df in dataframes.items():
|
|
df.to_csv(f"{needappr_base_dir}{name}.csv", index=False)
|
|
df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
|
|
ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html")
|
|
|
|
|
|
|
|
publishers = hashf.generatePublist(f"{parq_base_dir}all_hashes.parquet",bad_publisher_list)
|
|
|
|
publishers.to_csv(f"{needappr_base_dir}publishers.csv", index=False)
|
|
publishers.to_parquet(f"{parq_base_dir}publishers.parquet", index=False)
|
|
ct.style_dataframe_dark(publishers, f"{needappr_base_dir}publishers.html")
|
|
|
|
else:
|
|
print(ct.colorText(f"Please manually approve hashes prior to this step","red"))
|
|
|