Cleaned up menu by functionalizing
This commit is contained in:
=
+159
-8
@@ -12,15 +12,15 @@
|
||||
#
|
||||
# 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/>.
|
||||
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import pandas as pd
|
||||
import requests
|
||||
import os
|
||||
import json
|
||||
import utils.pathfunctions as pathf
|
||||
import utils.hashfunctions as hashf
|
||||
import utils.pretty as ct
|
||||
import gc
|
||||
|
||||
from AirlockTools import tryToReadCSV
|
||||
|
||||
def aggregateHashes(executions_json) -> pd.DataFrame:
|
||||
"""
|
||||
@@ -104,7 +104,7 @@ def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
|
||||
|
||||
return aug_df
|
||||
|
||||
def categorizeHashes(df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list):
|
||||
def categorizeHashes(first_policy, second_policy, df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list):
|
||||
if untrusted_publishers is None: untrusted_publishers = []
|
||||
if pups is None: pups = []
|
||||
|
||||
@@ -144,9 +144,14 @@ def categorizeHashes(df: pd.DataFrame, threat_tolerance: int, untrusted_publishe
|
||||
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(',')
|
||||
@@ -221,4 +226,150 @@ def combineHashAndHist(path, first_policy, second_policy):
|
||||
df.to_parquet(path, index=False)
|
||||
del df
|
||||
del condensed_combo
|
||||
gc.collect()
|
||||
|
||||
def combineHashes(url, first_policy, second_policy):
|
||||
combined_hashes = pd.DataFrame(columns=['sha256', 'publisher'])
|
||||
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}")
|
||||
|
||||
if hashes:
|
||||
combined_hashes = pd.concat(hashes, ignore_index=True)
|
||||
print(f"✅ Combined {len(combined_hashes)} hashes.")
|
||||
else:
|
||||
print("⚠️ No valid dataframes to combine.")
|
||||
|
||||
combined_hashes = combined_hashes.drop_duplicates(subset=['sha256'])
|
||||
augmented_combo = hashf.augmentAggregatedHashes(url, combined_hashes)
|
||||
|
||||
numeric_reputation_cols = [
|
||||
'reputation_scannermatch',
|
||||
'reputation_scannercount',
|
||||
'reputation_threatlevel'
|
||||
]
|
||||
|
||||
for col in numeric_reputation_cols:
|
||||
if col in augmented_combo.columns:
|
||||
augmented_combo[col] = pd.to_numeric(augmented_combo[col].replace('N/A', pd.NA), errors='coerce')
|
||||
|
||||
augmented_combo = augmented_combo.rename(columns={'publisher_x': 'publisher'})
|
||||
augmented_combo = augmented_combo[['sha256', 'publisher', 'description', 'productname', 'productversion',
|
||||
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
|
||||
'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"))
|
||||
|
||||
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}")
|
||||
|
||||
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}")
|
||||
|
||||
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
|
||||
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):
|
||||
|
||||
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)
|
||||
|
||||
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")
|
||||
|
||||
# Build regex pattern once
|
||||
pattern = pathf.regulator(pups)
|
||||
|
||||
# Move matching rows from unknown and good to bad
|
||||
bad = pd.concat([
|
||||
bad,
|
||||
unknown[unknown["filename"].str.contains(pattern, na=False)],
|
||||
good[good["filename"].str.contains(pattern, na=False)]
|
||||
], ignore_index=True)
|
||||
|
||||
# Remove matching rows from unknown and good
|
||||
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)
|
||||
|
||||
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")
|
||||
|
||||
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)
|
||||
|
||||
allowbyhash = allhashes[~allhashes['sha256'].isin(pathexclusions['sha256'])]
|
||||
|
||||
allowbyhash.to_parquet(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet", index=False)
|
||||
|
||||
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")
|
||||
|
||||
del allowbyhash
|
||||
del pathexclusions
|
||||
gc.collect()
|
||||
Reference in New Issue
Block a user