417 lines
16 KiB
Python
417 lines
16 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 logging
|
||
import os
|
||
import os.path
|
||
import re
|
||
from typing import List
|
||
|
||
import dotenv
|
||
import pandas as pd
|
||
|
||
from models.execution import ExecutionHistoryRecord
|
||
from models.policy import Allowlist, Policy
|
||
from services.API import AirlockAPIWrapper
|
||
from utils.configmanager import get_protected_value, load_env, load_env_json
|
||
from utils.selector import Selector
|
||
from utils.utils import colorText, formatHTML, print_x_wide, regulator
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
dotenv.load_dotenv()
|
||
|
||
|
||
|
||
def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
|
||
|
||
policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
|
||
logger.debug("Prompting for Policies")
|
||
print(colorText("Please select policy/policies", "white"))
|
||
selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
|
||
|
||
if selected is None:
|
||
return []
|
||
|
||
# Normalize to always return a list
|
||
logger.debug("Returning {selected.dict}")
|
||
return selected if isinstance(selected, list) else [selected]
|
||
|
||
|
||
def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]:
|
||
if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
|
||
else: allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()]
|
||
logger.debug("Prompting for Allowlist(s)")
|
||
print(colorText("Please select allowlist(s)", "white"))
|
||
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
|
||
|
||
if selected is None:
|
||
return []
|
||
|
||
# Normalize to always return a list
|
||
logger.debug(f"Returning {selected}")
|
||
return selected if isinstance(selected, list) else [selected]
|
||
|
||
|
||
def sortHashes(
|
||
api: AirlockAPIWrapper,
|
||
selected_policies: List[Policy],
|
||
type=[1, 2, 6, 7]
|
||
):
|
||
working_dir = load_env("WORKING_DIR")
|
||
history_days = Selector.select_value(
|
||
prompt="Enter how many days of history to pull (1–150): ",
|
||
value_type=int,
|
||
valid_range=(1, 150),
|
||
)
|
||
|
||
logger.debug(f"{history_days} day selected for history")
|
||
|
||
if history_days is None:
|
||
logging.warning("No history range selected. Aborting.")
|
||
return
|
||
|
||
policy_executions = ExecutionHistoryRecord.from_policies(
|
||
api, selected_policies, type_=type, history_days=history_days
|
||
)
|
||
|
||
logger.debug(f"Executions contains {policy_executions}")
|
||
|
||
enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(api, policy_executions)
|
||
categorized_executions = ExecutionHistoryRecord.categorize_executions_by_hash_decision(enriched_executions)
|
||
approved, unapproved, needs_review, unknown = ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions)
|
||
|
||
categories = {
|
||
"needs_review": needs_review,
|
||
"approved": approved,
|
||
"unapproved": unapproved,
|
||
"leftover" : unknown
|
||
}
|
||
|
||
|
||
for label, records in categories.items():
|
||
if not records:
|
||
continue # Skip empty or falsy categories
|
||
|
||
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv"
|
||
html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html"
|
||
|
||
# Convert ExecutionHistoryRecord objects to dictionaries
|
||
df = pd.DataFrame([r.__dict__ for r in records])
|
||
|
||
# Optional: flatten hash_obj if needed
|
||
if not df.empty and 'hash_obj' in df.columns:
|
||
hash_df = df['hash_obj'].apply(lambda h: h.to_dict() if h else {})
|
||
df = pd.concat([df.drop(columns=['hash_obj']), hash_df], axis=1)
|
||
|
||
# Save to CSV
|
||
df.to_csv(csv_path, index=False)
|
||
logger.info(f"Saved {label} executions to {csv_path}")
|
||
|
||
# Generate HTML
|
||
formatHTML(df, html_path)
|
||
logger.info(f"Generated HTML report at {html_path}")
|
||
|
||
|
||
def buildPathsandPublishers(selected_policies: List[Policy], split):
|
||
working_dir = load_env("WORKING_DIR")
|
||
df1 = pd.DataFrame()
|
||
df2 = pd.DataFrame()
|
||
all_approved_hashes = pd.DataFrame()
|
||
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv"
|
||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv"
|
||
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int)
|
||
|
||
if os.path.exists(path1):
|
||
df1 = pd.read_csv(path1)
|
||
else:
|
||
logger.warning(f"File not found: {path1}")
|
||
|
||
if os.path.exists(path2):
|
||
df2 = pd.read_csv(path2)
|
||
else:
|
||
logger.warning(f"File not found: {path2}")
|
||
|
||
if df1.empty and df2.empty:
|
||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||
all_approved_hashes = pd.DataFrame()
|
||
logger.debug(all_approved_hashes.head)
|
||
else:
|
||
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
|
||
if "filename" in all_approved_hashes.columns:
|
||
all_approved_hashes = all_approved_hashes.sort_values(by="filename")
|
||
else:
|
||
logger.warning("Warning: 'filename' column not found in concatenated DataFrame.")
|
||
|
||
if not all_approved_hashes.empty and path_exclusion_constant:
|
||
|
||
primary_path_exclusions = calculatePath(
|
||
all_approved_hashes, path_exclusion_constant,
|
||
split,
|
||
)
|
||
remaining_hashes = all_approved_hashes[
|
||
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
|
||
]
|
||
secondary_path_exclusions = calculatePath(
|
||
remaining_hashes,(path_exclusion_constant - 1), split
|
||
)
|
||
remaining_hashes = remaining_hashes[
|
||
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
|
||
]
|
||
dataframes = {
|
||
"all_approved_hashes" : all_approved_hashes,
|
||
"primary_Paths": primary_path_exclusions,
|
||
"secondary_Paths": secondary_path_exclusions,
|
||
"hashes_not_approvable_by_path": remaining_hashes
|
||
}
|
||
logger.debug("Preparing to sort dataframes")
|
||
for name, df in dataframes.items():
|
||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||
if "hashes" in name : df.sort_values(by="filename", inplace=True)
|
||
else: df.sort_values(by="longestcfp", inplace=True)
|
||
|
||
df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", index=False)
|
||
formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html")
|
||
|
||
if not all_approved_hashes.empty:
|
||
# Drop all not signed, only keep unique values
|
||
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 = regulator(load_env_json("BAD_PUBLISHERS","[]"))
|
||
publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
|
||
publist = publist[["publisher"]]
|
||
publist.sort_values(by="publisher", inplace=True)
|
||
publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", index=False)
|
||
else:
|
||
logger.debug("Approved Hashes list appears empty")
|
||
|
||
def buildPreflights(selected_policies: List[Policy]):
|
||
working_dir = load_env("WORKING_DIR")
|
||
|
||
df1 = pd.DataFrame()
|
||
df2 = pd.DataFrame()
|
||
approved_hashes = pd.DataFrame()
|
||
approved_publishers = pd.DataFrame()
|
||
|
||
hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv"
|
||
path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
|
||
path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv"
|
||
publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv"
|
||
|
||
|
||
#Read in and combine the two path generations
|
||
if os.path.exists(path1):
|
||
df1 = pd.read_csv(path1)
|
||
else:
|
||
logger.warning(f"File not found: {path1}")
|
||
|
||
if os.path.exists(path2):
|
||
df2 = pd.read_csv(path2)
|
||
else:
|
||
logger.warning(f"File not found: {path2}")
|
||
|
||
if df1.empty and df2.empty:
|
||
logger.warning("Both DataFrames are empty. Skipping sort.")
|
||
approved_paths = pd.DataFrame()
|
||
else:
|
||
approved_paths = pd.concat([df1, df2], ignore_index=True)
|
||
|
||
approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep ="first")
|
||
|
||
#We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions.
|
||
if os.path.exists(hash):
|
||
hashes = pd.read_csv(hash)
|
||
approved_hashes = hashes[~hashes['filename'].isin(approved_paths['longestcfp'])]
|
||
|
||
approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep ="first")
|
||
|
||
else:
|
||
logger.warning(f"File not found: {hash}")
|
||
|
||
|
||
if os.path.exists(publishers):
|
||
approved_publishers = pd.read_csv(publishers)
|
||
|
||
else:
|
||
logger.warning(f"File not found: {publishers}")
|
||
|
||
dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers}
|
||
|
||
for name, df in dataframes.items():
|
||
logger.debug(f" DataFrame headers: {list(df.columns)}")
|
||
if name == "approved_paths":df.sort_values(by="longestcfp", inplace=True)
|
||
elif name == "approved_hashes":df.sort_values(by="filename", inplace=True)
|
||
elif name == "approved_publishers" : df.sort_values(by="publisher", inplace=True)
|
||
|
||
df.to_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", index=False)
|
||
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html")
|
||
|
||
def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"):
|
||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int)
|
||
|
||
def clean_split(path):
|
||
if not isinstance(path, (str, bytes, os.PathLike)):
|
||
return []
|
||
parts = str(os.path.normpath(path)).split(os.sep)
|
||
parts = [p for p in parts if p] # Remove empty strings
|
||
return parts
|
||
|
||
# Diagnostic: log any non-string entries
|
||
non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
|
||
if not non_string_entries.empty:
|
||
print(f"[WARNING] Non-string entries found in column '{col}':")
|
||
print(non_string_entries)
|
||
|
||
df = df.copy()
|
||
split_paths = df[col].apply(clean_split)
|
||
|
||
if min_files_for_path is not None:
|
||
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
|
||
split_paths = split_paths[df.index]
|
||
|
||
df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
|
||
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 calculatePath(approved_hashes, path_exclusion_constant, split):
|
||
if split:
|
||
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
|
||
else:
|
||
dfs_by_policy = [approved_hashes]
|
||
|
||
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
|
||
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int)
|
||
|
||
processed_dfs = []
|
||
|
||
for df in dfs_by_policy:
|
||
haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename")
|
||
haslcp = haslcp.drop_duplicates()
|
||
|
||
forbidden = regulator(badpathparts, True)
|
||
forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
|
||
|
||
logger.debug("Removing forbidden filepaths for path exceptions")
|
||
print(colorText("Removing forbidden filepaths for path exceptions", "green"))
|
||
lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
|
||
|
||
lcp_not_forbidden_review = lcp_not_forbidden[
|
||
[
|
||
"policyname",
|
||
"longestcfp",
|
||
"middle",
|
||
"filename_only",
|
||
"file_extension",
|
||
"sha256",
|
||
]
|
||
]
|
||
|
||
unique_sha_counts = (
|
||
lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index()
|
||
)
|
||
unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
|
||
|
||
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
|
||
]
|
||
processed_dfs.append(lcp_not_forbidden_review)
|
||
|
||
pathExclusions = pd.concat(processed_dfs, ignore_index=True)
|
||
|
||
return pathExclusions
|
||
|
||
def testChange(selected_policies, destination_policy, destination_allowlist):
|
||
working_dir = load_env("WORKING_DIR")
|
||
|
||
logger.info("These path exclusions would be added to:")
|
||
logger.info(destination_policy)
|
||
|
||
pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv")
|
||
hashes = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv")
|
||
|
||
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
|
||
|
||
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||
processed_paths = [
|
||
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
|
||
for path, ext in unique_combinations.itertuples(index=False, name=None)
|
||
]
|
||
|
||
for path in processed_paths:
|
||
logger.info(path)
|
||
|
||
print(colorText("These publishers would added", "yellow"))
|
||
processed_publishers = []
|
||
if os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"):
|
||
publishers = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv")
|
||
if publishers.empty:
|
||
print(colorText("The publishers list is empty.", "red"))
|
||
else:
|
||
processed_publishers = (
|
||
publishers[publishers["publisher"] != "Not Signed"]
|
||
["publisher"]
|
||
.drop_duplicates()
|
||
.tolist()
|
||
)
|
||
for publisher in processed_publishers:
|
||
print(publisher)
|
||
|
||
print(colorText("These hashes would be added to:", "yellow"))
|
||
print(destination_allowlist)
|
||
|
||
processed_hashes = hashes["sha256"].unique().tolist()
|
||
print_x_wide(processed_hashes, 3)
|
||
|
||
return processed_paths, processed_hashes, processed_publishers
|
||
|