246 lines
8.6 KiB
Python
246 lines
8.6 KiB
Python
# 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 datetime
|
|
import gc
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
from bson import ObjectId
|
|
import pandas as pd
|
|
import tqdm
|
|
|
|
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 areYouSure, colorText, get_sanitized_input
|
|
|
|
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.")
|
|
|
|
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]
|
|
)
|
|
|
|
# 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.total = len(histories)
|
|
|
|
if not histories:
|
|
break
|
|
|
|
for index, history_item in enumerate(histories):
|
|
if (
|
|
"checkpoint" not in history_item
|
|
or "datetime" not in history_item
|
|
):
|
|
continue # Skip malformed entries
|
|
|
|
# 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", ""), # 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)
|
|
|
|
filebar.update(1)
|
|
filebar.refresh()
|
|
|
|
# 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
|
|
|
|
deduplicated = list(seen.values())
|
|
with open(file_path, "w") as file:
|
|
json.dump(
|
|
{
|
|
"error": "Success",
|
|
"response": {"exechistories": deduplicated},
|
|
},
|
|
file,
|
|
)
|
|
|
|
json_output["response"]["exechistories"].clear()
|
|
|
|
# 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
|
|
|
|
# Final output
|
|
with open(file_path, "r") as file:
|
|
final_output = json.load(file)
|
|
os.remove(file_path)
|
|
|
|
return json.dumps(final_output) if outputjson else None
|
|
|
|
|
|
def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
|
|
import airlock_libs
|
|
|
|
executionhist_policy = pd.DataFrame()
|
|
exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
|
|
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 # 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",
|
|
)
|
|
)
|
|
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)
|
|
|
|
|
|
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
|
|
policy_relationship_map = 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")
|
|
|
|
|
|
def confirmUpdateAfromE(api: AirlockAPIWrapper):
|
|
areYouSure()
|
|
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
|
if confirmation.strip() == "I AGREE":
|
|
updateAuditPoliciesFromEnforcementPolices(api) |