import asyncio import inspect import json import logging import os import re from dataclasses import dataclass from datetime import datetime from typing import List, Optional import aiofiles import pandas as pd from services.PolicyHandler import pullPolicyExechistories from utils.configmanager import get_protected_value, load_env_json from utils.utils import regulator logger = logging.getLogger(__name__) class Hash: 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"" def __eq__(self, other): return isinstance(other, Hash) and self.sha256 == other.sha256 def __hash__(self): return hash(self.sha256) def to_dict(self): return self.__dict__ @staticmethod def safe_int(value, default=0): try: return int(value) except (TypeError, ValueError): return default @classmethod def deduplicate(cls, hash_list): seen = set() deduped = [] for h in hash_list: if h.sha256 not in seen: seen.add(h.sha256) deduped.append(h) return deduped @classmethod 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 "" description = hash_obj.description or "" reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} scannermatch = reputation.get("scannermatch") if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): unapproved.append(hash_obj) continue if re.search(pups_pattern, description, re.IGNORECASE): unapproved.append(hash_obj) continue if publisher != "Not Signed": approved.append(hash_obj) continue try: score = int(scannermatch) # pyright: ignore[reportArgumentType] if score > threat_tolerance: # type: ignore unapproved.append(hash_obj) else: approved.append(hash_obj) except (ValueError, TypeError): needs_review.append(hash_obj) return needs_review, approved, unapproved @classmethod 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: 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" 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 = 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) 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: username: str hostname: str netdomain: str filename: str ppolicy: str policyname: str policyver: str commandline: str publisher: str sha256: str datetime: str 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 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 async def enrich_with_hashes(executions: list, hashes: list): 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: merged_df = exec_df.copy() logger.debug("Hash dataframe appears empty") else: merged_df = await asyncio.to_thread( pd.merge, exec_df, hash_df, on="sha256", how="left", suffixes=("_exec", "_hash") ) # 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] 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"), )