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

Bugfix, plus some QoL upgrades from the Async branch

Browse Source
This commit is contained in:
Zarithas
2025-10-20 14:25:03 -04:00
parent aa8380ac0a
commit f1f33948e3
7 changed files with 299 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/execution.py
+225 -153
View File
@@ -13,18 +13,20 @@
# 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 dataclasses
import inspect
import json
import logging
import os
import re
from dataclasses import dataclass
from dataclasses import asdict, dataclass
from datetime import datetime
from typing import List, Optional
from typing import List, Optional, Tuple
import dotenv
import pandas as pd
from services.API import AirlockAPIWrapper
from services.policyhandler import pullPolicyExechistories
from utils.configmanager import get_protected_value, load_env_json
from utils.utils import colorText, regulator
@@ -33,80 +35,39 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv()
@dataclass
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,
):
self.sha256 = sha256
self.applications = applications
self.baselines = baselines
self.blocklists = blocklists
self.createtime = createtime
self.datetime = datetime
self.description = description
self.filename = filename
self.filepath = filepath
self.filesize = filesize
self.md5 = md5
self.modtime = modtime
self.origname = origname
self.productname = productname
self.productversion = productversion
self.publisher = publisher
self.reputation = reputation
self.sha128 = sha128
self.sha384 = sha384
self.sha512 = sha512
def __repr__(self):
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
return f"<Hash({attrs})>"
def __eq__(self, other):
if isinstance(other, Hash):
return self.sha256 == other.sha256
return False
def __hash__(self):
return hash(self.sha256)
sha256: str
applications: str
baselines: str
blocklists: str
createtime: str
datetime: str
description: str
filename: str
filepath: str
filesize: str
md5: str
modtime: str
origname: str
productname: str
productversion: str
publisher: str
reputation: str
sha128: str
sha384: str
sha512: str
at_decision: Optional[str] = None
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
return asdict(self)
@classmethod
def from_dict(cls, data: dict):
return cls(**data)
@classmethod
def deduplicate(cls, hash_list):
@@ -131,9 +92,9 @@ class Hash:
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
pups_pattern = regulator(load_env_json("PUPS", "[]"))
needs_review = []
approved = []
unapproved = []
approved_count = 0
unapproved_count = 0
needs_review_count = 0
for hash_obj in hashes:
publisher = hash_obj.publisher or ""
@@ -147,18 +108,21 @@ class Hash:
# 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)
hash_obj.at_decision = "unapproved"
unapproved_count += 1
continue
if re.search(pups_pattern, description, re.IGNORECASE):
logger.debug("Unapproved: Description matches PUP pattern.")
unapproved.append(hash_obj)
hash_obj.at_decision = "unapproved"
unapproved_count += 1
continue
# 2. Approved: signed
if publisher != "Not Signed":
logger.debug("Approved: File is signed and not flagged.")
approved.append(hash_obj)
hash_obj.at_decision = "approved"
approved_count += 1
continue
# 3. Approved or Unapproved based on threat level
@@ -167,17 +131,20 @@ class Hash:
logger.debug(f"Parsed scannermatch score: {score}")
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
logger.debug("Unapproved: Unsigned file with high threat score.")
unapproved.append(hash_obj)
hash_obj.at_decision = "unapproved"
unapproved_count += 1
else:
logger.debug("Approved: Unsigned file with low threat score.")
approved.append(hash_obj)
hash_obj.at_decision = "approved"
approved_count += 1
except (ValueError, TypeError):
logger.debug("Needs Review: Scannermatch score is missing or invalid.")
needs_review.append(hash_obj)
logger.debug("Needs Review: Scannermatch score is missing or invalid. — {e}")
hash_obj.at_decision = "needs_review"
needs_review_count += 1
logger.debug(f"Final counts — Needs Review: {len(needs_review)}, Approved: {len(approved)}, Unapproved: {len(unapproved)}")
return needs_review, approved, unapproved
logger.debug(f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}")
return hashes
@classmethod
@@ -209,43 +176,6 @@ class Hash:
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)
"""
@dataclass
class ExecutionHistoryRecord:
# Mandatory fields
@@ -275,37 +205,7 @@ class ExecutionHistoryRecord:
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
@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])
logger.debug(f"Execution DataFrame columns: {exec_df.columns}")
logger.debug(f"Hash DataFrame columns: {hash_df.columns}")
if hash_df.empty:
logger.warning(f"hash_df is empty for label: {label}. Skipping merge.")
merged_df = exec_df.copy()
else:
merged_df = pd.merge(
exec_df,
hash_df,
on="sha256",
how="left", # Preserve all executions, enrich where possible
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}")
hash_obj: Optional[Hash] = None
@classmethod
@@ -354,6 +254,7 @@ class ExecutionHistoryRecord:
extname=data.get("extname"),
exttype=data.get("exttype"),
extbrowser=data.get("extbrowser"),
hash_obj=data.get("hash_obj")
)
@classmethod
@@ -385,10 +286,181 @@ class ExecutionHistoryRecord:
)
return executions
@staticmethod
def enrich_with_hashes(
api: AirlockAPIWrapper,
executions: List["ExecutionHistoryRecord"]
) -> List["ExecutionHistoryRecord"]:
"""
Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API.
"""
sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256})
logger.info(f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records.")
if not sha256_list:
logger.warning("No sha256 values found in execution records. Skipping enrichment.")
return executions
logger.debug("Querying hash data from API...")
hash_df = api.hash_query(sha256_list)
logger.info(f"Retrieved {len(hash_df)} hash records from API.")
hash_objects = []
required_fields = {
f.name for f in dataclasses.fields(Hash)
if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING
}
for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()):
row_dict = row.to_dict()
# Unwrap nested 'data' field if present
if "data" in row_dict and isinstance(row_dict["data"], dict):
row_dict = row_dict["data"]
# Inject the sha256 back into the row
row_dict["sha256"] = sha256
missing = required_fields - row_dict.keys()
if missing:
logger.warning(f"Skipping hash row due to missing fields: {missing}")
logger.debug(f"Row content: {row_dict}")
continue
try:
hash_obj = Hash.from_dict(row_dict)
hash_objects.append(hash_obj)
except Exception as e:
logger.warning(f"Failed to create Hash from row: {e}")
logger.debug(f"Row content: {row_dict}")
logger.debug("Converted hash DataFrame to Hash objects.")
hash_lookup = {h.sha256.strip().lower(): h for h in hash_objects}
logger.debug("Built hash lookup table.")
enriched_count = 0
for exec_record in executions:
hash_obj = hash_lookup.get(exec_record.sha256.strip().lower())
if hash_obj:
exec_record.hash_obj = hash_obj
enriched_count += 1
logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.")
return executions
@staticmethod
def categorize_executions_by_hash_decision(executions: List["ExecutionHistoryRecord"]) -> List["ExecutionHistoryRecord"]:
"""
Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation.
Modifies the `at_decision` field of each associated Hash object in-place.
Returns:
List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated.
"""
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", "[]"))
approved_count = 0
unapproved_count = 0
needs_review_count = 0
for record in executions:
hash_obj = record.hash_obj
if not hash_obj:
continue # Skip if no hash object is attached
publisher = hash_obj.publisher or ""
description = hash_obj.description or ""
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.")
hash_obj.at_decision = "unapproved"
unapproved_count += 1
continue
if re.search(pups_pattern, description, re.IGNORECASE):
logger.debug("Unapproved: Description matches PUP pattern.")
hash_obj.at_decision = "unapproved"
unapproved_count += 1
continue
# 2. Approved: signed
if publisher != "Not Signed":
logger.debug("Approved: File is signed and not flagged.")
hash_obj.at_decision = "approved"
approved_count += 1
continue
# 3. Approved or Unapproved based on threat level
try:
score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}")
if threat_tolerance is not None and score >= threat_tolerance:
logger.debug("Unapproved: Unsigned file with high threat score.")
hash_obj.at_decision = "unapproved"
unapproved_count += 1
else:
logger.debug("Approved: Unsigned file with low threat score.")
hash_obj.at_decision = "approved"
approved_count += 1
except (ValueError, TypeError) as e:
logger.debug(f"Needs Review: Scannermatch score is missing or invalid. — {e}")
hash_obj.at_decision = "needs_review"
needs_review_count += 1
logger.debug(
f"Final counts — Needs Review: {needs_review_count}, "
f"Approved: {approved_count}, Unapproved: {unapproved_count}"
)
return executions
@classmethod
def sort_by_hash_decision(
cls, executions: List["ExecutionHistoryRecord"]
) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]:
"""
Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups
based on the value of hash_obj.at_decision.
Returns:
Tuple of lists: (approved, unapproved, needs_review, unknown)
"""
approved = []
unapproved = []
needs_review = []
unknown = []
for record in executions:
decision = getattr(record.hash_obj, "at_decision", None)
if decision == "approved":
approved.append(record)
elif decision == "unapproved":
unapproved.append(record)
elif decision == "needs_review":
needs_review.append(record)
else:
unknown.append(record)
logger.info(f"[ExecutionHistoryRecord] Sorted {len(executions)} records by hash_obj.at_decision:")
logger.info(f" Approved: {len(approved)}")
logger.info(f" Unapproved: {len(unapproved)}")
logger.info(f" Needs Review: {len(needs_review)}")
logger.info(f" Unknown/Unset: {len(unknown)}")
return approved, unapproved, needs_review, unknown
def __repr__(self):
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
return f"<Execution({attrs})>"
"""