Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance
- Migrated codebase to class-based architecture for better modularity and maintainability - Introduced system_config.json for centralized configuration (required for runtime) - Added structured working directories for improved file organization - Significantly reduced reliance on Parquet; replaced with alternative data handling - Implemented security improvements across modules - Several TODOs remain in the main script for future enhancements - Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
@@ -0,0 +1,236 @@
|
||||
# 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
|
||||
|
||||
import pandas as pd
|
||||
import tqdm
|
||||
from bson import ObjectId
|
||||
|
||||
from models.policy import Policy
|
||||
from services.API import AirlockAPIWrapper
|
||||
from utils.utils import colorText, load_env, load_env_json
|
||||
from services.setup import get_base_directory
|
||||
|
||||
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):
|
||||
executionhist_policy = pd.DataFrame()
|
||||
exehist = 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 # 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 = load_env_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")
|
||||
Reference in New Issue
Block a user