mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

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

Browse Source
This commit is contained in:
Zarithas
2025-10-16 16:43:56 -04:00
parent fa0c18ee02
commit 6f2355fea9
21 changed files with 903 additions and 1647 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/execution.py
+65 -217
View File
@@ -1,18 +1,4 @@
# 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 asyncio
import inspect
import json
import logging
@@ -22,46 +8,20 @@ from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional
import dotenv
import aiofiles
import pandas as pd
from services.policyhandler import pullPolicyExechistories
from services.PolicyHandler import pullPolicyExechistories
from utils.configmanager import get_protected_value, load_env_json
from utils.utils import colorText, regulator
from utils.utils import regulator
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
class Hash:
"""
Hash model representing Hash data
"""
def __init__(
self,
sha256,
applications=None,
baselines=None,
blocklists=None,
createtime=None,
datetime=None,
description=None,
filename=None,
filepath=None,
filesize=None,
md5=None,
modtime=None,
origname=None,
productname=None,
productversion=None,
publisher=None,
reputation=None,
sha128=None,
sha384=None,
sha512=None,
):
def __init__(self, sha256, applications=None, baselines=None, blocklists=None, createtime=None,
datetime=None, description=None, filename=None, filepath=None, filesize=None,
md5=None, modtime=None, origname=None, productname=None, productversion=None,
publisher=None, reputation=None, sha128=None, sha384=None, sha512=None):
self.sha256 = sha256
self.applications = applications
self.baselines = baselines
@@ -88,35 +48,23 @@ class Hash:
return f"<Hash({attrs})>"
def __eq__(self, other):
if isinstance(other, Hash):
return self.sha256 == other.sha256
return False
return isinstance(other, Hash) and self.sha256 == other.sha256
def __hash__(self):
return hash(self.sha256)
def to_dict(self):
"""Returns a dictionary representation of the hash."""
return self.__dict__
@staticmethod
def safe_int(value, default=0):
"""Safely convert a value to int, returning default on failure."""
try:
return int(value)
except (TypeError, ValueError):
return default
@classmethod
def deduplicate(cls, hash_list):
"""
Deduplicates a list of Hash objects based on sha256.
Args:
hash_list (list): List of Hash instances.
Returns:
list: Deduplicated list of Hash instances.
"""
seen = set()
deduped = []
for h in hash_list:
@@ -126,14 +74,11 @@ class Hash:
return deduped
@classmethod
def categorize_hashes(cls, hashes):
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
pups_pattern = regulator(load_env_json("PUPS", "[]"))
needs_review = []
approved = []
unapproved = []
async def categorize_hashes(cls, hashes):
threat_tolerance = await get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
bad_publishers_pattern = regulator(await load_env_json("BAD_PUBLISHERS", "[]"))
pups_pattern = regulator(await load_env_json("PUPS", "[]"))
needs_review, approved, unapproved = [], [], []
for hash_obj in hashes:
publisher = hash_obj.publisher or ""
@@ -141,57 +86,32 @@ class Hash:
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
scannermatch = reputation.get("scannermatch")
logger.debug(f"Evaluating hash: {hash_obj}")
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}")
# 1. Unapproved: bad publisher or PUP
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
logger.debug("Unapproved: Publisher matches bad publisher pattern.")
unapproved.append(hash_obj)
continue
if re.search(pups_pattern, description, re.IGNORECASE):
logger.debug("Unapproved: Description matches PUP pattern.")
unapproved.append(hash_obj)
continue
# 2. Approved: signed
if publisher != "Not Signed":
logger.debug("Approved: File is signed and not flagged.")
approved.append(hash_obj)
continue
# 3. Approved or Unapproved based on threat level
try:
score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}")
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
logger.debug("Unapproved: Unsigned file with high threat score.")
if score > threat_tolerance: # type: ignore
unapproved.append(hash_obj)
else:
logger.debug("Approved: Unsigned file with low threat score.")
approved.append(hash_obj)
except (ValueError, TypeError):
logger.debug("Needs Review: Scannermatch score is missing or invalid.")
needs_review.append(hash_obj)
logger.debug(f"Final counts — Needs Review: {len(needs_review)}, Approved: {len(approved)}, Unapproved: {len(unapproved)}")
return needs_review, approved, unapproved
@classmethod
def export_to_csv(cls, hash_list, directory_path):
"""
Exports a list of Hash objects to a CSV file in the specified directory.
The filename is derived from the variable name of the list if possible,
and includes a timestamp to ensure uniqueness.
"""
async def export_to_csv(cls, hash_list, directory_path):
filename = "hashes_export.csv"
frame = inspect.currentframe()
if frame is not None and frame.f_back is not None:
callers_local_vars = frame.f_back.f_locals.items()
for var_name, var_val in callers_local_vars:
for var_name, var_val in frame.f_back.f_locals.items():
if var_val is hash_list:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"{var_name}_{timestamp}.csv"
@@ -203,52 +123,16 @@ class Hash:
os.makedirs(directory_path, exist_ok=True)
file_path = os.path.join(directory_path, filename)
df = pd.DataFrame([h.to_dict() for h in hash_list])
df.to_csv(file_path, index=False)
df = await asyncio.to_thread(pd.DataFrame, [h.to_dict() for h in hash_list])
await asyncio.to_thread(df.to_csv, file_path, index=False)
logger.info(f"CSV file saved to: {file_path}")
"""
#Example - Convert Dataframe returned by hash query into hash objects
hash_objects = []
for _, row in df.iterrows():
try:
parsed_data = ast.literal_eval(row['data'])
hash_obj = Hash(sha256=row['sha256'], **parsed_data)
hash_objects.append(hash_obj)
except Exception as e:
print(f"Error parsing row: {e}")
# Display the created Hash objects
for obj in hash_objects:
print(obj)
# Categorize hashes
needs_review, approved, unapproved = Hash.categorize_hashes(
hashes=hash_objects,
threat_tolerance=3,
untrusted_pattern=untrusted_pattern,
pups_pattern=pups_pattern
)
# Deduplicate
deduped_hashes = Hash.deduplicate(hash_list)
# Specify the directory where you want to save the CSV
output_directory = "C:/Users/Brandon/Documents/HashExports"
# Call the export method
Hash.export_to_csv(hashes_for_export, output_directory)
"""
async with aiofiles.open(file_path, mode='r') as f:
preview = await f.read()
print(f"CSV file saved to: {file_path}\nPreview:\n{preview[:500]}")
@dataclass
class ExecutionHistoryRecord:
# Mandatory fields
username: str
hostname: str
netdomain: str
@@ -260,8 +144,6 @@ class ExecutionHistoryRecord:
publisher: str
sha256: str
datetime: str
# Optional fields
type: Optional[int] = None
pprocess: Optional[str] = None
gprocess: Optional[str] = None
@@ -273,62 +155,69 @@ class ExecutionHistoryRecord:
localip: Optional[str] = None
extid: Optional[str] = None
extname: Optional[str] = None
exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge
exttype: Optional[int] = None
extbrowser: Optional[int] = None
@classmethod
async def from_policies(cls, api, selected_policies, type_: list, history_days: int) -> List["ExecutionHistoryRecord"]:
async def fetch_and_parse(policy):
execs = await pullPolicyExechistories(api, policy, type_, history_days, True)
if not execs:
return []
data = json.loads(execs)
exechistories = data.get("response", {}).get("exechistories", [])
if not exechistories:
return []
df = await asyncio.to_thread(pd.DataFrame, exechistories)
df = await asyncio.to_thread(df.drop_duplicates, subset=["sha256", "filename", "hostname"])
df = await asyncio.to_thread(df.sort_values, by=["sha256", "filename"])
return [cls.from_dict(row.to_dict()) for _, row in df.iterrows()]
tasks = [fetch_and_parse(policy) for policy in selected_policies]
results = await asyncio.gather(*tasks)
return [record for sublist in results for record in sublist]
@staticmethod
def enrich_with_hashes_and_export(
executions: list, hashes: list, directory_path: str, label: str = "enriched"
):
exec_df = pd.DataFrame([e.__dict__ for e in executions])
hash_df = pd.DataFrame([h.to_dict() for h in hashes])
async def enrich_with_hashes(executions: list, hashes: list):
logger.debug(f"Execution DataFrame columns: {exec_df.columns}")
logger.debug(f"Hash DataFrame columns: {hash_df.columns}")
exec_task = asyncio.to_thread(pd.DataFrame, [e.__dict__ for e in executions])
hash_task = asyncio.to_thread(pd.DataFrame, [h.to_dict() for h in hashes])
exec_df, hash_df = await asyncio.gather(exec_task, hash_task)
if hash_df.empty:
logger.warning(f"hash_df is empty for label: {label}. Skipping merge.")
merged_df = exec_df.copy()
logger.debug("Hash dataframe appears empty")
else:
merged_df = pd.merge(
merged_df = await asyncio.to_thread(
pd.merge,
exec_df,
hash_df,
on="sha256",
how="left", # Preserve all executions, enrich where possible
how="left",
suffixes=("_exec", "_hash")
)
merged_df.sort_values(by="filename_exec", inplace=True)
logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}")
filename = f"{label}_executions.csv"
os.makedirs(directory_path, exist_ok=True)
file_path = os.path.join(directory_path, filename)
merged_df.to_csv(file_path, index=False)
logger.info(f"CSV file saved to: {file_path}")
# Log available columns for debugging
logger.debug(f"Merged DataFrame columns: {merged_df.columns.tolist()}")
# Only sort if the column exists
if "filename_exec" in merged_df.columns:
merged_df = await asyncio.to_thread(merged_df.sort_values, by="filename_exec")
else:
merged_df = await asyncio.to_thread(merged_df.sort_values, by="filename")
return merged_df
@classmethod
def from_dict(cls, data: dict):
mandatory_fields = [
"username",
"hostname",
"netdomain",
"filename",
"ppolicy",
"policyname",
"policyver",
"commandline",
"publisher",
"sha256",
"datetime",
]
missing_fields = [
field for field in mandatory_fields if field not in data or data[field] is None
"username", "hostname", "netdomain", "filename", "ppolicy",
"policyname", "policyver", "commandline", "publisher", "sha256", "datetime"
]
missing_fields = [field for field in mandatory_fields if field not in data or data[field] is None]
if missing_fields:
raise ValueError(f"Missing mandatory fields: {missing_fields}")
return cls(
username=data["username"],
hostname=data["hostname"],
@@ -354,45 +243,4 @@ class ExecutionHistoryRecord:
extname=data.get("extname"),
exttype=data.get("exttype"),
extbrowser=data.get("extbrowser"),
)
@classmethod
def from_policies(
cls, api, selected_policies, type_: list, history_days: int
) -> List["ExecutionHistoryRecord"]:
executions = []
for policy in selected_policies:
execs = pullPolicyExechistories(
api, policy, type_, history_days, True
)
if execs:
data = json.loads(execs)
exechistories = data.get("response", {}).get("exechistories", [])
if not exechistories:
continue
df = pd.DataFrame(exechistories)
df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
df = df.sort_values(by=["sha256", "filename"])
executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()])
logger.debug(f"Staging of Execution history for policy: {policy.name} is complete")
print(
colorText(
f"Staging of Execution history for policy: {policy.name} is complete",
"green",
)
)
return executions
def __repr__(self):
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
return f"<Execution({attrs})>"
"""
executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)
ExecutionHistoryRecord.enrich_with_hashes_and_export(executions, hash_objects, "C:/Users/Brandon/Documents/EnrichedExports")
"""
)