Files
AirlockTools/services/policyhandler.py
T

238 lines
8.3 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
import pandas as pd
import tqdm
from bson import ObjectId
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__)
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 = 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")