Major Refactor now allows multiple policies to be selected
This commit is contained in:
+75
-182
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user