143 lines
5.3 KiB
Python
143 lines
5.3 KiB
Python
import asyncio
|
|
import datetime
|
|
import gc
|
|
import json
|
|
import logging
|
|
import os
|
|
import uuid
|
|
|
|
import aiofiles
|
|
import pandas as pd
|
|
from bson import ObjectId
|
|
from tqdm.asyncio import tqdm_asyncio
|
|
|
|
from models.policy import Policy
|
|
from services.API import AirlockAPIWrapper
|
|
from utils.configmanager import get_protected_json
|
|
from utils.setup import get_base_directory
|
|
from utils.utils import colorText
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
async def pullPolicyExechistories(api: AirlockAPIWrapper, policy: Policy, type, days, outputjson):
|
|
file_path = f"{get_base_directory()}\\cache\\chunkinator_{policy.name}_{uuid.uuid4().hex}.json"
|
|
checkpoint = str(skipback(days))
|
|
json_output = {"error": "Success", "response": {"exechistories": []}}
|
|
|
|
if not os.path.exists(file_path):
|
|
async with aiofiles.open(file_path, "w") as file:
|
|
await file.write(json.dumps(json_output))
|
|
|
|
filebar = tqdm_asyncio(total=10000, desc=f"Checkpoint Progress: {checkpoint}", colour="blue")
|
|
pbar = tqdm_asyncio(total=100, desc=f"Total of {policy.name} Complete: ")
|
|
|
|
while True:
|
|
histories = await api.history_logging(type=type, checkpoint=checkpoint, policy=[policy.name])
|
|
if not histories:
|
|
break
|
|
|
|
for index, history_item in enumerate(histories):
|
|
if "checkpoint" not in history_item or "datetime" not in history_item:
|
|
continue
|
|
|
|
if index == len(histories) - 1:
|
|
checkpoint = history_item["checkpoint"]
|
|
filebar.set_description(f"Checkpoint Progress: {checkpoint}")
|
|
break
|
|
|
|
try:
|
|
history_date = datetime.datetime.strptime(
|
|
history_item["datetime"].replace(" +0000 UTC", ""),
|
|
"%Y-%m-%dT%H:%M:%SZ"
|
|
).date()
|
|
except ValueError:
|
|
continue
|
|
|
|
if datetime.date.today() - datetime.timedelta(days=days) <= history_date:
|
|
json_output["response"]["exechistories"].append(history_item)
|
|
|
|
filebar.update(1)
|
|
await asyncio.sleep(0)
|
|
|
|
# Deduplication
|
|
seen = {}
|
|
if os.path.exists(file_path):
|
|
async with aiofiles.open(file_path, "r") as file:
|
|
content = await file.read()
|
|
existing_data = json.loads(content)
|
|
combined = existing_data["response"]["exechistories"] + json_output["response"]["exechistories"]
|
|
else:
|
|
combined = json_output["response"]["exechistories"]
|
|
|
|
for entry in combined:
|
|
key = (entry.get("sha256"), entry.get("filename"), entry.get("hostname"))
|
|
seen[key] = entry
|
|
|
|
deduplicated = list(seen.values())
|
|
async with aiofiles.open(file_path, "w") as file:
|
|
await file.write(json.dumps({"error": "Success", "response": {"exechistories": deduplicated}}))
|
|
|
|
json_output["response"]["exechistories"].clear()
|
|
|
|
try:
|
|
last_date = datetime.datetime.strptime(
|
|
history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportPossiblyUnboundVariable]
|
|
"%Y-%m-%dT%H:%M:%SZ"
|
|
).date()
|
|
date_diff = datetime.date.today() - last_date
|
|
percentage_diff = (((days + 10) - date_diff.days) / (days + 10)) * 100
|
|
pbar.n = round(percentage_diff)
|
|
pbar.set_description(f"Total of {policy.name} Complete: ")
|
|
except Exception:
|
|
pass
|
|
|
|
filebar.n = 1
|
|
|
|
async with aiofiles.open(file_path, "r") as file:
|
|
final_output = await file.read()
|
|
os.remove(file_path)
|
|
|
|
return final_output if outputjson else None
|
|
|
|
|
|
async def getPolicyInfo(api, policy, type, days):
|
|
executionhist_policy = pd.DataFrame()
|
|
exehist = await pullPolicyExechistories(api, policy, type, days, True)
|
|
if exehist is not None:
|
|
data = json.loads(exehist)
|
|
executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
|
|
if not executionhist_policy.empty:
|
|
executionhist_policy = executionhist_policy[
|
|
["datetime", "sha256", "publisher", "filename", "hostname", "username", "pprocess", "gprocess", "commandline"]
|
|
]
|
|
executionhist_policy["policy"] = policy.name
|
|
executionhist_policy = executionhist_policy.drop_duplicates(subset=["sha256", "filename", "hostname"])
|
|
executionhist_policy = executionhist_policy.sort_values(by=["sha256", "filename"])
|
|
print(colorText(f"Staging of Execution history for policy: {policy.name} is complete", "green"))
|
|
del data
|
|
del exehist
|
|
gc.collect()
|
|
return executionhist_policy
|
|
|
|
|
|
def skipback(days):
|
|
"""
|
|
Generate a MongoDB ObjectId for a given number of days ago from today.
|
|
"""
|
|
adjusted_days = days
|
|
date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
|
days=adjusted_days
|
|
)
|
|
timestamp = int(date_days_ago.timestamp())
|
|
hex_timestamp = format(timestamp, "08x")
|
|
objectid_hex = hex_timestamp + "0000000000000000"
|
|
return ObjectId(objectid_hex)
|
|
|
|
|
|
async def updateAuditPoliciesFromEnforcementPolices(api):
|
|
policy_relationship_map = await get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
|
for enforcement_policy, audit_policy in policy_relationship_map.items():
|
|
await api.policy_clone(enforcement_policy, audit_policy)
|
|
await api.policy_set_auditmode(audit_policy, "1") |