First Round Async. Much work left to do, dont trust results of hash categorization presently.

This commit is contained in:
2025-10-16 16:43:56 -04:00
parent fa0c18ee02
commit 6f2355fea9
21 changed files with 903 additions and 1647 deletions
+181 -185
View File
@@ -13,10 +13,11 @@
# 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 asyncio
import logging
import os
import os.path
from typing import List
from typing import List, Optional
import dotenv
import pandas as pd
@@ -24,12 +25,12 @@ import pandas as pd
from models.execution import ExecutionHistoryRecord, Hash
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread
from utils.configmanager import get_protected_value, load_env, load_env_json
from utils.selector import Selector
from utils.Selector import Selector
from utils.utils import (
colorText,
formatHTML,
import_to_dataframe,
regulator,
)
@@ -37,79 +38,86 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv()
async def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
df = await api.policy_find_all()
policies = [Policy(**row.to_dict()) for _, row in df.iterrows()]
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)
selected = await 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}")
logger.debug("Returning selected policies")
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()]
async def selectAllowlists(api: AirlockAPIWrapper, policy="all", allow_multiple=True) -> List[Allowlist]:
if policy == "all":
df = await api.allowlist_find_all()
else:
df = await api.policy_list_allowlists(policy[0].groupid) # pyright: ignore[reportAttributeAccessIssue]
allowlists = [Allowlist(**row.to_dict()) for _, row in df.iterrows()]
logger.debug("Prompting for Allowlist(s)")
print(colorText("Please select allowlist(s)", "white"))
selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
selected = await 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(
async def sortHashes(
api: AirlockAPIWrapper,
queue: AsyncTaskQueue,
selected_policies: List[Policy],
type=[1, 2, 6, 7]
type=[1, 2, 6, 7],
history_days: Optional[int] = None
):
working_dir = load_env("WORKING_DIR")
history_days = Selector.select_value(
prompt="Enter how many days of history to pull (1150): ",
value_type=int,
valid_range=(1, 150),
)
if history_days is None:
history_days = await Selector.select_value(
prompt="Enter how many days of history to pull (1150): ",
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
executions = []
hashes = []
working_dir = await load_env("WORKING_DIR")
# Pull execution histories for each policy
policy_executions = ExecutionHistoryRecord.from_policies(
policy_executions = await ExecutionHistoryRecord.from_policies(
api, selected_policies, type_=type, history_days=history_days
)
logger.debug(f"Policy_executions is {policy_executions}")
logger.debug(f"Policy_executions is {policy_executions}")
executions.extend(policy_executions)
logger.debug(f"Executions contains {executions}")
if executions:
hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
sha_list = [record.sha256 for record in executions]
hash_df = await api.hash_query(sha_list)
hashes = [
Hash(sha256=row["sha256"], **row["data"])
for _, row in hash_df.iterrows()
]
if hashes:
unique_hashes = Hash.deduplicate(hashes)
needs_review, approved, unapproved = Hash.categorize_hashes(
hashes=unique_hashes
)
needs_review, approved, unapproved = await Hash.categorize_hashes(hashes=unique_hashes)
categories = {
"needs_review": needs_review,
@@ -118,157 +126,158 @@ def sortHashes(
}
for label, category in categories.items():
csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv"
html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
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"
ExecutionHistoryRecord.enrich_with_hashes_and_export(
executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
)
df = import_to_dataframe(csv_path)
formatHTML(df, html_path)
df = await ExecutionHistoryRecord.enrich_with_hashes(executions, category)
asyncio.create_task(queue.enqueue(
f"DF TO CSV {selected_policies[0].name}_{label}",
run_sync_task_in_thread,
df.to_csv,
csv_path,
index=False,
encoding='utf-8'
))
def buildPathsandPublishers(split):
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
all_approved_hashes = pd.DataFrame()
path1 = f"{working_dir}\\Approved\\approved_executions.csv"
path2 = f"{working_dir}\\Approved\\needs_review_executions.csv"
if os.path.exists(path1):
df1 = pd.read_csv(path1)
else:
logger.warning(f"File not found: {path1}")
asyncio.create_task(queue.enqueue(
f"DF TO HTML {selected_policies[0].name}_{label}",
run_sync_task_in_thread,
formatHTML,
df,
html_path
))
if os.path.exists(path2):
df2 = pd.read_csv(path2)
else:
logger.warning(f"File not found: {path2}")
print("sortHashes completed successfully.")
async def buildPathsandPublishers(queue: AsyncTaskQueue, split):
working_dir = await load_env("WORKING_DIR")
path1 = f"{working_dir}/Approved/approved_executions.csv"
path2 = f"{working_dir}/Approved/needs_review_executions.csv"
df1 = await asyncio.to_thread(pd.read_csv, path1) if os.path.exists(path1) else pd.DataFrame()
if df1.empty:
logger.warning(f"File not found or empty: {path1}")
df2 = await asyncio.to_thread(pd.read_csv, path2) if os.path.exists(path2) else pd.DataFrame()
if df2.empty:
logger.warning(f"File not found or empty: {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)
logger.debug(all_approved_hashes.head())
return
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
if "filename_exec" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
else:
all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
if "filename_exec" in all_approved_hashes.columns:
all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
else:
logger.warning("Warning: 'filename_exec' column not found in concatenated DataFrame.")
logger.warning("'filename_exec' column not found in concatenated DataFrame.")
if not all_approved_hashes.empty:
primary_path_exclusions = calculatePath(
all_approved_hashes,
split,
)
remaining_hashes = all_approved_hashes[
~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
]
secondary_path_exclusions = calculatePath(
remaining_hashes, split
)
remaining_hashes = remaining_hashes[
~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
]
dataframes = {
"primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions,
"hashes_to_add": remaining_hashes,
}
logger.debug("Preparing to sort dataframes")
for name, df in dataframes.items():
logger.debug(f" DataFrame headers: {list(df.columns)}")
if name == "hashes_to_add": df.sort_values(by="filename_exec", inplace=True)
else: df.sort_values(by="longestcfp", inplace=True)
df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False)
formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html")
primary_path_exclusions = await calculatePath(all_approved_hashes, split)
remaining_hashes = all_approved_hashes[~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])]
secondary_path_exclusions = await calculatePath(remaining_hashes, split)
remaining_hashes = remaining_hashes[~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])]
if not all_approved_hashes.empty:
# Drop all not signed, only keep unique values
publist = all_approved_hashes[
all_approved_hashes["publisher_hash"] != "Not Signed"
].drop_duplicates(subset=["publisher_hash"])
# Remove Bad publisher if somehow they made it this far
pattern = regulator(load_env_json("BAD_PUBLISHERS","[]"))
publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
publist = publist[["publisher_hash"]]
publist.sort_values(by="publisher_hash", inplace=True)
publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False)
def buildPreflights():
working_dir = load_env("WORKING_DIR")
df1 = pd.DataFrame()
df2 = pd.DataFrame()
approved_hashes = pd.DataFrame()
approved_publishers = pd.DataFrame()
hash = f"{working_dir}\\Approved\\hashes_to_add.csv"
path1 = f"{working_dir}\\Approved\\primary_Paths.csv"
path2 = f"{working_dir}\\Approved\\secondary_Paths.csv"
publishers = f"{working_dir}\\Approved\\publishers.csv"
if os.path.exists(hash):
approved_hashes = pd.read_csv(hash)
else:
logger.warning(f"File not found: {hash}")
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)
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}
dataframes = {
"primary_Paths": primary_path_exclusions,
"secondary_Paths": secondary_path_exclusions,
"hashes_to_add": remaining_hashes,
}
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_exec", inplace=True)
elif name == "approved_publishers" : df.sort_values(by="publisher_hash", inplace=True)
df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False)
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")
logger.debug(f"DataFrame headers for {name}: {list(df.columns)}")
sort_column = "filename_exec" if name == "hashes_to_add" else "longestcfp"
df.sort_values(by=sort_column, inplace=True)
csv_path = f"{working_dir}/Needs_Review/Review_Second/{name}.csv"
html_path = f"{working_dir}/Needs_Review/HTML/{name}.html"
await asyncio.to_thread(df.to_csv, csv_path, index=False)
await asyncio.to_thread(formatHTML, df, html_path)
def splitFilepathsGrouped(df, col="filename"):
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int)
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int)
publist = all_approved_hashes[all_approved_hashes["publisher_hash"] != "Not Signed"].drop_duplicates(subset=["publisher_hash"])
pattern = regulator(await load_env_json("BAD_PUBLISHERS", "[]"))
publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
publist = publist[["publisher_hash"]]
publist.sort_values(by="publisher_hash", inplace=True)
pub_csv_path = f"{working_dir}/Needs_Review/Review_Second/publishers.csv"
await asyncio.to_thread(publist.to_csv, pub_csv_path, index=False)
print("buildPathsandPublishers completed asynchronously.")
async def buildPreflights():
working_dir = await load_env("WORKING_DIR")
hash_path = f"{working_dir}/Approved/hashes_to_add.csv"
path1 = f"{working_dir}/Approved/primary_Paths.csv"
path2 = f"{working_dir}/Approved/secondary_Paths.csv"
publishers_path = f"{working_dir}/Approved/publishers.csv"
df1 = await asyncio.to_thread(pd.read_csv, path1) if os.path.exists(path1) else pd.DataFrame()
if df1.empty:
logger.warning(f"File not found or empty: {path1}")
df2 = await asyncio.to_thread(pd.read_csv, path2) if os.path.exists(path2) else pd.DataFrame()
if df2.empty:
logger.warning(f"File not found or empty: {path2}")
approved_hashes = await asyncio.to_thread(pd.read_csv, hash_path) if os.path.exists(hash_path) else pd.DataFrame()
if approved_hashes.empty:
logger.warning(f"File not found or empty: {hash_path}")
approved_publishers = await asyncio.to_thread(pd.read_csv, publishers_path) if os.path.exists(publishers_path) else pd.DataFrame()
if approved_publishers.empty:
logger.warning(f"File not found or empty: {publishers_path}")
approved_paths = pd.concat([df1, df2], ignore_index=True) if not (df1.empty and df2.empty) else pd.DataFrame()
dataframes = {
"approved_paths": approved_paths,
"approved_hashes": approved_hashes,
"approved_publishers": approved_publishers
}
for name, df in dataframes.items():
logger.debug(f"DataFrame headers for {name}: {list(df.columns)}")
if name == "approved_paths":
df.sort_values(by="longestcfp", inplace=True)
elif name == "approved_hashes":
df.sort_values(by="filename_exec", inplace=True)
elif name == "approved_publishers":
df.sort_values(by="publisher_hash", inplace=True)
csv_path = f"{working_dir}/Preflight/{name}.csv"
html_path = f"{working_dir}/Preflight/HTML/{name}.html"
await asyncio.to_thread(df.to_csv, csv_path, index=False)
await asyncio.to_thread(formatHTML, df, html_path)
print("buildPreflights completed asynchronously.")
async def splitFilepathsGrouped(df, col="filename"):
path_task = asyncio.create_task(get_protected_value("PATH_EXCLUSION_CONST", int))
min_files_task = asyncio.create_task(get_protected_value("MIN_FILES_FOR_PATH", int))
path_exclusion_constant = await path_task
min_files_for_path = await min_files_task
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
return []
parts = os.path.normpath(path).split(os.sep)
parts = [p for p in parts if p] # Remove empty strings
return parts
return [p for p in parts if p]
# 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)
logger.warning(f"Non-string entries found in column '{col}':")
logger.debug(non_string_entries)
df = df.copy()
split_paths = df[col].apply(clean_split)
# Filter out paths with fewer than `min_files_for_path` components
df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
split_paths = split_paths[df.index]
@@ -295,11 +304,7 @@ def splitFilepathsGrouped(df, col="filename"):
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 ""
)
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
@@ -309,19 +314,16 @@ def splitFilepathsGrouped(df, col="filename"):
return pd.DataFrame(new_rows).drop(columns=["group_key"])
def calculatePath(approved_hashes, 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)
async def calculatePath(approved_hashes, split):
dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] if split else [approved_hashes]
badpathparts = await asyncio.to_thread(load_env_json, "BAD_PATH_PARTS", "[]")
min_files_for_path = await asyncio.to_thread(get_protected_value, "MIN_FILES_FOR_PATH", int)
processed_dfs = []
for df in dfs_by_policy:
haslcp = splitFilepathsGrouped(df, "filename_exec")
haslcp = await splitFilepathsGrouped(df, "filename_exec")
haslcp = haslcp.drop_duplicates()
forbidden = regulator(badpathparts, True)
@@ -329,18 +331,12 @@ def calculatePath(approved_hashes, split):
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",
]
]
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()
@@ -353,8 +349,8 @@ def calculatePath(approved_hashes, split):
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
return pathExclusions