=
136 lines
4.5 KiB
Python
136 lines
4.5 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/>.
|
|
|
|
import pandas as pd
|
|
import os
|
|
from itertools import chain
|
|
import ast
|
|
|
|
|
|
|
|
|
|
def split_path(path):
|
|
parts = []
|
|
while True:
|
|
head, tail = os.path.split(path)
|
|
if tail:
|
|
parts.insert(0, tail)
|
|
path = head
|
|
else:
|
|
if head:
|
|
parts.insert(0, head)
|
|
break
|
|
return parts
|
|
|
|
def local_common_pass(paths, min_parts=3):
|
|
results = {}
|
|
paths_sorted = sorted(paths)
|
|
for i, path in enumerate(paths_sorted):
|
|
candidates = []
|
|
|
|
if i > 0:
|
|
try:
|
|
candidates.append(os.path.commonpath([path, paths_sorted[i-1]]))
|
|
except ValueError:
|
|
# different drives, skip
|
|
pass
|
|
if i < len(paths_sorted) - 1:
|
|
try:
|
|
candidates.append(os.path.commonpath([path, paths_sorted[i+1]]))
|
|
except ValueError:
|
|
# different drives, skip
|
|
pass
|
|
|
|
best = path
|
|
best_len = 0
|
|
for c in candidates:
|
|
parts = split_path(c)
|
|
if len(parts) >= min_parts and len(parts) > best_len:
|
|
best = c
|
|
best_len = len(parts)
|
|
results[path] = best
|
|
return results
|
|
|
|
|
|
def add_longest_common_two_local(df, col="filename_x", new_col="longestcfp", min_parts=3):
|
|
dirs_series = df[col].astype(str).apply(os.path.dirname)
|
|
first_pass = local_common_pass(dirs_series.tolist(), min_parts)
|
|
second_pass = local_common_pass(list(first_pass.values()), min_parts)
|
|
df[new_col] = dirs_series.map(lambda d: second_pass[first_pass[d]])
|
|
return df
|
|
|
|
|
|
def export_groups_for_review(df, col, group_col, min_number_in_group, path_length_constant):
|
|
"""
|
|
Compute longest common paths, group filepaths, write CSV for review.
|
|
"""
|
|
df = add_longest_common_two_local(df, col=col, new_col=group_col)
|
|
grouped = df.groupby(group_col)[col].apply(list).reset_index()
|
|
grouped = grouped.sort_values(by=col)
|
|
print("Before filtering:", len(grouped))
|
|
grouped = grouped[grouped[col].apply(lambda x: len(x) >= min_number_in_group)]
|
|
filtered = grouped[grouped[group_col].apply(lambda x: len(os.path.normpath(x).split(os.sep)) >= path_length_constant)]
|
|
print("After filtering:", len(grouped))
|
|
|
|
return filtered, df
|
|
|
|
|
|
|
|
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 masked_df, 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
|