Files
AirlockTools/models/execution.py
T

246 lines
8.9 KiB
Python

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"<Hash({attrs})>"
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"),
)