First Round Async. Much work left to do, dont trust results of hash categorization presently.
This commit is contained in:
+75
-169
@@ -1,29 +1,15 @@
|
||||
# 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 asyncio
|
||||
import datetime
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
|
||||
import aiofiles
|
||||
import pandas as pd
|
||||
import tqdm
|
||||
from bson import ObjectId
|
||||
from tqdm.asyncio import tqdm_asyncio
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
@@ -35,181 +21,101 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
def pullPolicyExechistories(
|
||||
api: AirlockAPIWrapper,
|
||||
policy: Policy,
|
||||
type: list,
|
||||
days,
|
||||
outputjson: bool,
|
||||
):
|
||||
|
||||
file_path = f"{get_base_directory()}\\cache\\chunkinator.json"
|
||||
|
||||
# Ensure the file exists
|
||||
if not os.path.exists(file_path):
|
||||
with open(file_path, "w") as file:
|
||||
json.dump({"error": "Success", "response": {"exechistories": []}}, file)
|
||||
logger.debug(f"File '{file_path}' has been created.")
|
||||
else:
|
||||
logger.debug(f"File '{file_path}' already exists.")
|
||||
|
||||
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": []}}
|
||||
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=10000,
|
||||
desc=f"Checkpoint Progress: {checkpoint}",
|
||||
colour="blue",
|
||||
initial=1,
|
||||
) as filebar:
|
||||
with tqdm.tqdm(
|
||||
file=sys.stdout,
|
||||
leave=True,
|
||||
total=100,
|
||||
desc=f"Total of {policy} Complete: ",
|
||||
) as pbar:
|
||||
while True:
|
||||
histories = api.history_logging(
|
||||
type=type, checkpoint=checkpoint, policy= [policy.name]
|
||||
)
|
||||
if not os.path.exists(file_path):
|
||||
async with aiofiles.open(file_path, "w") as file:
|
||||
await file.write(json.dumps(json_output))
|
||||
|
||||
# Ensure histories is a list of dictionaries
|
||||
if not isinstance(histories, list) or not all(
|
||||
isinstance(h, dict) for h in histories
|
||||
):
|
||||
logger.error(
|
||||
"Unexpected response format from API. Expected list of dictionaries."
|
||||
)
|
||||
break
|
||||
filebar = tqdm_asyncio(total=10000, desc=f"Checkpoint Progress: {checkpoint}", colour="blue")
|
||||
pbar = tqdm_asyncio(total=100, desc=f"Total of {policy.name} Complete: ")
|
||||
|
||||
filebar.total = len(histories)
|
||||
while True:
|
||||
histories = await api.history_logging(type=type, checkpoint=checkpoint, policy=[policy.name])
|
||||
if not histories:
|
||||
break
|
||||
|
||||
if not histories:
|
||||
break
|
||||
for index, history_item in enumerate(histories):
|
||||
if "checkpoint" not in history_item or "datetime" not in history_item:
|
||||
continue
|
||||
|
||||
for index, history_item in enumerate(histories):
|
||||
if (
|
||||
"checkpoint" not in history_item
|
||||
or "datetime" not in history_item
|
||||
):
|
||||
continue # Skip malformed entries
|
||||
if index == len(histories) - 1:
|
||||
checkpoint = history_item["checkpoint"]
|
||||
filebar.set_description(f"Checkpoint Progress: {checkpoint}")
|
||||
break
|
||||
|
||||
# Update checkpoint on last item
|
||||
if index == len(histories) - 1:
|
||||
checkpoint = history_item["checkpoint"] # pyright: ignore[reportArgumentType]
|
||||
filebar.desc = 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
|
||||
|
||||
try:
|
||||
history_date = datetime.datetime.strptime(
|
||||
history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportArgumentType]
|
||||
"%Y-%m-%dT%H:%M:%SZ",
|
||||
).date()
|
||||
except ValueError:
|
||||
continue # Skip if date format is invalid
|
||||
if datetime.date.today() - datetime.timedelta(days=days) <= history_date:
|
||||
json_output["response"]["exechistories"].append(history_item)
|
||||
|
||||
if (
|
||||
datetime.date.today() - datetime.timedelta(days=days)
|
||||
) <= history_date:
|
||||
json_output["response"]["exechistories"].append(history_item)
|
||||
filebar.update(1)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
filebar.update(1)
|
||||
filebar.refresh()
|
||||
# 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"]
|
||||
|
||||
# Deduplicate entries
|
||||
seen = {}
|
||||
if os.path.exists(file_path):
|
||||
with open(file_path, "r") as file:
|
||||
existing_data = json.load(file)
|
||||
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
|
||||
|
||||
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}}))
|
||||
|
||||
deduplicated = list(seen.values())
|
||||
with open(file_path, "w") as file:
|
||||
json.dump(
|
||||
{
|
||||
"error": "Success",
|
||||
"response": {"exechistories": deduplicated},
|
||||
},
|
||||
file,
|
||||
)
|
||||
json_output["response"]["exechistories"].clear()
|
||||
|
||||
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
|
||||
|
||||
# Update progress bar based on last valid item
|
||||
try:
|
||||
last_date = datetime.datetime.strptime(
|
||||
history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore
|
||||
"%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_str(f"Total of {policy} Complete: ")
|
||||
pbar.refresh()
|
||||
except Exception:
|
||||
pass
|
||||
filebar.n = 1
|
||||
|
||||
filebar.n = 1
|
||||
|
||||
# Final output
|
||||
with open(file_path, "r") as file:
|
||||
final_output = json.load(file)
|
||||
async with aiofiles.open(file_path, "r") as file:
|
||||
final_output = await file.read()
|
||||
os.remove(file_path)
|
||||
|
||||
return json.dumps(final_output) if outputjson else None
|
||||
return final_output if outputjson else None
|
||||
|
||||
|
||||
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
|
||||
async def getPolicyInfo(api, policy, type, days):
|
||||
executionhist_policy = pd.DataFrame()
|
||||
exehist = pullPolicyExechistories(api, policy, type, days, True)
|
||||
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",
|
||||
]
|
||||
["datetime", "sha256", "publisher", "filename", "hostname", "username", "pprocess", "gprocess", "commandline"]
|
||||
]
|
||||
executionhist_policy["policy"] = policy # Add policy column here
|
||||
executionhist_policy = executionhist_policy.drop_duplicates(
|
||||
subset=["sha256", "filename", "hostname"]
|
||||
)
|
||||
executionhist_policy = executionhist_policy.sort_values(
|
||||
by=["sha256", "filename"]
|
||||
)
|
||||
logger.debug( f"Staging of Execution history for policy: {policy} is complete")
|
||||
print(
|
||||
colorText(
|
||||
f"Staging of Execution history for policy: {policy} is complete",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
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()
|
||||
@@ -230,8 +136,8 @@ def skipback(days):
|
||||
return ObjectId(objectid_hex)
|
||||
|
||||
|
||||
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
||||
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
|
||||
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():
|
||||
api.policy_clone(enforcement_policy, audit_policy)
|
||||
api.policy_set_auditmode(audit_policy, "1")
|
||||
await api.policy_clone(enforcement_policy, audit_policy)
|
||||
await api.policy_set_auditmode(audit_policy, "1")
|
||||
Reference in New Issue
Block a user