Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance
- Migrated codebase to class-based architecture for better modularity and maintainability - Introduced system_config.json for centralized configuration (required for runtime) - Added structured working directories for improved file organization - Significantly reduced reliance on Parquet; replaced with alternative data handling - Implemented security improvements across modules - Several TODOs remain in the main script for future enhancements - Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
@@ -0,0 +1,64 @@
|
||||
# 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/>.
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import ClassVar, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class Agent:
|
||||
agentid: str
|
||||
clientversion: str
|
||||
domain: str
|
||||
freespace: int
|
||||
groupid: int
|
||||
hostname: str
|
||||
ip: str
|
||||
localip: str
|
||||
lastcheckin: str
|
||||
os: str
|
||||
policyversion: str
|
||||
status: int # raw status code
|
||||
username: str
|
||||
groupname: Optional[str] = field(default=None)
|
||||
status_text: Optional[str] = field(default=None)
|
||||
|
||||
# Class-level status map
|
||||
status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"}
|
||||
|
||||
def enrich(self, groupid_to_name: dict):
|
||||
"""Enrich the agent with groupname and human-readable status."""
|
||||
self.groupname = groupid_to_name.get(self.groupid, None)
|
||||
self.status_text = self.status_map.get(self.status, "Unknown")
|
||||
|
||||
|
||||
"""
|
||||
|
||||
from models.agent import Agent
|
||||
from modesls.policy
|
||||
|
||||
# Step 1: Load data from API
|
||||
policies = [Policy(**row['data']) for _, row in api.policy_find_all().iterrows()]
|
||||
agents = [Agent(**row['data']) for _, row in api.agent_find_all().iterrows()]
|
||||
|
||||
# Step 2: Create groupid → groupname map
|
||||
groupid_to_name = {policy.groupid: policy.name for policy in policies}
|
||||
|
||||
# Step 3: Enrich agents
|
||||
for agent in agents:
|
||||
agent.enrich(groupid_to_name)
|
||||
|
||||
|
||||
"""
|
||||
@@ -0,0 +1,389 @@
|
||||
# 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 inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import List, Optional
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
from services.policyhandler import pullPolicyExechistories
|
||||
from utils.utils import colorText, load_env, load_env_json, 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,
|
||||
):
|
||||
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)
|
||||
|
||||
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:
|
||||
if h.sha256 not in seen:
|
||||
seen.add(h.sha256)
|
||||
deduped.append(h)
|
||||
return deduped
|
||||
|
||||
@classmethod
|
||||
def categorize_hashes(cls, hashes):
|
||||
import re
|
||||
from utils.utils import load_env, load_env_json, regulator
|
||||
|
||||
threat_tolerance = load_env("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 = []
|
||||
|
||||
def reputationtool(hash_obj):
|
||||
val = hash_obj.reputation.get("scannermatch") if isinstance(hash_obj.reputation, dict) else None
|
||||
if val in [None, "N/A"]:
|
||||
return hash_obj.publisher == "Not Signed"
|
||||
try:
|
||||
return int(val) > threat_tolerance
|
||||
except (ValueError, TypeError):
|
||||
return hash_obj.publisher == "Not Signed"
|
||||
|
||||
for hash_obj in hashes:
|
||||
publisher = hash_obj.publisher or ""
|
||||
description = hash_obj.description or ""
|
||||
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
|
||||
rep_status = reputation.get("status")
|
||||
|
||||
rep_flag = reputationtool(hash_obj)
|
||||
|
||||
is_signed = publisher != "Not Signed"
|
||||
is_untrusted = re.search(bad_publishers_pattern, publisher, re.IGNORECASE) is not None
|
||||
is_pup = re.search(pups_pattern, description, re.IGNORECASE) is not None
|
||||
has_known_status = rep_status == "KNOWN"
|
||||
|
||||
if (not is_signed and rep_flag) or rep_status == "UNKNOWN":
|
||||
needs_review.append(hash_obj)
|
||||
elif (
|
||||
(is_signed and not is_untrusted and has_known_status and not is_pup) or
|
||||
(not is_signed and not rep_flag and not is_untrusted and has_known_status and not is_pup)
|
||||
):
|
||||
approved.append(hash_obj)
|
||||
else:
|
||||
unapproved.append(hash_obj)
|
||||
|
||||
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.
|
||||
"""
|
||||
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:
|
||||
if var_val is hash_list:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{var_name}_{timestamp}.csv"
|
||||
break
|
||||
else:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"hashes_export_{timestamp}.csv"
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
username: str
|
||||
hostname: str
|
||||
netdomain: str
|
||||
filename: str
|
||||
ppolicy: str
|
||||
policyname: str
|
||||
policyver: str
|
||||
commandline: str
|
||||
publisher: str
|
||||
sha256: str
|
||||
datetime: str
|
||||
|
||||
# Optional fields
|
||||
type: Optional[int] = None
|
||||
pprocess: Optional[str] = None
|
||||
gprocess: Optional[str] = None
|
||||
md5: Optional[str] = None
|
||||
sha128: Optional[str] = None
|
||||
sha384: Optional[str] = None
|
||||
sha512: Optional[str] = None
|
||||
ip: Optional[str] = None
|
||||
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
|
||||
|
||||
@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")
|
||||
)
|
||||
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}")
|
||||
|
||||
|
||||
@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
|
||||
]
|
||||
if missing_fields:
|
||||
raise ValueError(f"Missing mandatory fields: {missing_fields}")
|
||||
|
||||
return cls(
|
||||
username=data["username"],
|
||||
hostname=data["hostname"],
|
||||
netdomain=data["netdomain"],
|
||||
filename=data["filename"],
|
||||
ppolicy=data["ppolicy"],
|
||||
policyname=data["policyname"],
|
||||
policyver=data["policyver"],
|
||||
commandline=data["commandline"],
|
||||
publisher=data["publisher"],
|
||||
sha256=data["sha256"],
|
||||
datetime=data["datetime"],
|
||||
type=data.get("type"),
|
||||
pprocess=data.get("pprocess"),
|
||||
gprocess=data.get("gprocess"),
|
||||
md5=data.get("md5"),
|
||||
sha128=data.get("sha128"),
|
||||
sha384=data.get("sha384"),
|
||||
sha512=data.get("sha512"),
|
||||
ip=data.get("ip"),
|
||||
localip=data.get("localip"),
|
||||
extid=data.get("extid"),
|
||||
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")
|
||||
"""
|
||||
@@ -0,0 +1,65 @@
|
||||
# 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 json
|
||||
|
||||
"""
|
||||
Policy model representing policy data and relationships.
|
||||
"""
|
||||
|
||||
|
||||
class Policy:
|
||||
def __init__(self, groupid, hidden, name, parent):
|
||||
self.groupid = groupid
|
||||
self.hidden = hidden
|
||||
self.name = name
|
||||
self.parent = parent
|
||||
|
||||
def __repr__(self):
|
||||
# Show all current attributes, including dynamically added ones
|
||||
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
|
||||
return f"<Execution({attrs})>"
|
||||
|
||||
def to_dict(self):
|
||||
# Return all attributes as a dictionary
|
||||
return self.__dict__
|
||||
|
||||
def to_json(self):
|
||||
# Convert to JSON string, handling non-serializable types gracefully
|
||||
return json.dumps(self.to_dict(), default=str)
|
||||
|
||||
|
||||
class Allowlist:
|
||||
"""
|
||||
Represents Allowlist
|
||||
"""
|
||||
|
||||
def __init__(self, applicationid, name, version):
|
||||
self.applicationid = applicationid
|
||||
self.name = name
|
||||
self.version = version
|
||||
|
||||
def __repr__(self):
|
||||
# Show all current attributes, including dynamically added ones
|
||||
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
|
||||
return f"<Execution({attrs})>"
|
||||
|
||||
def to_dict(self):
|
||||
# Return all attributes as a dictionary
|
||||
return self.__dict__
|
||||
|
||||
def to_json(self):
|
||||
# Convert to JSON string, handling non-serializable types gracefully
|
||||
return json.dumps(self.to_dict(), default=str)
|
||||
Reference in New Issue
Block a user