Major Refactor now allows multiple policies to be selected

This commit is contained in:
2025-09-26 14:57:12 -04:00
parent 74b3f7b2e1
commit 0a26ed71c1
8 changed files with 777 additions and 1133 deletions
+75 -182
View File
@@ -14,10 +14,8 @@
# 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.pretty as ct
from AirlockTools import tryToReadCSV
import utils.utils as ct
#Standard Libary Imports:
import gc
@@ -99,7 +97,7 @@ def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
df = agg_df.merge(df_api, on="sha256", how="left")
# Only include columns that exist to avoid KeyErrors
expected_columns = ['sha256', 'filename_x', 'description', 'productname', 'productversion',
expected_columns = ['policy','sha256', 'filename_x', 'description', 'productname', 'productversion',
'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username',
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
@@ -110,7 +108,7 @@ def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
return aug_df
def categorizeHashes(first_policy, second_policy, df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list):
def categorizeHashes(df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
if untrusted_publishers is None: untrusted_publishers = []
if pups is None: pups = []
@@ -149,122 +147,56 @@ def categorizeHashes(first_policy, second_policy, df: pd.DataFrame, threat_toler
needsreview_df = df[mask_needsreview]
approved_df = df[mask_approved]
unapproved_df = df[~(mask_needsreview | mask_approved)]
return needsreview_df, approved_df, unapproved_df
needsreview_df.to_parquet(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet", index=False)
approved_df.to_parquet(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet", index=False)
unapproved_df.to_parquet(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet", index=False)
del needsreview_df
del approved_df
del unapproved_df
gc.collect()
def explode_and_deduplicate(df):
df['sha256'] = df['sha256'].str.split(',')
df = df.explode('sha256')
return df.drop_duplicates().reset_index(drop=True)
def clean_sha256(df, column='sha256'):
"""Discard quotes, brackets, and whitespace from sha256 values."""
df[column] = df[column].astype(str).str.strip("'[]\" ")
return df
def destinationHashes(
df_approved_paths: pd.DataFrame,
df_approved_hashes: pd.DataFrame,
df_hashes_auto_approved: pd.DataFrame,
df_hashes_manually_approved: pd.DataFrame,
):
# Deduplicate and explode all input DataFrames
df_approved_paths = explode_and_deduplicate(df_approved_paths)
df_approved_hashes = explode_and_deduplicate(df_approved_hashes)
df_hashes_auto_approved = explode_and_deduplicate(df_hashes_auto_approved)
df_hashes_manually_approved = explode_and_deduplicate(df_hashes_manually_approved)
# Clean sha256 values in all relevant DataFrames
df_approved_hashes = clean_sha256(df_approved_hashes)
df_hashes_auto_approved = clean_sha256(df_hashes_auto_approved)
df_hashes_manually_approved = clean_sha256(df_hashes_manually_approved)
# Create sets for faster lookup
auto_approved_sha256 = set(df_hashes_auto_approved['sha256'].values)
manually_approved_sha256 = set(df_hashes_manually_approved['sha256'].values)
# Debug: Print unmatched hashes
unmatched = set(df_approved_hashes['sha256']) - (auto_approved_sha256 | manually_approved_sha256)
print(f"Unmatched hashes: {unmatched}")
# Process df_approved_paths
df_paths = df_approved_paths.assign(destination='Path Exclusion')
df_paths = df_paths[['sha256', 'description', 'destination', 'grouped_directory', 'filename']]
# Process df_approved_hashes
df_hashes = df_approved_hashes.copy()
df_hashes['destination'] = df_hashes['sha256'].apply(
lambda x: 'Parent Policy Baseline' if x in auto_approved_sha256
else ('Child Policy Allowlist' if x in manually_approved_sha256 else None)
)
df_hashes = df_hashes.dropna(subset=['destination'])
df_hashes = df_hashes.assign(grouped_directory=None)
# Use 'filename_x' only if it exists, otherwise fallback to 'filename'
filename_col = 'filename_x' if 'filename_x' in df_hashes.columns else 'filename'
selected_cols = ['sha256', 'description', 'destination', 'grouped_directory', filename_col]
df_hashes = df_hashes[selected_cols]
# Concatenate results
df_hashdestination = pd.concat([df_paths, df_hashes], ignore_index=True)
return df_hashdestination
def combineHashAndHist(path, first_policy, second_policy):
condensed_combo = pd.read_parquet(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet")
df = pd.read_parquet(path)
#Pull hash info for the entries in the needs approval table
def combineHashAndHist(hash_path, condensed_path):
# Load both datasets
condensed_combo = pd.read_parquet(condensed_path)
df = pd.read_parquet(hash_path)
# Merge on sha256
df = pd.merge(condensed_combo, df, on='sha256', how='inner')
#Rename Publisher, Keep and reorder columns we want
df.to_csv("testing4.csv")
# Rename and reorder columns
df = df.rename(columns={'publisher_x': 'publisher'})
df = df[['sha256', 'publisher', 'description', 'filename', 'hostname', 'username', 'productname', 'productversion','reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount','reputation_status', 'reputation_threatlevel', 'reputation_threatname','reputation_timestamp', 'pprocess', 'gprocess', 'commandline']]
df = df.rename(columns={'policy_x': 'policy'})
df = df[['policy','sha256', 'publisher', 'description', 'filename', 'hostname', 'username',
'productname', 'productversion', 'reputation_lastseen', 'reputation_scannermatch',
'reputation_scannercount', 'reputation_status', 'reputation_threatlevel',
'reputation_threatname', 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']]
df = df.sort_values(by='filename')
df.to_parquet(path, index=False)
# Overwrite the original hash file
df.to_parquet(hash_path, index=False)
# Cleanup
del df
del condensed_combo
gc.collect()
def combineHashes(url, first_policy, second_policy):
combined_hashes = pd.DataFrame(columns=['sha256', 'publisher'])
def combineHashes(url, parquet_files) -> pd.DataFrame:
combined_hashes = pd.DataFrame()
hashes = []
try:
hash1 = pd.read_parquet(f"parquet\\execution_history_{first_policy}.parquet", columns=['sha256', 'publisher'])
pathf.inspect_parquet(f"parquet\\execution_history_{first_policy}.parquet")
if not hash1.empty:
hashes.append(hash1)
else:
print("⚠️ First dataframe is empty.")
except Exception as e:
print(f"❌ Error reading first Parquet file: {e}")
try:
hash2 = pd.read_parquet(f"parquet\\execution_history_{second_policy}.parquet", columns=['sha256', 'publisher'])
pathf.inspect_parquet(f"parquet\\execution_history_{second_policy}.parquet")
if not hash2.empty:
hashes.append(hash2)
else:
print("⚠️ Second dataframe is empty.")
except Exception as e:
print(f"❌ Error reading second Parquet file: {e}")
for file_path in parquet_files:
try:
hash_df = pd.read_parquet(file_path)
pathf.inspect_parquet(file_path)
if not hash_df.empty:
hashes.append(hash_df)
else:
print(f"⚠️ Dataframe is empty: {file_path}")
except Exception as e:
print(f"❌ Error reading Parquet file '{file_path}': {e}")
if hashes:
combined_hashes = pd.concat(hashes, ignore_index=True)
print(f"✅ Combined {len(combined_hashes)} hashes.")
print(f"✅ Combined {len(combined_hashes)} hashes from {len(hashes)} files.")
else:
print("⚠️ No valid dataframes to combine.")
combined_hashes = combined_hashes.drop_duplicates(subset=['sha256'])
augmented_combo = hashf.augmentAggregatedHashes(url, combined_hashes)
augmented_combo = augmentAggregatedHashes(url, combined_hashes)
numeric_reputation_cols = [
'reputation_scannermatch',
@@ -282,62 +214,45 @@ def combineHashes(url, first_policy, second_policy):
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
'reputation_timestamp']]
augmented_combo = augmented_combo.sort_values(by=['publisher', 'description', 'productname'])
augmented_combo.to_parquet(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet", index=False)
del combined_hashes
del augmented_combo
gc.collect()
print(ct.colorText("Hash reputation info added to dataframe", "green"))
return augmented_combo
def condenseExecutions(first_policy,second_policy):
exe1 = pd.DataFrame()
exe2 = pd.DataFrame()
condensed_combo = pd.DataFrame()
try:
exe1 = pd.read_parquet(f"parquet\\execution_history_{first_policy}.parquet")
pathf.inspect_parquet(f"parquet\\execution_history_{first_policy}.parquet")
if not exe1.empty:
print()
else:
print("⚠️ First dataframe is empty.")
except Exception as e:
print(f"❌ Error reading first Parquet file: {e}")
def condenseExecutions(parquet_paths):
combined_df = pd.DataFrame()
valid_files = []
try:
exe2 = pd.read_parquet(f"parquet\\execution_history_{second_policy}.parquet")
pathf.inspect_parquet(f"parquet\\execution_history_{second_policy}.parquet")
if not exe2.empty:
print()
else:
print("⚠️ Second dataframe is empty.")
except Exception as e:
print(f"❌ Error reading second Parquet file: {e}")
for file_path in parquet_paths:
try:
df = pd.read_parquet(file_path)
# Optional: pathf.inspect_parquet(file_path)
if not df.empty:
combined_df = pd.concat([combined_df, df], ignore_index=True)
valid_files.append(file_path)
print(f"✅ Loaded {len(df)} rows from {file_path}")
else:
print(f"⚠️ DataFrame from '{file_path}' is empty.")
except Exception as e:
print(f"❌ Error reading Parquet file '{file_path}': {e}")
if not exe1.empty and not exe2.empty:
condensed_combo = pd.concat([exe1, exe2], ignore_index=True)
print(f"✅ Combined {len(condensed_combo)} hashes.")
elif exe1.empty:
condensed_combo = exe2
elif exe2.empty:
condensed_combo = exe1
if not combined_df.empty:
print(f"✅ Combined {len(combined_df)} rows from {len(valid_files)} files.")
else:
print("⚠️ No valid dataframes to combine.")
condensed_combo.to_parquet(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet", index=False)
del condensed_combo
gc.collect()
def divideSortedHashExecutions(first_policy,second_policy, pups):
return combined_df
combineHashAndHist(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet", first_policy, second_policy)
combineHashAndHist(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet", first_policy, second_policy)
combineHashAndHist(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet", first_policy, second_policy)
def divideSortedHashExecutions(unknown_parq, good_parq, bad_parq, condensed_parq, pups) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
# Run combineHashAndHist on each file
combineHashAndHist(unknown_parq, condensed_parq)
combineHashAndHist(good_parq, condensed_parq)
combineHashAndHist(bad_parq, condensed_parq)
unknown = pd.read_parquet(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet")
good = pd.read_parquet(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet")
bad = pd.read_parquet(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet")
# Load data
unknown = pd.read_parquet(unknown_parq)
good = pd.read_parquet(good_parq)
bad = pd.read_parquet(bad_parq)
# Build regex pattern once
pattern = pathf.regulator(pups)
@@ -353,52 +268,30 @@ def divideSortedHashExecutions(first_policy,second_policy, pups):
unknown = unknown[~unknown["filename"].str.contains(pattern, na=False)]
good = good[~good["filename"].str.contains(pattern, na=False)]
unknown.to_csv(f"needs_approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv",index=False)
good.to_csv(f"needs_approved\\hashes_rep_good_{first_policy}_{second_policy}.csv",index=False)
bad.to_csv(f"needs_approved\\hashes_rep_bad_{first_policy}_{second_policy}.csv",index=False)
return unknown, good, bad
ct.style_dataframe_dark(unknown, f"needs_approved\\hashes_rep_unknown_{first_policy}_{second_policy}.html")
ct.style_dataframe_dark(good, f"needs_approved\\hashes_rep_good_{first_policy}_{second_policy}.html")
ct.style_dataframe_dark(bad, f"needs_approved\\hashes_rep_bad_{first_policy}_{second_policy}.html")
def generatePreflights(first_policy, second_policy):
allhashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet")
primarypathexclusions = tryToReadCSV(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv")
secondarypathexclusions = tryToReadCSV(f"approved\\secondary_paths_{first_policy}_{second_policy}.csv")
def generatePreflights(hashes, primary_path, secondary_path):
all_hashes = pd.read_parquet(hashes)
primarypathexclusions = ct.tryToReadCSV(primary_path)
secondarypathexclusions = ct.tryToReadCSV(secondary_path)
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)
allowbyhash = allhashes[~allhashes['sha256'].isin(pathexclusions['sha256'])]
allowbyhash.to_parquet(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet", index=False)
allowbyhash = all_hashes[~all_hashes['sha256'].isin(pathexclusions['sha256'])]
allowbyhash.sort_values(by=["filename"])
ct.style_dataframe_dark(allowbyhash, f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html")
ct.style_dataframe_dark(pathexclusions, f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html")
ct.style_dataframe_dark(publishers, f"preflight\\publishers_{first_policy}_{second_policy}.html")
return pathexclusions, allowbyhash
del allowbyhash
del pathexclusions
gc.collect()
def generatePublist(first_policy,second_policy,bad_publisher_list):
try:
publist = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet", columns=['publisher'])
except Exception as e:
print(f"Error reading parquet file: {e}")
publist = pd.DataFrame()
def generatePublist(all_hashes, bad_publisher_list):
all_approved_hashes = ct.tryToReadParquet(all_hashes)
#Drop all not signed, only keep unique values
publist = publist[publist['publisher'] != "Not Signed"].drop_duplicates(subset='publisher')
publist = all_approved_hashes[all_approved_hashes['publisher'] != "Not Signed"].drop_duplicates(subset='publisher')
#Remove Bad publisher if somehow they made it this far
pattern = pathf.regulator(bad_publisher_list)
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
publist.to_csv(f"needs_approved\\publishers_{first_policy}_{second_policy}.csv", index=False)
return publist