Files
AirlockTools/utils/pathfunctions.py
T
2025-09-18 15:19:19 -04:00

324 lines
13 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.pretty as ct
from AirlockTools import tryToReadCSV
#Standard Libary Imports:
import ast
import gc
import os
import re
#3rd Party Imports:
import pandas as pd
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 mask_from_csv(df, csv_path, filepath_col):
"""
Reads reviewed CSV of groups, keeps only files in approved groups.
"""
review_df = pd.read_csv(csv_path)
def parse_paths(val):
if isinstance(val, str):
try:
# Try to parse as a list
parsed = ast.literal_eval(val)
# If it's not a list, wrap it
return parsed if isinstance(parsed, list) else [parsed]
except (ValueError, SyntaxError):
# If parsing fails, treat it as a single path
return [val]
return [val]
review_df[filepath_col] = review_df[filepath_col].apply(parse_paths)
# Flatten all approved file paths into a set for masking
approved_files = set()
for paths in review_df[filepath_col]:
approved_files.update(paths)
# Keep only rows in df that are in approved_files
masked_df = df[df[filepath_col].isin(approved_files)].copy()
remainder = df[~df[filepath_col].isin(approved_files)].copy()
return remainder
def filter_and_drop(approved, eligiblepaths, min_hashes):
"""
Filters eligiblepaths to rows where all hashes are in approved,
then drops rows with fewer than min_hashes hashes.
"""
approved_hashes = set(approved['sha256'])
def all_hashes_approved(row):
return all(h in approved_hashes for h in row['sha256'])
filtered = eligiblepaths[eligiblepaths.apply(all_hashes_approved, axis=1)]
filtered = filtered[filtered['sha256'].apply(len) >= min_hashes]
return filtered
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 generatePathReview(first_policy, second_policy, badpathparts, path_exclusion_constant, min_files_for_path):
if not os.path.exists(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet"):
df1 = tryToReadCSV(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv")
df2 = tryToReadCSV(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv")
all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename'])
print(ct.colorText(f"Approved hash lists have been combined","green"))
all_approved_hashes.to_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet", index=False)
del all_approved_hashes
gc.collect()
if not os.path.exists(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet"):
all_approved_hashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet")
print(ct.colorText(f"Beginning calculating longest common filepaths for path exceptions","green"))
haslcp = pathf.split_filepaths_grouped(all_approved_hashes,"filename",path_exclusion_constant, 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\\path_needs_approved_{first_policy}_{second_policy}.parquet",index=False)
lcp_not_forbidden_review.to_csv(f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.csv",index=False)
ct.style_dataframe_dark(lcp_not_forbidden_review,f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.html", True)
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
# Load raw data
hashes = pd.read_csv("allATPolicyExecs.csv")
# Deduplicate hashes before augmentation
deduped_hashes = hashes.drop_duplicates(subset=['sha256']).copy()
# Save policy name mapping (before deduplication)
policyname_map = hashes[['hostname', 'PolicyName']].drop_duplicates()
# Handle None inputs
untrusted_publishers = untrusted_publishers or []
pups = pups or []
# Augment deduplicated hashes
augmented = hashf.augmentAggregatedHashes(url, deduped_hashes)
# Merge policy names
augmented = augmented.merge(policyname_map, on='hostname', how='left')
# Clean numeric reputation fields
for col in ['reputation_scannermatch', 'reputation_scannercount', 'reputation_threatlevel']:
if col in augmented.columns:
augmented[col] = pd.to_numeric(augmented[col].replace('N/A', pd.NA), errors='coerce')
# Rename and select relevant columns
augmented = augmented.rename(columns={'publisher_x': 'publisher'})
augmented = augmented[[
'PolicyName', 'sha256', 'publisher', 'description', 'productname', 'productversion',
'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
'reputation_timestamp'
]].sort_values(by=['publisher', 'description', 'productname'])
# Merge with deduplicated hashes to enrich data
final_augmented = augmented.merge(deduped_hashes, on='sha256', how='left')
# Clean up column names before applying reputation logic
if 'publisher_y' in final_augmented.columns:
final_augmented = final_augmented.drop(columns=['publisher_y'])
if 'publisher_x' in final_augmented.columns:
final_augmented = final_augmented.rename(columns={'publisher_x': 'publisher'})
if 'PolicyName_x' in final_augmented.columns:
final_augmented = final_augmented.rename(columns={'PolicyName_x': 'PolicyName'})
final_augmented.to_csv("testing.csv", index=False)
# Reputation flag logic
def reputationtool(row):
val = row["reputation_scannermatch"]
if pd.isna(val):
return row["publisher"] == "Not Signed"
try:
return int(val) > threat_tolerance
except (ValueError, TypeError):
return row["publisher"] == "Not Signed"
df = final_augmented.copy()
df["reputation_flag"] = df.apply(reputationtool, axis=1)
# Filtering logic
mask_needsreview = (
((df["publisher"] == "Not Signed") & df["reputation_flag"]) |
(df["reputation_status"] == "UNKNOWN")
)
mask_approved = (
(
(df["publisher"] != "Not Signed") &
~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
~df["reputation_status"].isna() &
~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
) |
(
(df["publisher"] == "Not Signed") &
~df["reputation_flag"] &
~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
~df["reputation_status"].isna() &
~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
)
)
needsreview_df = df[mask_needsreview]
approved_df = df[mask_approved]
all_approved_hashes = pd.concat([needsreview_df, approved_df], ignore_index=True)
print(ct.colorText("Beginning calculating longest common filepaths for path exceptions", "green"))
# Path analysis
haslcp = pathf.split_filepaths_grouped(all_approved_hashes, "filename", path_exclusion_constant, min_files_for_path)
haslcp = haslcp.drop_duplicates()
# Remove forbidden paths
forbidden = pathf.regulator(badpathparts, True)
lcp_not_forbidden = haslcp[~haslcp["longestcfp"].str.contains(forbidden, na=False)].copy()
print(ct.colorText("Removing forbidden filepaths for path exceptions", "green"))
# Reviewable paths
review_df = lcp_not_forbidden[['longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']]
sha_counts = review_df.groupby('longestcfp')['sha256'].nunique().reset_index()
sha_counts.columns = ['longestcfp', 'unique_sha256_count']
review_df = review_df.merge(sha_counts, on='longestcfp', how='left')
review_df = review_df[review_df['unique_sha256_count'] >= min_files_for_path]
review_df.to_csv("ALL_AT_PATHS.csv", index=False)
# Cleanup
del lcp_not_forbidden, sha_counts, review_df
def mergeTesting():
# Load the two CSVs
testing_df = pd.read_csv("testing.csv")
paths_df = pd.read_csv("ALL_AT_PATHS.csv")
# Merge on 'sha256' with testing as the left DataFrame
merged_df = testing_df.merge(paths_df, on="sha256", how="left")
# Save the merged result
merged_df.to_csv("merged_output.csv", index=False)
print(f"Merged DataFrame saved with {len(merged_df)} rows.")