diff --git a/.gitignore b/.gitignore
index 37c7540..33dfcb9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,7 @@
chunkinator.json
jobs.json
*.xl*
-.exe
\ No newline at end of file
+.exe
+securitytest.py
+*.toml
+system_config.json
\ No newline at end of file
diff --git a/AirlockTools.py b/AirlockTools.py
index 958cec8..5f9416c 100644
--- a/AirlockTools.py
+++ b/AirlockTools.py
@@ -12,57 +12,80 @@
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
+
+
+#TODO Continue implementing logger
+#TODO Add input sanitation and CSV injection prevention
+#TODO Continue OTP and Local approval rewrites
+#TODO Explore pywin32
+#TODO Fix Requirements.txt
+#TODO Create Generic system_config.json for gitea
+
+
import argparse
-import dotenv
+import json
+import logging
import os
+
+import dotenv
import urllib3
-import utils.clientfunctions
-import utils.localapproval as la
-import utils.otpfunctions
-import utils.policyfunctions
-import utils.utils as ct
-import pandas as pd
+import utils.menus as menus
+from services.API import AirlockAPIWrapper
+from services.security import getAPI
+from services.setup import setup
-from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler
+urllib3.disable_warnings(
+ urllib3.exceptions.InsecureRequestWarning
+)
-urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-dotenv.load_dotenv()
-#Constants
-url = os.getenv('url')
-bad_publisher_list = ["Brave","Zoom", "GlavSoft", "VNC"]
-pups = ["logmein", "invalid", "nmap", "LTSvc", "VNC", "Kaseya", "Solarwinds", "mRemoteNG"]
-badpathparts = ["users", "wwwroot", "windows\\temp", "windows\\task", "windows\\system32", "startup", "windows\\fonts", "Recycle.Bin", "AppData", "programdata", "Solarwinds", "kaseya"]
-path_exclusion_constant = 4
-min_files_for_path = 4
-threat_tolerance_constant = 4
-
-policy_relationship_map ={ #Enforcement : Audit
-"bf0b1f9b-bfea-4f44-97c0-80e27ff61712" : "538d3218-92f4-4943-a6ee-db9267ab62d8", #AT Servers General, AT Servers General Audit
-"31ababac-65de-4c6a-86dd-6691d7e3ee3b" : "fc05b42a-b846-4e72-88ca-c35d416e699f", #AT Epic, #AT Epic Audit
-"d1f58960-f866-49e0-848a-a5b09fffd4cd" : "d55c03a6-c376-4391-8626-4f843b882a7c", #AT DMZ Enforced, #AT DMZ Audit
-"504dd011-86b6-489a-b78f-eff589cef8aa" : "88a1cfdc-3b30-448b-b309-be16fe437ca3", #AT Workstations BCA, #AT Workstations BCA Audit
-"d126db36-72ed-4937-adc7-d88b7509a5b5" : "5aebf6a0-1d67-47b4-9c5f-2866ffca5671" #AT Testing, AT Testing
-}
-
def main():
- parser = argparse.ArgumentParser(description="Your script description")
- parser.add_argument('--monitorOTP', action='store_true', help='Run in non-interactive mode')
+
+ #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
+
+ working_dir = setup()
+ logger = logging.getLogger(__name__)
+
+ dotenv.load_dotenv(dotenv_path=working_dir / ".env")
+
+ try:
+ url = os.getenv("URL")
+ username = os.getenv("USERNAME")
+
+ if not url:
+ raise ValueError("Missing URL in environment variables.")
+ if not username:
+ raise ValueError("Missing USERNAME in environment variables.")
+
+ logger.debug(f"Retrieved URL: {url}")
+ logger.debug(f"Retrieved Username: {username}")
+
+ except ValueError as e:
+ logger.error(f"Configuration error: {e}", exc_info=True)
+ raise
+
+
+ api = AirlockAPIWrapper(
+ base_url=str(os.getenv("URL")),
+ api_key = getAPI(username, "AirlockTools"),
+ )
+
+ parser = argparse.ArgumentParser(description="Program for Managing Airlock via API & CMD")
+ parser.add_argument("--monitor", action="store_true", help="Run in non-interactive mode")
# Add other arguments as needed
args = parser.parse_args()
- if args.monitorOTP:
+ if args.monitor:
# Non-interactive logic
- print(f"Running non-interactively to start monitoring OTP")
-
+ logger.info("Running non-interactively to start monitoring Airlock Changes")
+ """
os.makedirs("scheduling", exist_ok=True)
os.makedirs("OTP/HTML", exist_ok=True)
os.makedirs("OTP/PARQ", exist_ok=True)
os.makedirs("Local_Approval/HTML", exist_ok=True)
os.makedirs("Local_Approval/PARQ", exist_ok=True)
- ct.apivalidation()
-
+
register_function("monitorOTP", utils.otpfunctions.monitorOTP)
register_function("monitorLA", la.scheduleAddingLAHashes)
register_function("updateAuditPolicies", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices)
@@ -75,80 +98,27 @@ def main():
else:
reload_jobs()
start_scheduler()
-
+ """
else:
- # Interactive logic
- ct.apivalidation()
- menu_main()
-
-def menu_main():
- while True:
- ct.displayIntro();
- print(ct.colorText("1. 🖥️ - Get All Events for Single Device", "yellow"))
- print(ct.colorText("2. 🎫 - OTP", "yellow"))
- print(ct.colorText("3. 🔇 - Find Quiet Hosts", "yellow"))
- print(ct.colorText("4. 🔒 - Prepare Policy For Enforcement", "yellow"))
- print(ct.colorText("5. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
- print(ct.colorText("6. 🔍 - Device Search", "yellow"))
- print(ct.colorText("7. ➡️ - Move Devices to Local Approval", "yellow"))
- print(ct.colorText("Q. 🔚 - Quit", "yellow"))
-
- choice = input(ct.colorText("\nEnter Menu Item: ", "white"))
- if choice == '1':
- utils.clientfunctions.devicehistory(url,False)
- elif choice == "2":
- menu_otp()
- elif choice == "3":
- utils.clientfunctions.findQuietAgents(url)
- elif choice == "4":
- utils.policyfunctions.prepare_to_enforce(url, bad_publisher_list, pups, badpathparts, threat_tolerance_constant, path_exclusion_constant, min_files_for_path)
- elif choice == "5":
- ct.areYouSure()
- confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white"))
- if confirmation.strip().upper() == "I AGREE": utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices(url, policy_relationship_map)
- elif choice == "6":
- devicelist = utils.clientfunctions.promptForDevices()
- utils.clientfunctions.findAgents(url, devicelist, False)
-
- elif choice == "7":
- la.moveToLocalApproval(url, policy_relationship_map)
-
- elif choice == "8":
- las = la.getLocalApprovals(url)
- las.to_csv("la.csv", index=False)
-
- elif choice == "9":
- pass
- elif choice == "10":
- utils.clientfunctions.moveAgentToAudit(url,"6a221ece-0c10-4eb8-b1e5-06a1000a5696",policy_relationship_map)
- elif choice == "11":
- utils.clientfunctions.moveAgentToEnforcement(url,"6a221ece-0c10-4eb8-b1e5-06a1000a5696",policy_relationship_map)
- elif choice == "Q":
- break
- else:
- print(ct.colorText("Invalid choice. Please try again.","red"))
-
-def menu_otp():
- while True:
- print(ct.colorText("\n--- 🎫 OTP Submenu 🎫 ---","cyan"))
- print(ct.colorText("1. Generate OTP","cyan"))
- #print(ct.colorText("2. Sub-option B","cyan"))
- print(ct.colorText("Q. Return to Main Menu","cyan"))
- choice = input("Enter your choice: ")
-
- if choice == "1":
- utils.otpfunctions.generateOTP(url, utils.clientfunctions.findAgentID(url))
- break
-
- elif choice == "2":
- print("You selected Sub-option B")
- elif choice == "Q":
- print("Returning to Main Menu...")
- break
- else:
- print("Invalid choice. Please try again.")
-
+ # Interactive logic
+
+
+ raw = os.getenv("POLICY_MAP_ENF_AUD", "{}")
+
+ try:
+ # Escape backslashes before parsing
+ escaped = raw.encode('unicode_escape').decode('utf-8')
+ badpathparts = json.loads(escaped)
+ except Exception as e:
+ logging.error(f"Failed to parse BAD_PATH_PARTS: {e}")
+ badpathparts = []
+ print(badpathparts)
+
+
+
+ menus.menu_main(api)
+
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/README.md b/README.md
index 86edfdb..6207bdc 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,9 @@ Python based Carbon Black App Control feature implementation for Airlock
- "Local Approval Initialization"
This programmatically scans devices in audit mode within Airlock and subsequently adds the identified blocks to a user-specified whitelist.
+- KDE Wallet must be running on Linux (`kwalletd`)
+
+
## License
**AirlockTools** is licensed under the **GNU Affero General Public License v3.0**.
diff --git a/defalt_system_config.json b/defalt_system_config.json
new file mode 100644
index 0000000..b09f901
--- /dev/null
+++ b/defalt_system_config.json
@@ -0,0 +1,14 @@
+{
+ "APPNAME": "AirlockTools",
+ "URL": "https://server:3129",
+ "LOG_LEVEL": "INFO",
+ "BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"],
+ "BAD_PUBLISHERS": ["Brave", "Zoom", "GlavSoft", "VNC"],
+ "PUPS":["logmein","invalid","nmap","LTSvc","VNC","Kaseya","Solarwinds","mRemoteNG"],
+ "PATH_EXCLUSION_CONST": 4,
+ "MIN_FILES_FOR_PATH": 4,
+ "VT_THREAT_TOLERANCE": 4,
+ "POLICY_MAP_ENF_AUD": {
+
+ }
+}
\ No newline at end of file
diff --git a/flows/localApproval.py b/flows/localApproval.py
new file mode 100644
index 0000000..1419c65
--- /dev/null
+++ b/flows/localApproval.py
@@ -0,0 +1,292 @@
+# 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 .
+
+
+import datetime
+import logging
+import os
+import re
+import time
+
+import dotenv
+import numpy as np
+import pandas as pd
+
+from models.agent import Agent
+from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
+from services.API import AirlockAPIWrapper
+from services.scheduler import (
+ register_function,
+ run_once_job,
+)
+from utils.utils import colorText, load_env, load_env_json
+
+logger = logging.getLogger(__name__)
+
+dotenv.load_dotenv()
+
+
+def getLocalApprovals(api: AirlockAPIWrapper):
+ working_dir = load_env("WORKING_DIR")
+ result = api.otp_find_awaiting()
+ local_approval = pd.DataFrame(result["response"]["otpusage"])
+ if os.path.exists(f"{working_dir}\\Scheduling\\newest_local_approval.parquet"):
+ previous_run = pd.read_parquet(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
+ previous_run.to_parquet(
+ f"{working_dir}\\Scheduling\\last_local_approval.parquet", index=False
+ )
+ os.remove(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
+
+ # Only keep rows presumably created by the generate local approval function
+ local_approval = local_approval[
+ local_approval["purpose"].str.startswith("🎫 Local Approval 🎫")
+ ]
+
+ local_approval["batchid"] = local_approval["purpose"].apply(
+ lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
+ )
+
+ if not local_approval.empty:
+ local_approval.to_parquet(
+ f"{working_dir}\\Scheduling\\newest_local_approval.parquet", index=False
+ )
+
+ return local_approval
+
+
+def scheduleAddingLAHashes(api: AirlockAPIWrapper):
+
+ policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
+ bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
+ pups = load_env_json("PUPS", "[]")
+ threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int)
+
+ try:
+ register_function("add_hash", returnFromLocalApproval)
+ register_function("move_device", moveAgentToRelatedPolicy)
+ except Exception as e:
+ logger.warning(f"Failed to register functions: {e}")
+ return
+
+ try:
+ approvals_df = getNewLocalApprovals(api)
+ if approvals_df.empty:
+ logger.debug("No new local approvals found. Nothing to schedule.")
+ return
+ batches = approvals_df.groupby("batchid")
+ except Exception as e:
+ logger.warning(f"Failed to retrieve or group local approvals: {e}")
+ return
+
+ for batchid, batch_df in batches:
+ try:
+ duration_minutes = int(batch_df["duration"].iloc[0])
+ start_time = datetime.datetime.now()
+ run_time = start_time + datetime.timedelta(minutes=duration_minutes)
+ early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))
+
+ early_timestamp = early_time.timestamp()
+ run_timestamp = run_time.timestamp()
+
+ # Schedule add_hash job
+ try:
+ run_once_job(
+ f"add_hash_{batchid}",
+ "add_hash",
+ early_timestamp,
+ [
+ api,
+ batch_df,
+ policy_relationship_map,
+ bad_publisher_list,
+ pups,
+ threat_tolerance_constant,
+ ],
+ None,
+ )
+ logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}")
+ except Exception:
+ logger.debug("Failed to schedule add_hash for batch {batchid}: {e}")
+
+ # Schedule move_device jobs
+ devices = batch_df["agentid"].drop_duplicates().tolist()
+ agents = []
+
+ for device in devices:
+ rows = api.agent_find_by_hostname(device).iterrows()
+ agents += [Agent(**row["data"]) for _, row in rows]
+
+ for agent in agents:
+ try:
+ run_once_job(
+ f"move_device_{agent.hostame}_{batchid}",
+ "move_device",
+ run_timestamp,
+ [api, agent, policy_relationship_map],
+ "enforcement",
+ )
+
+ print(
+ f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}"
+ )
+ except Exception as e:
+ print(
+ f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}"
+ )
+
+ except Exception as e:
+ logger.warning(f"Failed to process batch {batchid}: {e}")
+
+
+def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant
+):
+ """
+ # Get unique policy names from device list
+ policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
+
+ # Create inverse map to go from Audit to Enforcement
+ inverse_map = {v: k for k, v in policy_relationship_map.items()}
+
+ # Fetch all policies
+ all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()]
+
+ # Define policy types
+ policy_types = [1, 2, 6, 7]
+
+ #TODO finish logic for adding hashes
+ """
+ working_dir = load_env("WORKING_DIR")
+ policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
+ bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
+ pups = load_env_json("PUPS", "[]")
+ threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
+ print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}")
+
+def moveToLocalApproval(api: AirlockAPIWrapper):
+ possible_durations = [15, 60, 360, 1440, 10080]
+ duration_selected = None
+
+ print(colorText("Please select a duration:", "white"))
+ for i, option in enumerate(possible_durations, start=1):
+ print(f"{i}. {option}")
+
+ try:
+ choice = int(input("Enter the number of your choice: "))
+ if 1 <= choice <= len(possible_durations):
+ duration_selected = possible_durations[choice - 1]
+ print(colorText(f"You selected: {duration_selected}", "yellow"))
+ logger.debug(f"You selected: {duration_selected}")
+ else:
+ print(colorText("❌ Invalid choice.", "red"))
+ logger.debug("Invalid Input")
+ return
+ except ValueError:
+ print(colorText("❌ Invalid input. Please enter a number.", "red"))
+ logger.debug("Invalid Input")
+ return
+
+ agents = selectAgents(api)
+ batch = int(time.time())
+
+ if not agents:
+ print(colorText("❌ No agents found or error retrieving agents.", "red"))
+ logger.debug("No agents found or error retrieving agents")
+ return
+
+ for agent in agents:
+ try:
+ addLocalApproval(api, batch, duration_selected, agent.agentid)
+ moveAgentToRelatedPolicy(api, agent, "audit")
+ except Exception as e:
+ print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))
+
+
+def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
+
+ purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
+ api.otp_generate(agentid, duration_selected, purpose)
+
+
+
+def monitorAuditStatus(api: AirlockAPIWrapper):
+ current_agents = findAllAgents(api)
+ last_agents = []
+ if not last_agents:
+ last_agents = current_agents
+ policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
+
+ # Reverse map for audit → enforcement
+ reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
+ known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())
+
+ # Index last_agents by hostname for quick lookup
+ last_agent_map = {agent.hostname: agent for agent in last_agents}
+
+ # Result buckets
+ newly_added = []
+ same_policy = []
+ moved_to_audit = []
+ moved_to_enforcement = []
+ unusual_move = []
+
+ for current in current_agents:
+ previous = last_agent_map.get(current.hostname)
+
+ if not previous:
+ newly_added.append(current)
+ continue
+
+ if current.groupid == previous.groupid:
+ same_policy.append(current)
+ elif (previous.groupid, current.groupid) in known_transitions:
+ moved_to_audit.append(current)
+ elif (current.groupid, previous.groupid) in known_transitions:
+ moved_to_enforcement.append(current)
+ else:
+ unusual_move.append(current)
+
+ # Return all five DataFrames
+ return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
+
+
+def getNewLocalApprovals(api: AirlockAPIWrapper):
+
+ working_dir = load_env("WORKING_DIR")
+ current_la = getLocalApprovals(api)
+
+ # Load old approval list
+ old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet"
+ if os.path.exists(old_la_path):
+ old_la = pd.read_parquet(old_la_path)
+ else:
+ old_la = pd.DataFrame(columns=current_la.columns)
+
+ # Create composite keys
+ current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
+ old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)
+
+ # Find new entries
+ new_entries = current_la[~current_la["key"].isin(old_la["key"])]
+
+ # Convert 'granted' to datetime and filter by last 10 minutes
+ new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce")
+ ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
+ recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]
+
+ # Save current approvals for next run
+ current_la.drop(columns=["key"], inplace=True)
+ current_la.to_parquet(old_la_path, index=False)
+
+ return recent_entries
diff --git a/utils/otpfunctions.py b/flows/otp.py
similarity index 80%
rename from utils/otpfunctions.py
rename to flows/otp.py
index cd316cb..2966758 100644
--- a/utils/otpfunctions.py
+++ b/flows/otp.py
@@ -13,23 +13,14 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
-#Local Imports
-import utils.clientfunctions as clientf
-import utils.pathfunctions as pathf
-import utils.policyfunctions as policyf
-import utils.utils as ct
-from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid
-#Standard Libary Imports:
-import json
-import math
-import os
-import time
-#3rd Party Imports:
-import pandas as pd
-import requests
+import logging
+logger = logging.getLogger(__name__)
+
+
+"""
def getActiveOTP(url):
endpoint = url + f'/v1/otp/usage'
@@ -37,7 +28,7 @@ def getActiveOTP(url):
"status" : "1"
}
- headers = {"X-APIKey": os.getenv('APIKEY')}
+ headers = {"X-APIKey": load_env('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
@@ -49,12 +40,12 @@ def getActiveOTP(url):
os.remove("OTP\\PARQ\\newest_active_OTP.parquet")
otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False)
if not otp.empty:
- ct.style_dataframe_dark(otp, f"OTP\\HTML\\newest_active_OTP.html")
-
+ formatHTML(otp, f"OTP\\HTML\\newest_active_OTP.html")
+
def getOTPActivities(url, otpid):
endpoint = url + f'/v1/otp/activities'
payload = {"otpid": f"{otpid}"}
- headers = {"X-APIKey": os.getenv('APIKEY')}
+ headers = {"X-APIKey": load_env('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
@@ -77,12 +68,12 @@ def getOTPActivities(url, otpid):
# Optional: generate styled HTML if there's data
if not combined_data.empty:
- ct.style_dataframe_dark(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
+ formatHTML(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
def monitorOTP(url, pups):
-
+
getActiveOTP(url)
-
+
old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet"
new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet"
@@ -94,14 +85,14 @@ def monitorOTP(url, pups):
current_active_OTP = pd.read_parquet(new_otp_path)
-
+
if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])]
still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
no_longer_OTP = old_active_OTP[~old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
-
+
register_function("addhash", addOTPHashes)
@@ -115,7 +106,7 @@ def monitorOTP(url, pups):
#If newly added to the list - schedule adding the majority of the executions prior to the expiration of OTP period.
run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, [url, clientid, pid, pups], None)
print(f"Processing: {pid} with other data: {row}")
-
+
for _, row in still_in_OTP.iterrows():
pid = row['otpid']
@@ -129,11 +120,12 @@ def monitorOTP(url, pups):
pid = row['otpid']
allowlist = clientf.getDestAllowlistFromClientID(url,clientid)
policy, policyid = clientf.getPolicyFromClientID(url,clientid)
-
- """
- Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
- But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
- """
+
+"""
+
+# Devices can come out of OTP either by timeout, or by early move out of OTP. If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
+# But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
+"""
getOTPActivities(url,pid)
find_and_prioritize_jobs_by_pid(pid, 1)
@@ -150,7 +142,7 @@ def monitorOTP(url, pups):
history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
history = pd.concat([history, finalhashesadded], ignore_index=True)
- ct.style_dataframe_dark(history, f"localapproval_history.html")
+ formatHTML(history, f"localapproval_history.html")
history.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
@@ -161,7 +153,7 @@ def addOTPHashes(url, clientid, otpid, pups):
path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
activities = pd.read_parquet(path)
- pattern = pathf.regulator(pups)
+ pattern = regulator(pups)
allowlist = clientf.getDestAllowlistFromClientID(url, clientid)
# Initialize or preserve 'hash_added' column
@@ -176,7 +168,8 @@ def addOTPHashes(url, clientid, otpid, pups):
# Add hashes to policy
if hashes_to_add:
- policyf.addHash(url, allowlist, hashes_to_add)
+ #TODO add the api call
+ pass
# Update 'hash_added' column
activities["hash_added"] = activities.apply(
@@ -189,12 +182,12 @@ def addOTPHashes(url, clientid, otpid, pups):
activities.to_parquet(path)
def generateOTP(url, agentid):
- purpose = input(ct.colorText(" Please enter the purpose for the OTP: ", "white"))
+ purpose = input(colorText(" Please enter the purpose for the OTP: ", "white"))
possible_durations = [15, 60, 360, 1440, 10080]
duration_selected = " "
- print(ct.colorText("Please select a duration:", "white"))
+ print(colorText("Please select a duration:", "white"))
for i, option in enumerate(possible_durations, start=1):
print(f"{i}. {option}")
@@ -202,11 +195,11 @@ def generateOTP(url, agentid):
choice = int(input("Enter the number of your choice: "))
if 1 <= choice <= len(possible_durations):
duration_selected = possible_durations[choice - 1]
- print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
+ print(colorText(f"You selected: {duration_selected}", "yellow"))
else:
- print(ct.colorText("Invalid choice.", "red"))
+ print(colorText("Invalid choice.", "red"))
except ValueError:
- print(ct.colorText("Invalid input. Please enter a number.", "red"))
+ print(colorText("Invalid input. Please enter a number.", "red"))
endpoint = url + '/v1/otp/retrieve'
payload = {
@@ -215,12 +208,13 @@ def generateOTP(url, agentid):
"purpose" : f"{purpose}"
}
- headers = {"X-APIKey": os.getenv('APIKEY')}
+ headers = {"X-APIKey": load_env('APIKEY')}
payload = json.dumps(payload)
response = requests.post(endpoint, headers=headers, data=payload, verify=False)
result = json.loads(response.text)
otpcode = result["response"]["otpcode"]
- print(ct.colorText(f"The OPT code is: {otpcode}", "yellow"))
+ print(colorText(f"The OPT code is: {otpcode}", "yellow"))
+"""
diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py
new file mode 100644
index 0000000..910e629
--- /dev/null
+++ b/flows/prepPolicy.py
@@ -0,0 +1,351 @@
+# 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 .
+
+import logging
+import os
+import os.path
+from typing import List
+
+import dotenv
+import pandas as pd
+
+from models.execution import ExecutionHistoryRecord, Hash
+from models.policy import Allowlist, Policy
+from services.API import AirlockAPIWrapper
+from services.selector import Selector
+from utils.utils import (
+ colorText,
+ formatHTML,
+ import_to_dataframe,
+ load_env,
+ load_env_json,
+ regulator,
+)
+
+logger = logging.getLogger(__name__)
+
+dotenv.load_dotenv()
+
+
+
+def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
+
+ allowlists = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
+ logger.debug("Prompting for Policies")
+ print(colorText("Please select policy/policies", "white"))
+ selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
+
+ if selected is None:
+ return []
+
+ # Normalize to always return a list
+ logger.debug("Returning {selected.dict}")
+ return selected if isinstance(selected, list) else [selected]
+
+
+def selectAllowlists(api: AirlockAPIWrapper, allow_multiple=True) -> List[Allowlist]:
+
+ allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
+ logger.debug("Prompting for Allowlist(s)")
+ print(colorText("Please select allowlist(s)", "white"))
+ selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)
+
+ if selected is None:
+ return []
+
+ # Normalize to always return a list
+ logger.debug(f"Returning {selected}")
+ return selected if isinstance(selected, list) else [selected]
+
+
+def sortHashes(
+ api: AirlockAPIWrapper,
+ selected_policies: List[Policy],
+ type=[1, 2, 6, 7]
+):
+ working_dir = load_env("WORKING_DIR")
+ history_days = Selector.select_value(
+ prompt="Enter how many days of history to pull (1–150): ",
+ value_type=int,
+ valid_range=(1, 150),
+ )
+
+ logger.debug(f"{history_days} day selected for history")
+
+ if history_days is None:
+ logging.warning("No history range selected. Aborting.")
+ return
+
+ executions = []
+ hashes = []
+
+ # Pull execution histories for each policy
+
+ policy_executions = ExecutionHistoryRecord.from_policies(
+ api, selected_policies, type_=type, history_days=history_days
+ )
+
+ logger.debug(f"Policy_executions is {policy_executions}")
+
+ executions.extend(policy_executions)
+ logger.debug(f"Executions contains {executions}")
+ if executions:
+ hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
+ ]
+
+ if hashes:
+ unique_hashes = Hash.deduplicate(hashes)
+
+ needs_review, approved, unapproved = Hash.categorize_hashes(
+ hashes=unique_hashes
+ )
+
+ categories = {
+ "needs_review": needs_review,
+ "approved": approved,
+ "unapproved": unapproved,
+ }
+
+ for label, category in categories.items():
+ csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv"
+ html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html"
+
+ ExecutionHistoryRecord.enrich_with_hashes_and_export(
+ executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
+ )
+ df = import_to_dataframe(csv_path)
+ formatHTML(df, html_path)
+
+def buildPathsandPublishers(split):
+ working_dir = load_env("WORKING_DIR")
+ df1 = pd.DataFrame()
+ df2 = pd.DataFrame()
+ all_approved_hashes = pd.DataFrame()
+ path1 = f"{working_dir}\\Approved\\approved_executions.csv"
+ path2 = f"{working_dir}\\Approved\\needs_review_executions.csv"
+
+ if os.path.exists(path1):
+ df1 = pd.read_csv(path1)
+ else:
+ logger.warning(f"File not found: {path1}")
+
+ if os.path.exists(path2):
+ df2 = pd.read_csv(path2)
+ else:
+ logger.warning(f"File not found: {path2}")
+
+ if df1.empty and df2.empty:
+ logger.warning("Both DataFrames are empty. Skipping sort.")
+ all_approved_hashes = pd.DataFrame()
+ logger.debug(all_approved_hashes.head)
+ else:
+ all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
+ if "filename_exec" in all_approved_hashes.columns:
+ all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
+ else:
+ logger.warning("Warning: 'filename_exec' column not found in concatenated DataFrame.")
+
+ if not all_approved_hashes.empty:
+ primary_path_exclusions = calculatePath(
+ all_approved_hashes,
+ split,
+ )
+ remaining_hashes = all_approved_hashes[
+ ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
+ ]
+ secondary_path_exclusions = calculatePath(
+ remaining_hashes, split
+ )
+ remaining_hashes = remaining_hashes[
+ ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])
+ ]
+ dataframes = {
+ "primary_Paths": primary_path_exclusions,
+ "secondary_Paths": secondary_path_exclusions,
+ "hashes_to_add": remaining_hashes,
+ }
+
+ for name, df in dataframes.items():
+ df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False)
+ formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html")
+
+ if not all_approved_hashes.empty:
+ # Drop all not signed, only keep unique values
+ publist = all_approved_hashes[
+ all_approved_hashes["publisher_hash"] != "Not Signed"
+ ].drop_duplicates(subset=["publisher_hash"])
+ # Remove Bad publisher if somehow they made it this far
+ pattern = regulator(load_env_json("BAD_PUBLISHERS","[]"))
+ publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
+ publist = publist[["publisher_hash"]]
+ publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False)
+
+def buildPreflights():
+ working_dir = load_env("WORKING_DIR")
+
+ df1 = pd.DataFrame()
+ df2 = pd.DataFrame()
+ approved_hashes = pd.DataFrame()
+ approved_publishers = pd.DataFrame()
+
+ hash = f"{working_dir}\\Approved\\hashes_to_add.csv"
+ path1 = f"{working_dir}\\Approved\\primary_Paths.csv"
+ path2 = f"{working_dir}\\Approved\\secondary_Paths.csv"
+ publishers = f"{working_dir}\\Approved\\publishers.csv"
+
+ if os.path.exists(hash):
+ approved_hashes = pd.read_csv(hash)
+
+ else:
+ logger.warning(f"File not found: {hash}")
+
+ if os.path.exists(path1):
+ df1 = pd.read_csv(path1)
+ else:
+ logger.warning(f"File not found: {path1}")
+
+ if os.path.exists(path2):
+ df2 = pd.read_csv(path2)
+ else:
+ logger.warning(f"File not found: {path2}")
+
+ if df1.empty and df2.empty:
+ logger.warning("Both DataFrames are empty. Skipping sort.")
+ approved_paths = pd.DataFrame()
+ else:
+ approved_paths = pd.concat([df1, df2], ignore_index=True)
+
+ if os.path.exists(publishers):
+ approved_publishers = pd.read_csv(publishers)
+
+ else:
+ logger.warning(f"File not found: {publishers}")
+
+ dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers}
+
+ for name, df in dataframes.items():
+ df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False)
+ formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")
+
+def splitFilepathsGrouped(df, col="filename"):
+ path_exclusion_constant = load_env("PATH_EXCLUSION_CONST", cast_type= int)
+ min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type= int)
+
+ def clean_split(path):
+ if not isinstance(path, (str, bytes, os.PathLike)):
+ return []
+ parts = os.path.normpath(path).split(os.sep)
+ parts = [p for p in parts if p] # Remove empty strings
+ return parts
+
+ # Diagnostic: log any non-string entries
+ non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
+ if not non_string_entries.empty:
+ print(f"[WARNING] Non-string entries found in column '{col}':")
+ print(non_string_entries)
+
+ df = df.copy()
+ split_paths = df[col].apply(clean_split)
+
+ # Filter out paths with fewer than `min_files_for_path` components
+ df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
+ split_paths = split_paths[df.index]
+
+ df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
+ grouped = df.groupby("group_key")
+ new_rows = []
+
+ for _, group_df in grouped:
+ paths = group_df[col].tolist()
+ split_parts = [clean_split(p) for p in paths]
+
+ def longest_common_prefix(paths):
+ if not paths:
+ return []
+ prefix = paths[0]
+ for path in paths[1:]:
+ prefix = [a for a, b in zip(prefix, path) if a == b]
+ if not prefix:
+ break
+ return prefix
+
+ common_prefix = longest_common_prefix(split_parts)
+ prefix_str = os.sep.join(common_prefix)
+
+ for i, parts in enumerate(split_parts):
+ filename = parts[-1]
+ middle = (
+ os.sep.join(parts[len(common_prefix):-1])
+ if len(parts) > len(common_prefix) + 1
+ else ""
+ )
+ row = group_df.iloc[i].copy()
+ row["longestcfp"] = prefix_str
+ row["middle"] = middle
+ row["filename_only"] = filename
+ row["file_extension"] = os.path.splitext(filename)[1].lower()
+ new_rows.append(row)
+
+ return pd.DataFrame(new_rows).drop(columns=["group_key"])
+
+def calculatePath(approved_hashes, split):
+ if split:
+ dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
+ else:
+ dfs_by_policy = [approved_hashes]
+
+ badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
+ min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type = int)
+
+ processed_dfs = []
+
+ for df in dfs_by_policy:
+ haslcp = splitFilepathsGrouped(df, "filename_exec")
+ haslcp = haslcp.drop_duplicates()
+
+ forbidden = regulator(badpathparts, True)
+ forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
+
+ logger.debug("Removing forbidden filepaths for path exceptions")
+ print(colorText("Removing forbidden filepaths for path exceptions", "green"))
+ lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
+
+ lcp_not_forbidden_review = lcp_not_forbidden[
+ [
+ "policyname",
+ "longestcfp",
+ "middle",
+ "filename_only",
+ "file_extension",
+ "sha256",
+ ]
+ ]
+
+ unique_sha_counts = (
+ lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index()
+ )
+ unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]
+
+ lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
+ unique_sha_counts, on="longestcfp", how="left"
+ )
+ lcp_not_forbidden_review = lcp_not_forbidden_review[
+ lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
+ ]
+ processed_dfs.append(lcp_not_forbidden_review)
+
+ pathExclusions = pd.concat(processed_dfs, ignore_index=True)
+
+ return pathExclusions
diff --git a/flows/quietAgent.py b/flows/quietAgent.py
new file mode 100644
index 0000000..0e1c3c3
--- /dev/null
+++ b/flows/quietAgent.py
@@ -0,0 +1,123 @@
+# 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 .
+
+import datetime
+import logging
+
+import pandas as pd
+
+from flows.prepPolicy import selectPolicies
+from services.API import AirlockAPIWrapper
+from services.policyhandler import getPolicyInfo
+from services.selector import Selector
+from utils.utils import colorText, load_env
+
+logger = logging.getLogger(__name__)
+import dotenv
+
+dotenv.load_dotenv()
+
+
+def findQuietAgents(api: AirlockAPIWrapper):
+ working_dir = load_env("WORKING_DIR")
+ # Get policy selection and agent list
+ selected_policy = selectPolicies(api, False)
+ if selected_policy:
+ agents = api.agents_find_by_group(selected_policy[0].groupid)
+
+ # Prompt user for history range
+ history_days = Selector.select_value(
+ prompt="Enter how many days of history to pull (1–150): ",
+ value_type=int,
+ valid_range=(1, 150),
+ )
+ required_quiet = Selector.select_value(
+ prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
+ value_type=int,
+ valid_range=(1, 365),
+ )
+
+ # Get execution history as a DataFrame
+ policy_exec_history = getPolicyInfo(
+ api, selected_policy[0], [1, 2, 6, 7], history_days
+ )
+
+
+ if policy_exec_history.empty:
+ logging.info("No execution history found for the selected policy and time range.")
+ return
+
+
+ # Convert 'datetime' column to timezone-aware datetime objects
+ policy_exec_history["datetime"] = pd.to_datetime(
+ policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
+ )
+
+ # Get current UTC time
+ now = datetime.datetime.now(datetime.timezone.utc)
+
+ # Calculate days ago
+ policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
+ lambda dt: (now - dt).days
+ )
+
+ # Count total executions per hostname
+ hostname_counts = policy_exec_history["hostname"].value_counts()
+
+ # Map execution counts to agents
+ agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)
+
+ # Find most recent execution per hostname
+ most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
+ subset="hostname", keep="first"
+ )
+
+ # Map most recent execution age to agents
+ agents["days_since"] = agents["hostname"].map(
+ most_recent_exec.set_index("hostname")["days_ago"]
+ )
+
+ # Check for enforcement readiness
+ agents["required_quiet"] = required_quiet
+ agents["enforce_ready"] = agents["days_since"].apply(
+ lambda x: True if pd.isna(x) or x > required_quiet else False
+ )
+
+ # Sort agents by execution count and hostname
+ agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])
+
+ # Save to CSV
+ filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
+ logging.debug(f"Saving CSV to {filename}")
+ print(colorText(f"Saving CSV to {filename}", "green"))
+ agents.to_csv(filename, index=False)
+
+ # Summary statistics
+ total_agents = len(agents)
+ ready_agents = agents["enforce_ready"].sum()
+ not_ready_agents = total_agents - ready_agents
+ ready_percentage = (ready_agents / total_agents) * 100
+
+ # Print results
+
+
+ message = (
+ f"Total agents: {total_agents}\n"
+ f"Agents marked as 'enforce_ready': {ready_agents}\n"
+ f"Agents not ready: {not_ready_agents}\n"
+ f"Percentage ready for enforcement: {ready_percentage:.2f}%"
+ )
+ logger.debug(message)
+ colorText(message,"green")
diff --git a/models/agent.py b/models/agent.py
new file mode 100644
index 0000000..def8e80
--- /dev/null
+++ b/models/agent.py
@@ -0,0 +1,64 @@
+# 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 .
+
+from dataclasses import dataclass, field
+from typing import ClassVar, Optional
+
+
+@dataclass
+class Agent:
+ agentid: str
+ clientversion: str
+ domain: str
+ freespace: int
+ groupid: int
+ hostname: str
+ ip: str
+ localip: str
+ lastcheckin: str
+ os: str
+ policyversion: str
+ status: int # raw status code
+ username: str
+ groupname: Optional[str] = field(default=None)
+ status_text: Optional[str] = field(default=None)
+
+ # Class-level status map
+ status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"}
+
+ def enrich(self, groupid_to_name: dict):
+ """Enrich the agent with groupname and human-readable status."""
+ self.groupname = groupid_to_name.get(self.groupid, None)
+ self.status_text = self.status_map.get(self.status, "Unknown")
+
+
+"""
+
+from models.agent import Agent
+from modesls.policy
+
+# Step 1: Load data from API
+policies = [Policy(**row['data']) for _, row in api.policy_find_all().iterrows()]
+agents = [Agent(**row['data']) for _, row in api.agent_find_all().iterrows()]
+
+# Step 2: Create groupid → groupname map
+groupid_to_name = {policy.groupid: policy.name for policy in policies}
+
+# Step 3: Enrich agents
+for agent in agents:
+ agent.enrich(groupid_to_name)
+
+
+"""
diff --git a/models/execution.py b/models/execution.py
new file mode 100644
index 0000000..f980c91
--- /dev/null
+++ b/models/execution.py
@@ -0,0 +1,389 @@
+# 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 .
+
+import inspect
+import json
+import logging
+import os
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional
+
+import dotenv
+import pandas as pd
+
+from services.policyhandler import pullPolicyExechistories
+from utils.utils import colorText, load_env, load_env_json, regulator
+
+logger = logging.getLogger(__name__)
+
+dotenv.load_dotenv()
+
+
+class Hash:
+ """
+ Hash model representing Hash data
+ """
+
+ def __init__(
+ self,
+ sha256,
+ applications=None,
+ baselines=None,
+ blocklists=None,
+ createtime=None,
+ datetime=None,
+ description=None,
+ filename=None,
+ filepath=None,
+ filesize=None,
+ md5=None,
+ modtime=None,
+ origname=None,
+ productname=None,
+ productversion=None,
+ publisher=None,
+ reputation=None,
+ sha128=None,
+ sha384=None,
+ sha512=None,
+ ):
+ self.sha256 = sha256
+ self.applications = applications
+ self.baselines = baselines
+ self.blocklists = blocklists
+ self.createtime = createtime
+ self.datetime = datetime
+ self.description = description
+ self.filename = filename
+ self.filepath = filepath
+ self.filesize = filesize
+ self.md5 = md5
+ self.modtime = modtime
+ self.origname = origname
+ self.productname = productname
+ self.productversion = productversion
+ self.publisher = publisher
+ self.reputation = reputation
+ self.sha128 = sha128
+ self.sha384 = sha384
+ self.sha512 = sha512
+
+ def __repr__(self):
+ attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
+ return f""
+
+ def __eq__(self, other):
+ if isinstance(other, Hash):
+ return self.sha256 == other.sha256
+ return False
+
+ def __hash__(self):
+ return hash(self.sha256)
+
+ def to_dict(self):
+ """Returns a dictionary representation of the hash."""
+ return self.__dict__
+
+ @staticmethod
+ def safe_int(value, default=0):
+ """Safely convert a value to int, returning default on failure."""
+ try:
+ return int(value)
+ except (TypeError, ValueError):
+ return default
+
+
+ @classmethod
+ def deduplicate(cls, hash_list):
+ """
+ Deduplicates a list of Hash objects based on sha256.
+ Args:
+ hash_list (list): List of Hash instances.
+ Returns:
+ list: Deduplicated list of Hash instances.
+ """
+ seen = set()
+ deduped = []
+ for h in hash_list:
+ if h.sha256 not in seen:
+ seen.add(h.sha256)
+ deduped.append(h)
+ return deduped
+
+ @classmethod
+ def categorize_hashes(cls, hashes):
+ import re
+ from utils.utils import load_env, load_env_json, regulator
+
+ threat_tolerance = load_env("VT_THREAT_TOLERANCE", cast_type=int)
+ bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
+ pups_pattern = regulator(load_env_json("PUPS", "[]"))
+
+ needs_review = []
+ approved = []
+ unapproved = []
+
+ def reputationtool(hash_obj):
+ val = hash_obj.reputation.get("scannermatch") if isinstance(hash_obj.reputation, dict) else None
+ if val in [None, "N/A"]:
+ return hash_obj.publisher == "Not Signed"
+ try:
+ return int(val) > threat_tolerance
+ except (ValueError, TypeError):
+ return hash_obj.publisher == "Not Signed"
+
+ for hash_obj in hashes:
+ publisher = hash_obj.publisher or ""
+ description = hash_obj.description or ""
+ reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
+ rep_status = reputation.get("status")
+
+ rep_flag = reputationtool(hash_obj)
+
+ is_signed = publisher != "Not Signed"
+ is_untrusted = re.search(bad_publishers_pattern, publisher, re.IGNORECASE) is not None
+ is_pup = re.search(pups_pattern, description, re.IGNORECASE) is not None
+ has_known_status = rep_status == "KNOWN"
+
+ if (not is_signed and rep_flag) or rep_status == "UNKNOWN":
+ needs_review.append(hash_obj)
+ elif (
+ (is_signed and not is_untrusted and has_known_status and not is_pup) or
+ (not is_signed and not rep_flag and not is_untrusted and has_known_status and not is_pup)
+ ):
+ approved.append(hash_obj)
+ else:
+ unapproved.append(hash_obj)
+
+ return needs_review, approved, unapproved
+
+
+ @classmethod
+ def export_to_csv(cls, hash_list, directory_path):
+ """
+ Exports a list of Hash objects to a CSV file in the specified directory.
+ The filename is derived from the variable name of the list if possible,
+ and includes a timestamp to ensure uniqueness.
+ """
+ filename = "hashes_export.csv"
+ frame = inspect.currentframe()
+ if frame is not None and frame.f_back is not None:
+ callers_local_vars = frame.f_back.f_locals.items()
+ for var_name, var_val in callers_local_vars:
+ if var_val is hash_list:
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ filename = f"{var_name}_{timestamp}.csv"
+ break
+ else:
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ filename = f"hashes_export_{timestamp}.csv"
+
+ os.makedirs(directory_path, exist_ok=True)
+ file_path = os.path.join(directory_path, filename)
+
+ df = pd.DataFrame([h.to_dict() for h in hash_list])
+ df.to_csv(file_path, index=False)
+
+ logger.info(f"CSV file saved to: {file_path}")
+
+
+"""
+#Example - Convert Dataframe returned by hash query into hash objects
+hash_objects = []
+for _, row in df.iterrows():
+ try:
+ parsed_data = ast.literal_eval(row['data'])
+ hash_obj = Hash(sha256=row['sha256'], **parsed_data)
+ hash_objects.append(hash_obj)
+ except Exception as e:
+ print(f"Error parsing row: {e}")
+
+# Display the created Hash objects
+for obj in hash_objects:
+ print(obj)
+
+
+# Categorize hashes
+needs_review, approved, unapproved = Hash.categorize_hashes(
+ hashes=hash_objects,
+ threat_tolerance=3,
+ untrusted_pattern=untrusted_pattern,
+ pups_pattern=pups_pattern
+)
+
+ # Deduplicate
+ deduped_hashes = Hash.deduplicate(hash_list)
+
+# Specify the directory where you want to save the CSV
+output_directory = "C:/Users/Brandon/Documents/HashExports"
+
+# Call the export method
+Hash.export_to_csv(hashes_for_export, output_directory)
+
+
+"""
+
+
+@dataclass
+class ExecutionHistoryRecord:
+ # Mandatory fields
+ username: str
+ hostname: str
+ netdomain: str
+ filename: str
+ ppolicy: str
+ policyname: str
+ policyver: str
+ commandline: str
+ publisher: str
+ sha256: str
+ datetime: str
+
+ # Optional fields
+ type: Optional[int] = None
+ pprocess: Optional[str] = None
+ gprocess: Optional[str] = None
+ md5: Optional[str] = None
+ sha128: Optional[str] = None
+ sha384: Optional[str] = None
+ sha512: Optional[str] = None
+ ip: Optional[str] = None
+ localip: Optional[str] = None
+ extid: Optional[str] = None
+ extname: Optional[str] = None
+ exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
+ extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge
+
+ @staticmethod
+ def enrich_with_hashes_and_export(
+ executions: list, hashes: list, directory_path: str, label: str = "enriched"
+ ):
+ exec_df = pd.DataFrame([e.__dict__ for e in executions])
+ hash_df = pd.DataFrame([h.to_dict() for h in hashes])
+
+ logger.debug(f"Execution DataFrame columns: {exec_df.columns}")
+ logger.debug(f"Hash DataFrame columns: {hash_df.columns}")
+
+ if hash_df.empty:
+ logger.warning(f"hash_df is empty for label: {label}. Skipping merge.")
+ merged_df = exec_df.copy()
+ else:
+ merged_df = pd.merge(
+ exec_df,
+ hash_df,
+ on="sha256",
+ how="left", # Preserve all executions, enrich where possible
+ suffixes=("_exec", "_hash")
+ )
+ logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}")
+
+ filename = f"{label}_executions.csv"
+ os.makedirs(directory_path, exist_ok=True)
+ file_path = os.path.join(directory_path, filename)
+ merged_df.to_csv(file_path, index=False)
+
+ logger.info(f"CSV file saved to: {file_path}")
+
+
+ @classmethod
+ def from_dict(cls, data: dict):
+ mandatory_fields = [
+ "username",
+ "hostname",
+ "netdomain",
+ "filename",
+ "ppolicy",
+ "policyname",
+ "policyver",
+ "commandline",
+ "publisher",
+ "sha256",
+ "datetime",
+ ]
+ missing_fields = [
+ field for field in mandatory_fields if field not in data or data[field] is None
+ ]
+ if missing_fields:
+ raise ValueError(f"Missing mandatory fields: {missing_fields}")
+
+ return cls(
+ username=data["username"],
+ hostname=data["hostname"],
+ netdomain=data["netdomain"],
+ filename=data["filename"],
+ ppolicy=data["ppolicy"],
+ policyname=data["policyname"],
+ policyver=data["policyver"],
+ commandline=data["commandline"],
+ publisher=data["publisher"],
+ sha256=data["sha256"],
+ datetime=data["datetime"],
+ type=data.get("type"),
+ pprocess=data.get("pprocess"),
+ gprocess=data.get("gprocess"),
+ md5=data.get("md5"),
+ sha128=data.get("sha128"),
+ sha384=data.get("sha384"),
+ sha512=data.get("sha512"),
+ ip=data.get("ip"),
+ localip=data.get("localip"),
+ extid=data.get("extid"),
+ extname=data.get("extname"),
+ exttype=data.get("exttype"),
+ extbrowser=data.get("extbrowser"),
+ )
+
+ @classmethod
+ def from_policies(
+ cls, api, selected_policies, type_: list, history_days: int
+ ) -> List["ExecutionHistoryRecord"]:
+ executions = []
+ for policy in selected_policies:
+ execs = pullPolicyExechistories(
+ api, policy, type_, history_days, True
+ )
+ if execs:
+ data = json.loads(execs)
+ exechistories = data.get("response", {}).get("exechistories", [])
+ if not exechistories:
+ continue
+
+ df = pd.DataFrame(exechistories)
+ df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
+ df = df.sort_values(by=["sha256", "filename"])
+
+ executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()])
+ logger.debug(f"Staging of Execution history for policy: {policy.name} is complete")
+ print(
+ colorText(
+ f"Staging of Execution history for policy: {policy.name} is complete",
+ "green",
+ )
+ )
+
+ return executions
+
+ def __repr__(self):
+ attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
+ return f""
+
+
+"""
+executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)
+
+ExecutionHistoryRecord.enrich_with_hashes_and_export(executions, hash_objects, "C:/Users/Brandon/Documents/EnrichedExports")
+"""
diff --git a/models/policy.py b/models/policy.py
new file mode 100644
index 0000000..7040f32
--- /dev/null
+++ b/models/policy.py
@@ -0,0 +1,65 @@
+# 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 .
+
+import json
+
+"""
+Policy model representing policy data and relationships.
+"""
+
+
+class Policy:
+ def __init__(self, groupid, hidden, name, parent):
+ self.groupid = groupid
+ self.hidden = hidden
+ self.name = name
+ self.parent = parent
+
+ def __repr__(self):
+ # Show all current attributes, including dynamically added ones
+ attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
+ return f""
+
+ def to_dict(self):
+ # Return all attributes as a dictionary
+ return self.__dict__
+
+ def to_json(self):
+ # Convert to JSON string, handling non-serializable types gracefully
+ return json.dumps(self.to_dict(), default=str)
+
+
+class Allowlist:
+ """
+ Represents Allowlist
+ """
+
+ def __init__(self, applicationid, name, version):
+ self.applicationid = applicationid
+ self.name = name
+ self.version = version
+
+ def __repr__(self):
+ # Show all current attributes, including dynamically added ones
+ attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
+ return f""
+
+ def to_dict(self):
+ # Return all attributes as a dictionary
+ return self.__dict__
+
+ def to_json(self):
+ # Convert to JSON string, handling non-serializable types gracefully
+ return json.dumps(self.to_dict(), default=str)
diff --git a/requirements.txt b/requirements.txt
index 8a3c64f..bb00042 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,29 +1,11 @@
-bson==0.5.10
-certifi==2025.8.3
-charset-normalizer==3.4.3
-colorama==0.4.6
-cramjam==2.11.0
-docopt==0.6.2
-dotenv==0.9.9
-fastparquet==2024.11.0
-fsspec==2025.9.0
-idna==3.10
-ijson==3.4.0
-lxml==6.0.0
-markdown-it-py==4.0.0
-mdurl==0.1.2
+cryptography==46.0.1
+keyring==25.6.0
numpy==2.3.2
-packaging==25.0
pandas==2.3.1
-pretty-tables==3.1.0
-pyarrow==21.0.0
-Pygments==2.19.2
-python-dateutil==2.9.0.post0
python-dotenv==1.1.1
-pytz==2025.2
-requests==2.32.4
-six==1.17.0
+pymongo
+requests==2.32.5
+schedule==1.2.2
tqdm==4.67.1
-tzdata==2025.2
urllib3==2.5.0
-yarg==0.1.10
\ No newline at end of file
+bson==0.5.10
\ No newline at end of file
diff --git a/services/API.py b/services/API.py
new file mode 100644
index 0000000..a51eaf6
--- /dev/null
+++ b/services/API.py
@@ -0,0 +1,239 @@
+# 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 .
+
+import json
+import logging
+from typing import Dict, List, Optional
+
+import pandas as pd
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+
+class AirlockAPIWrapper:
+ """
+ A wrapper class for interacting with the Airlock API.
+ Provides methods for managing agents, policies, hashes, OTPs, and execution history.
+ """
+
+ def __init__(self, base_url: str, api_key: str):
+ """
+ Initialize the API wrapper.
+
+ Parameters:
+ - base_url (str): Base URL of the Airlock API.
+ - api_key (str): API key for authentication.
+ """
+ self.base_url = base_url.rstrip("/")
+ self.api_key = api_key
+ self.headers = {"X-APIKey": self.api_key}
+
+ def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
+ """
+ Internal method to send POST requests to the API.
+
+ Parameters:
+ - endpoint (str): API endpoint.
+ - payload (dict, optional): Request payload.
+
+ Returns:
+ - dict: JSON response from the API.
+ """
+ url = f"{self.base_url}{endpoint}"
+ data = json.dumps(payload or {})
+ try:
+ logger.debug(f"POST Request to {url} with payload: {payload}")
+ response = requests.post(url, headers=self.headers, data=data, verify=False)
+ response.raise_for_status()
+ logger.debug(f"Response received from {url}")
+ return response.json()
+ except requests.exceptions.RequestException as e:
+ logger.error(f"API request failed: {e}")
+ raise
+
+ # Allowlist Management
+ def allowlist_find_all(self) -> pd.DataFrame:
+ """
+ Retrieve all applications in the allowlist.
+
+ Returns:
+ - pd.DataFrame: DataFrame containing allowlisted applications.
+ """
+ result = self._post("/v1/application", {})
+ return pd.DataFrame(result["response"]["applications"])
+
+ # Agent Management
+ def agent_find_all(self) -> pd.DataFrame:
+ """Retrieve all agents."""
+ result = self._post("/v1/agent/find", {})
+ return pd.DataFrame(result["response"]["agents"])
+
+ def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
+ """Find agents by hostname."""
+ payload = {"hostname": hostname}
+ result = self._post("/v1/agent/find", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
+ """Find agents by agent ID."""
+ payload = {"agentid": agentid}
+ result = self._post("/v1/agent/find", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ def agent_find_by_status(self, status: int) -> pd.DataFrame:
+ """Find agents by status (0 = Offline, 1 = Online, 3 = Safemode)."""
+ payload = {"status": status}
+ result = self._post("/v1/agent/find", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ def agent_find_by_username(self, username: str) -> pd.DataFrame:
+ """Find agents by username."""
+ payload = {"username": username}
+ result = self._post("/v1/agent/find", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ def agent_move(self, agentid: str, groupid: str) -> dict:
+ """Move an agent to a different group."""
+ payload = {"agentid": agentid, "groupid": groupid}
+ return self._post("/v1/agent/move", payload)
+
+ def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
+ """Find agents by group ID."""
+ payload = {"groupid": groupid}
+ result = self._post("/v1/agent/find", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ # Hash Management
+ def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
+ """Add hashes to the allowlist for a specific application."""
+ payload = {"applicationid": applicationid, "hashes": hashes}
+ return self._post("/v1/hash/application/add", payload)
+
+ def hash_query(self, hashes: List[str]) -> pd.DataFrame:
+ """Query information about specific hashes."""
+ payload = {"hashes": hashes}
+ result = self._post("/v1/hash/query", payload)
+ return pd.DataFrame(result["response"]["results"])
+
+ # OTP Management
+ def otp_find_active(self) -> pd.DataFrame:
+ """Find active OTPs."""
+ payload = {"status": "1"}
+ result = self._post("/v1/otp/usage", payload)
+ return pd.DataFrame(result["response"]["otpusage"])
+
+ def otp_find_awaiting(self) -> pd.DataFrame:
+ """Find OTPs that are awaiting activation."""
+ payload = {"status": "0"}
+ result = self._post("/v1/otp/usage", payload)
+ return pd.DataFrame(result["response"]["otpusage"])
+
+ def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
+ """Generate a new OTP for an agent."""
+ payload = {
+ "duration": str(duration),
+ "agentid": str(agentid),
+ "purpose": purpose,
+ }
+ result = self._post("/v1/otp/retrieve", payload)
+ return result["response"]["otpcode"]
+
+ def otp_get_activities(self, otpid: str) -> pd.DataFrame:
+ """Retrieve activities associated with a specific OTP."""
+ payload = {"otpid": otpid}
+ result = self._post("/v1/otp/activities", payload)
+ return pd.DataFrame(result["response"]["otpactivities"])
+
+ # Policy Management
+ def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
+ """Add path exclusions to a policy group."""
+ payload = {"groupid": groupid, "path": paths}
+ return self._post("/v1/group/path/add", payload)
+
+ def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
+ """Add publishers to a policy group."""
+ payload = {"groupid": groupid, "publisher": publishers}
+ return self._post("/v1/group/publisher/add", payload)
+
+ def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
+ """Clone a policy from one group to another."""
+ payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
+ return self._post("/v1/group/assign", payload)
+
+ def policy_find_all(self) -> pd.DataFrame:
+ """Retrieve all policy groups."""
+ result = self._post("/v1/group")
+ return pd.DataFrame(result["response"]["groups"])
+
+ def policy_list_agents(self, groupid: str) -> pd.DataFrame:
+ """List agents assigned to a specific policy group."""
+ payload = {"groupid": groupid}
+ result = self._post("/v1/group/agents", payload)
+ return pd.DataFrame(result["response"]["agents"])
+
+ def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
+ """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
+ payload = {"groupid": groupid, "auditmode": auditmode}
+ return self._post("/v1/group/settings/auditmode", payload)
+
+ # Execution History
+ def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str:
+ """Retrieve execution history logs."""
+ payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
+ result = self._post("/v1/logging/exechistories", payload)
+ return result["response"]["exechistories"]
+
+ def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]:
+ """
+ Retrieve execution history logs.
+
+ "datefrom":"", //(Optional) Datefrom is for date range search, formatted as "YYYY-MM-DD"
+ "dateto":"", //(Optional) Dateto is for date range search, formatted as "YYYY-MM-DD"
+ "category":"", //(Optional) Category for filtering type
+ "hostname":"", //(Optional) Hostname to filter
+ "username":"admin", //(Optional) Username to filter
+ "netdomain":"", //(Optional) Domain (or group) to filter
+ "filename":"", //(Optional) Filename to filter
+ "ppolicy":"", //(Optional) Parent Policy name to filter
+ "policyname":"", //(Optional) Policy name to filter
+ "policyver":"", //(Optional) Policy version to filter (e.g. "v95")
+ "commandline":"", //(Optional) Commandline to filter
+ "publisher":"", //(Optional) Publisher to filter
+ "pprocess":"", //(Optional) Parent Process to filter
+ "sha256":"", //(Optional) SHA256 hash to filter
+ "contains":["hostname"], //(Optional) Contains is an array for wildcard searches on a filter
+ "limit":"5" //(Optional) Limit the amount of results returned, default set to 50
+
+ """
+
+ payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name}
+ result = self._post("/v1/getexechistory", payload)
+ return result["response"]["exechistory"]
+
+
+"""
+from services.API import AirlockAPIWrapper
+
+
+api = AirlockAPIWrapper(base_url="https://airlock.example.com/api", api_key="your_api_key_here")
+
+#Example: Get all agents
+
+agents_df = api.agent_find_all()
+print("All Agents:")
+print(agents_df)
+"""
diff --git a/services/agenthandler.py b/services/agenthandler.py
new file mode 100644
index 0000000..725a9c5
--- /dev/null
+++ b/services/agenthandler.py
@@ -0,0 +1,232 @@
+# 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 .
+
+
+import json
+import logging
+import os
+import re
+from dataclasses import asdict
+from datetime import datetime, timedelta
+from typing import List
+
+import pandas as pd
+
+from models.agent import Agent
+from models.policy import Policy
+from services.API import AirlockAPIWrapper
+from services.selector import Selector
+from utils.utils import colorText, load_env, load_env_json
+
+logger = logging.getLogger(__name__)
+
+
+def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
+ agents = selectAgents(api)
+ history_days = Selector.select_value(
+ prompt="Enter how many days of history to pull (1–150): ",
+ value_type=int,
+ valid_range=(1, 150),
+ )
+
+ if not agents or not history_days:
+ print(colorText("No agents selected or invalid history range.", "red"))
+ return
+
+ historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
+ today = datetime.now().strftime("%Y-%m-%d")
+
+ all_history = []
+
+ for agent in agents:
+ try:
+ exechistory = api.history_execution(today, historical_date, agent.hostname)
+ except Exception as e:
+ print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))
+ continue
+
+ if isinstance(exechistory, list):
+ for block in exechistory:
+ record = {
+ "Command": block.get("commandline", "N/A"),
+ "Date": block.get("datetime", "N/A"),
+ "Filename": block.get("filename", "N/A"),
+ "Policy Name": block.get("policyname", "N/A"),
+ "Hostname": block.get("hostname", "N/A"),
+ "Hash": block.get("sha256", "N/A"),
+ }
+ all_history.append(record)
+
+ if not outputjson:
+ for key, value in record.items():
+ print(colorText(f"{key}: {value}", "green"))
+ print("\n")
+ else:
+ print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))
+
+ if outputjson:
+ print(json.dumps(all_history, indent=2))
+
+
+def findAllAgents(api):
+ # Step 1: Load data from API
+ policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
+ agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]
+
+ # Step 2: Create groupid → groupname map
+ groupid_to_name = {policy.groupid: policy.name for policy in policies}
+
+ # Step 3: Enrich agents
+ for agent in agents:
+ agent.enrich(groupid_to_name)
+
+ return agents
+
+
+
+def findAgents(api, return_dataframe):
+ agents = selectAgents(api)
+ working_dir = load_env("WORKING_DIR")
+ if agents:
+ agent_dicts = [asdict(agent) for agent in agents]
+ agent_df = pd.DataFrame(agent_dicts)
+
+ if return_dataframe:
+ logging.debug("Returning DataFrame to caller.")
+ return agent_df
+ else:
+ print(agent_df)
+ logging.debug("Displayed DataFrame to console.")
+
+ # Ask user if they want to export
+ user_input = input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
+ if user_input == 'y':
+ timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ filename = f"agentsearch_{timestamp}.csv"
+ file_path = os.path.join(working_dir, filename)
+
+ agent_df.to_csv(file_path, index=False)
+ logging.info(f"Exported DataFrame to {file_path}")
+
+ print(
+ colorText(
+ f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
+ "green",
+ )
+ )
+ else:
+ logging.debug("User declined to export the DataFrame.")
+ else:
+ logging.warning("No agents found.")
+ print("No agents matched the criteria.")
+
+def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
+ print(colorText("🔍 Device Search", "cyan"))
+ print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
+ print(colorText("When you're done, press Enter twice.\n", "cyan"))
+ print(colorText("Example:", "cyan"))
+ print(colorText("H00000", "cyan"))
+ print(colorText("UTN00000", "cyan"))
+ print(colorText("i-hSuperSecretServer", "cyan"))
+ print(colorText("u-hVenderBroke\n", "cyan"))
+
+ print(colorText("Paste or type your device names below:", "white"))
+
+ device_input_lines = []
+ empty_line_count = 0
+
+ while True:
+ line = input()
+ if line.strip() == "":
+ empty_line_count += 1
+ if empty_line_count == 2:
+ break
+ else:
+ empty_line_count = 0
+ device_input_lines.append(line.strip())
+
+ device_names = [name for name in device_input_lines if name]
+ if not device_names:
+ logger.debug("No device names entered")
+ print(colorText("⚠️ No device names entered.", "red"))
+ return []
+
+ # Build regex pattern
+ pattern = "|".join(map(re.escape, device_names))
+ regex = re.compile(pattern, re.IGNORECASE)
+
+ # Fetch agents
+ agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
+ matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
+ matched_agents.sort(key=lambda agent: agent.hostname.lower())
+
+ # Show unmatched
+ unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)]
+ if unmatched:
+ logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
+ print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
+
+ if not matched_agents:
+ logger.debug("❌ No matching devices found.")
+ print(colorText("❌ No matching devices found.", "red"))
+ else:
+ logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")
+ print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
+
+ return matched_agents
+
+def moveAgentToRelatedPolicy(
+ api: AirlockAPIWrapper,
+ agent: Agent,
+ mode: str = "audit",
+):
+ """
+ Moves an agent between audit and enforcement policies based on the mode.
+
+ Args:
+ api: AirlockAPIWrapper instance.
+ agent: Agent object.
+ policy_relationship_map: Dict mapping enforcement → audit.
+ mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
+ """
+ policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
+
+ if mode == "audit":
+ if agent.groupid in policy_relationship_map:
+ target_policy = policy_relationship_map[agent.groupid]
+ elif agent.groupid in policy_relationship_map.values():
+ logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.")
+ print(f"Agent {agent.hostname} is already in an audit group. No action needed.")
+ return
+ else:
+ logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
+ return
+
+ elif mode == "enforcement":
+ inverse_map = {v: k for k, v in policy_relationship_map.items()}
+ if agent.groupid in inverse_map:
+ target_policy = inverse_map[agent.groupid]
+ elif agent.groupid in inverse_map.values():
+ logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.")
+ return
+ else:
+ logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
+ return
+
+ else:
+ logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
+ return
+
+ api.agent_move(agent.agentid, target_policy)
diff --git a/services/policyhandler.py b/services/policyhandler.py
new file mode 100644
index 0000000..c1698d1
--- /dev/null
+++ b/services/policyhandler.py
@@ -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 .
+
+
+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")
diff --git a/utils/perstscheduler.py b/services/scheduler.py
similarity index 72%
rename from utils/perstscheduler.py
rename to services/scheduler.py
index 81e6762..a653c3e 100644
--- a/utils/perstscheduler.py
+++ b/services/scheduler.py
@@ -1,15 +1,30 @@
-
-#Standard Libary Imports:
+# 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 .
import json
+import logging
import os
import time
-from typing import Callable, Any, List, Dict
+from typing import Any, Callable, Dict, List
-#3rd Party Imports:
import schedule
-# File where all jobs are persisted
-JOBS_FILE = "scheduling\\jobs.json"
+logger = logging.getLogger(__name__)
+
+
+# TODO - move this File where all jobs are persisted
+JOBS_FILE = f"{os.getenv('WORKING_DIR')}\\scheduling\\jobs.json"
# Ensure directory exists
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)
@@ -21,6 +36,7 @@ FUNCTION_MAP: Dict[str, Callable] = {}
# Function Registration
# -------------------------------
+
def register_function(name: str, func: Callable):
"""
Register a function so it can be called by name later.
@@ -29,10 +45,12 @@ def register_function(name: str, func: Callable):
"""
FUNCTION_MAP[name] = func
+
# -------------------------------
# Persistence Helpers
# -------------------------------
+
def load_jobs() -> List[Dict[str, Any]]:
"""Load jobs from the JSON file, or return [] if none exist."""
if not os.path.exists(JOBS_FILE):
@@ -40,6 +58,7 @@ def load_jobs() -> List[Dict[str, Any]]:
with open(JOBS_FILE, "r") as f:
return json.load(f)
+
def _atomic_save(path: str, data: Any):
"""Write JSON atomically to avoid partial writes."""
tmp = f"{path}.tmp"
@@ -47,18 +66,22 @@ def _atomic_save(path: str, data: Any):
json.dump(data, f, indent=4)
os.replace(tmp, path)
+
def save_jobs(jobs: List[Dict[str, Any]]):
"""Save jobs to the JSON file (overwrite)."""
_atomic_save(JOBS_FILE, jobs)
+
# -------------------------------
# Uniqueness Helpers
# -------------------------------
+
def job_in_store(job_id: str) -> bool:
"""Check if a job id exists in the persisted JSON file."""
return any(j.get("id") == job_id for j in load_jobs())
+
def job_in_scheduler(job_id: str) -> bool:
"""
Check if a job with this tag exists in the in-memory scheduler.
@@ -71,6 +94,7 @@ def job_in_scheduler(job_id: str) -> bool:
# Fallback for older versions
return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)
+
def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
"""
Ensure the job_id is unique across persistence and in-memory schedule.
@@ -85,7 +109,7 @@ def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
if on_conflict == "error":
raise ValueError(f"Job id '{job_id}' already exists.")
elif on_conflict == "skip":
- print(f"[INFO] Job '{job_id}' already exists. Skipping creation.")
+ logger.info(f"Job '{job_id}' already exists. Skipping creation.")
return False
elif on_conflict == "replace":
# Clear from scheduler
@@ -97,10 +121,12 @@ def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
else:
raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")
+
# -------------------------------
# Internal scheduling (no persistence)
# -------------------------------
+
def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
args = args or []
kwargs = kwargs or {}
@@ -108,7 +134,7 @@ def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=No
def job_wrapper():
"""Executes the job once, then removes it."""
if func_name not in FUNCTION_MAP:
- print(f"[ERROR] Function '{func_name}' is not registered.")
+ logger.error(f"Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
# Remove from persistence
@@ -120,18 +146,21 @@ def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=No
delay_seconds = run_at_timestamp - time.time()
if delay_seconds <= 0:
- print(f"[WARN] Job {job_id} scheduled in the past. Skipping.")
+ logger.info(f"Job {job_id} scheduled in the past. Skipping.")
return
# Schedule via schedule library
schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)
-def _schedule_recurring(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None):
+
+def _schedule_recurring(
+ job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None
+):
args = args or []
kwargs = kwargs or {}
def job_wrapper():
if func_name not in FUNCTION_MAP:
- print(f"[ERROR] Function '{func_name}' is not registered.")
+ logger.error(f"Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
@@ -146,10 +175,12 @@ def _schedule_recurring(job_id: str, func_name: str, interval: int, unit: str, a
else:
raise ValueError(f"Unsupported unit: {unit}")
+
# -------------------------------
# Public APIs (with uniqueness + persistence)
# -------------------------------
+
def run_once_job(
job_id: str,
func_name: str,
@@ -169,22 +200,24 @@ def run_once_job(
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
- else:
- if job_in_scheduler(job_id):
- schedule.clear(job_id)
+ elif job_in_scheduler(job_id):
+ schedule.clear(job_id)
_schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
- jobs.append({
- "id": job_id,
- "type": "once",
- "run_at": run_at_timestamp,
- "function": func_name,
- "args": args or [],
- "kwargs": kwargs or {}
- })
+ jobs.append(
+ {
+ "id": job_id,
+ "type": "once",
+ "run_at": run_at_timestamp,
+ "function": func_name,
+ "args": args or [],
+ "kwargs": kwargs or {},
+ }
+ )
save_jobs(jobs)
+
def recurring_job(
job_id: str,
func_name: str,
@@ -205,21 +238,22 @@ def recurring_job(
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
- else:
- if job_in_scheduler(job_id):
- schedule.clear(job_id)
+ elif job_in_scheduler(job_id):
+ schedule.clear(job_id)
_schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
- jobs.append({
- "id": job_id,
- "type": "recurring",
- "interval": interval,
- "unit": unit,
- "function": func_name,
- "args": args or [],
- "kwargs": kwargs or {}
- })
+ jobs.append(
+ {
+ "id": job_id,
+ "type": "recurring",
+ "interval": interval,
+ "unit": unit,
+ "function": func_name,
+ "args": args or [],
+ "kwargs": kwargs or {},
+ }
+ )
save_jobs(jobs)
@@ -231,14 +265,14 @@ def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float
matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]
if not matched_jobs:
- print(f"[INFO] No jobs found containing PID substring '{pid_substring}'.")
+ logger.info(f"No jobs found containing PID substring '{pid_substring}'.")
return
- print(f"[INFO] Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
+ logger.info(f"Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
for job in matched_jobs:
job_id = job["id"]
- print(f" - Prioritizing job: {job_id}")
+ logger.debug(f" - Prioritizing job: {job_id}")
# Clear existing job from scheduler
schedule.clear(job_id)
@@ -252,7 +286,7 @@ def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float
job.get("args"),
job.get("kwargs"),
replace=True,
- persist=True
+ persist=True,
)
elif job["type"] == "recurring":
recurring_job(
@@ -263,16 +297,17 @@ def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float
job.get("args"),
job.get("kwargs"),
replace=True,
- persist=True
+ persist=True,
)
else:
- print(f"[WARN] Unknown job type for job '{job_id}'")
+ logger.warning(f"Unknown job type for job '{job_id}'")
# -------------------------------
# Reload Saved Jobs
# -------------------------------
+
def reload_jobs():
"""Reload jobs from JSON and reschedule them (no re-persist)."""
jobs = load_jobs()
@@ -280,21 +315,30 @@ def reload_jobs():
if job["type"] == "once":
if job["run_at"] > time.time():
run_once_job(
- job["id"], job["function"], job["run_at"],
- job.get("args"), job.get("kwargs"),
- persist=False
+ job["id"],
+ job["function"],
+ job["run_at"],
+ job.get("args"),
+ job.get("kwargs"),
+ persist=False,
)
elif job["type"] == "recurring":
recurring_job(
- job["id"], job["function"], job["interval"], job["unit"],
- job.get("args"), job.get("kwargs"),
- persist=False
+ job["id"],
+ job["function"],
+ job["interval"],
+ job["unit"],
+ job.get("args"),
+ job.get("kwargs"),
+ persist=False,
)
+
# -------------------------------
# Scheduler Loop
# -------------------------------
+
def start_scheduler():
"""
Start the scheduler loop (blocking).
@@ -305,4 +349,4 @@ def start_scheduler():
schedule.run_pending()
time.sleep(0.5)
except KeyboardInterrupt:
- print("[INFO] Scheduler stopped.")
+ logger.critical("Scheduler stopped.")
diff --git a/services/security.py b/services/security.py
new file mode 100644
index 0000000..2eb7de1
--- /dev/null
+++ b/services/security.py
@@ -0,0 +1,148 @@
+# 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 .
+
+import base64
+import logging
+import os
+import platform
+import re
+from getpass import getpass
+
+import keyring
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+# Constants
+KDF_ITERATIONS = 200_000
+SALT_SIZE = 16 # 128-bit Salt
+NONCE_SIZE = 12 # AES-GCM
+KEY_SIZE = 32 # AES-256
+
+
+def _derive_key(password: bytes, salt: bytes) -> bytes:
+ kdf = PBKDF2HMAC(
+ algorithm=hashes.SHA256(),
+ length=KEY_SIZE,
+ salt=salt,
+ iterations=KDF_ITERATIONS,
+ )
+ return kdf.derive(password)
+
+
+def configure_keyring_backend():
+ system = platform.system()
+ if system == "Windows":
+ import keyring.backends.Windows
+ keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring())
+ elif system == "Linux":
+ import keyring.backends.kwallet
+ keyring.set_keyring(keyring.backends.kwallet.DBusKeyring())
+ else:
+ raise EnvironmentError(f"Unsupported OS: {system}")
+
+
+def store_api_key(service: str, username: str, api_key: str, password: str):
+ configure_keyring_backend()
+ salt = os.urandom(SALT_SIZE)
+ key = _derive_key(password.encode(), salt)
+ aesgcm = AESGCM(key)
+ nonce = os.urandom(NONCE_SIZE)
+ ct = aesgcm.encrypt(nonce, api_key.encode(), associated_data=None)
+ blob = salt + nonce + ct
+ b64 = base64.b64encode(blob).decode()
+ keyring.set_password(service, username, b64)
+
+
+def retrieve_api_key(service: str, username: str, password: str) -> str:
+ configure_keyring_backend()
+ b64 = keyring.get_password(service, username)
+ if b64 is None:
+ raise ValueError("No stored secret for this service/username.")
+ blob = base64.b64decode(b64)
+ salt = blob[:SALT_SIZE]
+ nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE]
+ ct = blob[SALT_SIZE + NONCE_SIZE:]
+ key = _derive_key(password.encode(), salt)
+ aesgcm = AESGCM(key)
+ pt = aesgcm.decrypt(nonce, ct, associated_data=None)
+ return pt.decode()
+
+
+def api_key_exists(service: str, username: str) -> bool:
+ configure_keyring_backend()
+ return keyring.get_password(service, username) is not None
+
+
+def check_password_complexity(password: str) -> bool:
+ if len(password) < 12:
+ return False
+ if not re.search(r"[A-Z]", password):
+ return False
+ if not re.search(r"[a-z]", password):
+ return False
+ if not re.search(r"[0-9]", password):
+ return False
+ if not re.search(r"[^A-Za-z0-9]", password):
+ return False
+ return True
+
+
+def getAPI(USERNAME, SERVICE_NAME):
+ logging.debug(
+ f"Checking for stored API key for user '{USERNAME}' in service '{SERVICE_NAME}'..."
+ )
+
+ if api_key_exists(SERVICE_NAME, USERNAME):
+ for attempt in range(1, 4):
+ password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ")
+ try:
+ apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password)
+ logging.debug("API key successfully retrieved.")
+ return apikey
+ except Exception as e:
+ logging.warning(f"Attempt {attempt} failed: {str(e)}")
+ logging.error("Failed to retrieve API key after 3 incorrect attempts.")
+ raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
+ else:
+ logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.")
+ api_key = input(f"🔑 No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
+
+ while True:
+ password = getpass("🔐 Create a password to encrypt your API key: ")
+ if check_password_complexity(password):
+ try:
+ store_api_key(SERVICE_NAME, USERNAME, api_key, password)
+ logging.info("API key stored securely.")
+ break
+ except Exception as e:
+ logging.error(f"Failed to store API key: {e}")
+ break
+ else:
+ print("❌ Password does not meet complexity requirements. Try again.")
+ return api_key
+
+class APIKeyManager:
+ _api_key = None
+
+ @classmethod
+ def load(cls, service: str, username: str, password: str):
+ cls._api_key = retrieve_api_key(service, username, password)
+
+ @classmethod
+ def get(cls) -> str:
+ if cls._api_key is None:
+ raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
+ return cls._api_key
\ No newline at end of file
diff --git a/services/selector.py b/services/selector.py
new file mode 100644
index 0000000..47fe758
--- /dev/null
+++ b/services/selector.py
@@ -0,0 +1,121 @@
+# 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 .
+
+
+import logging
+from typing import Any, List, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+class Selector:
+ @staticmethod
+ def select_objects(
+ objects: List[Any],
+ allow_multiple: bool = False,
+ prompt_each: bool = False
+ ) -> Union[Optional[Any], List[Any]]:
+ if not objects:
+ logger.warning("No objects available for selection.")
+ return None
+
+ # Sort objects alphabetically by their 'name' attribute
+ sorted_objects = sorted(objects, key=lambda obj: getattr(obj, "name", str(obj)).lower())
+
+ # Display in 4 columns with extra spacing
+ num_columns = 4
+ rows = (len(sorted_objects) + num_columns - 1) // num_columns
+ print("\nAvailable Choices:")
+ for row in range(rows):
+ line = ""
+ for col in range(num_columns):
+ idx = row + col * rows
+ if idx < len(sorted_objects):
+ obj = sorted_objects[idx]
+ name = getattr(obj, "name", str(obj))
+ line += f"{idx + 1}: {name:<30}"
+ print(line)
+
+ selected = []
+
+ if allow_multiple:
+ while True:
+ choice = input("Select an object by number (or Q to finish): ").strip().lower()
+ if choice == "q":
+ break
+ try:
+ index = int(choice)
+ if 1 <= index <= len(sorted_objects):
+ obj = sorted_objects[index - 1]
+ if obj not in selected:
+ selected.append(obj)
+ if prompt_each:
+ logger.info(f"Selected: {getattr(obj, 'name', str(obj))}")
+ else:
+ logger.warning("Object already selected.")
+ else:
+ logger.warning("Selection out of range. Try again.")
+ except ValueError:
+ logger.warning("Invalid input. Enter a number or 'Q' to quit.")
+ return selected if selected else None
+ else:
+ try:
+ choice = int(input("Select one object by number: "))
+ if 1 <= choice <= len(sorted_objects):
+ selected_obj = sorted_objects[choice - 1]
+ logger.info(f"Selected: {getattr(selected_obj, 'name', str(selected_obj))}")
+ return selected_obj
+ else:
+ logger.warning("Selection out of range.")
+ except ValueError:
+ logger.warning("Invalid input.")
+ return None
+
+ @staticmethod
+ def select_value(
+ prompt: str,
+ value_type: type = int,
+ valid_range: Optional[tuple] = None,
+ allow_quit: bool = False
+ ) -> Optional[Any]:
+ while True:
+ user_input = input(prompt).strip().lower()
+ if allow_quit and user_input == "q":
+ logger.info("User opted to quit value selection.")
+ return None
+ try:
+ value = value_type(user_input)
+ if valid_range:
+ min_val, max_val = valid_range
+ if not (min_val <= value <= max_val):
+ logger.warning(f"Value out of range ({min_val}–{max_val}).")
+ continue
+ logger.info(f"User selected value: {value}")
+ return value
+ except ValueError:
+ logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
+
+ @staticmethod
+ def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
+ while True:
+ response = input(prompt).strip().lower()
+ if response in ["y", "yes"]:
+ logger.info("User confirmed action.")
+ return True
+ elif response in ["n", "no"]:
+ logger.info("User declined action.")
+ return False
+ else:
+ logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
\ No newline at end of file
diff --git a/services/setup.py b/services/setup.py
new file mode 100644
index 0000000..0b67e47
--- /dev/null
+++ b/services/setup.py
@@ -0,0 +1,178 @@
+# 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 .
+
+
+
+import json
+import logging
+import logging.handlers
+import os
+import platform
+import sys
+from pathlib import Path
+from dotenv import load_dotenv, set_key
+
+PROTECTED_KEYS = [
+ "APPNAME",
+ "PATH_EXCLUSION_CONST",
+ "MIN_FILES_FOR_PATH",
+ "VT_THREAT_TOLERANCE",
+ "POLICY_MAP_ENF_AUD"
+]
+
+def get_base_directory() -> Path:
+ system = platform.system()
+ home = Path.home()
+ if system == 'Windows':
+ return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools"
+ elif system == 'Darwin':
+ return home / 'Library' / 'Application Support' / "AirlockTools"
+ else:
+ return home / '.local' / 'share' / "AirlockTools"
+
+def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
+ log_file = log_dir / "airlocktools.log"
+ logger = logging.getLogger()
+ logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG))
+
+ if not logger.handlers:
+ file_handler = logging.handlers.RotatingFileHandler(
+ log_file, maxBytes=5_000_000, backupCount=5, encoding='utf-8'
+ )
+ file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+ logger.addHandler(file_handler)
+
+ console_handler = logging.StreamHandler()
+ console_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s'))
+ logger.addHandler(console_handler)
+
+ if platform.system() == "Windows":
+ try:
+ event_handler = logging.handlers.NTEventLogHandler("AirlockTools")
+ event_handler.setLevel(logging.CRITICAL)
+ event_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
+ logger.addHandler(event_handler)
+ except Exception as e:
+ logger.warning(f"Could not attach Windows Event Log handler: {e}")
+
+ logger.debug("Logging configured.")
+
+def get_system_config_path() -> Path:
+ base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__))))
+ return base_path.parent / "system_config.json"
+
+def load_system_config() -> dict:
+ try:
+ config_path = get_system_config_path()
+ with open(config_path, "r") as f:
+ return json.load(f)
+ except FileNotFoundError:
+ logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
+ return {
+ "APPNAME": "AirlockTools",
+ "LOG_LEVEL": "DEBUG",
+ "PATH_EXCLUSION_CONST": 4,
+ "MIN_FILES_FOR_PATH": 4,
+ "VT_THREAT_TOLERANCE": 4,
+ "POLICY_MAP_ENF_AUD": {
+ "enforced_id": "audit_id"
+ }
+ }
+
+def load_user_config(config_dir: Path) -> dict:
+ user_config_path = config_dir / "user_config.json"
+ if not user_config_path.exists():
+ default_user_config = {
+ "URL": "",
+ "LOG_LEVEL": "INFO"
+ }
+ with open(user_config_path, "w") as f:
+ json.dump(default_user_config, f, indent=4)
+ logging.debug(f"Created user config at {user_config_path}")
+ with open(user_config_path, "r") as f:
+ return json.load(f)
+
+def write_config_to_env(config: dict, env_path: Path):
+ for key, value in config.items():
+ try:
+ serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value)
+ set_key(env_path, key, serialized)
+ except Exception as e:
+ logging.warning(f"Failed to write {key} to .env: {e}")
+
+def setup() -> Path:
+ base_dir = get_base_directory()
+ dirs = {
+ 'config': base_dir / 'config',
+ 'cache': base_dir / 'cache',
+ 'logs': base_dir / 'logs',
+ }
+
+ for name, path in dirs.items():
+ path.mkdir(parents=True, exist_ok=True)
+ logging.debug(f"{name.capitalize()} directory ensured at: {path}")
+
+ system_config = load_system_config()
+ configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG"))
+
+ env_path = base_dir / ".env"
+ if not env_path.exists():
+ env_path.touch()
+ load_dotenv(dotenv_path=env_path, override=True)
+
+ working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
+ working_dir.mkdir(parents=True, exist_ok=True)
+ set_key(env_path, "WORKING_DIR", str(working_dir))
+ os.environ["WORKING_DIR"] = str(working_dir)
+ logging.debug(f"Working directory set to: {working_dir}")
+
+ folders_structure = {
+ "Approved": [],
+ "Needs_Review": ["Review_First", "Review_Second", "HTML"],
+ "Preflight": ["HTML"],
+ "Archived": ["HTML"],
+ "Scheduling": []
+ }
+
+ for folder_name, subfolders in folders_structure.items():
+ folder_path = working_dir / folder_name
+ folder_path.mkdir(parents=True, exist_ok=True)
+ logging.debug(f"'{folder_name}' folder ensured at: {folder_path}")
+ for subfolder in subfolders:
+ subfolder_path = folder_path / subfolder
+ subfolder_path.mkdir(parents=True, exist_ok=True)
+ logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
+
+ user_config = load_user_config(dirs['config'])
+ merged_config = {**system_config, **user_config}
+
+ for key in PROTECTED_KEYS:
+ merged_config[key] = system_config.get(key, "")
+
+ # ✅ URL resolution order: system_config → .env → user prompt
+ url = system_config.get("URL")
+ if not url:
+ url = os.getenv("URL")
+ if not url:
+ url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip()
+
+ merged_config["URL"] = url
+ set_key(env_path, "URL", url)
+ os.environ["URL"] = url
+ logging.debug(f"Service URL set to: {url}")
+
+ write_config_to_env(merged_config, env_path)
+
+ return working_dir
\ No newline at end of file
diff --git a/utils/clientfunctions.py b/utils/clientfunctions.py
deleted file mode 100644
index 3e73fe2..0000000
--- a/utils/clientfunctions.py
+++ /dev/null
@@ -1,430 +0,0 @@
-# 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 .
-
-#Local Imports
-import utils.hashfunctions as hashf
-import utils.utils as ct
-import utils.pathfunctions as pathf
-import utils.policyfunctions as policyf
-
-#Standard Libary Imports:
-import datetime
-import json
-import os
-import re
-
-#3rd Party Imports:
-import pandas as pd
-import requests
-
-
-
-def findAgentID(url):
-
- print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
- hostname = input(ct.colorText("Enter Device Name: ", "white"))
-
- endpoint = url + '/v1/agent/find'
- payload = {
- "hostname" : f"{hostname}"
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["agents"])
- return data.loc[0, "agentid"]
-
-def getDestAllowlistFromClientID(url, clientid):
-
- #This is dependant on their being an allowlist in the policy containing the name "localapproval"
-
- endpoint = url + '/v1/agent/find'
- payload = {
- "agentid" : f"{clientid}"
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["agents"])
-
- allowlist = policyf.getDestAllowlist(url,data.loc[0, "groupid"])
-
- return allowlist
-
-def getPolicyFromClientID(url, clientid):
-
- endpoint = url + '/v1/agent/find'
- payload = {
- "agentid" : f"{clientid}"
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["agents"])
-
- policy_name = getPolicyName(url,data.loc[0, "groupid"])
- policy_id = data.loc[0, "groupid"]
- return policy_name, policy_id
-
-def getPolicyName(url, groupid):
- endpoint = url + '/v1/group/'
- payload = {}
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
-
- data = pd.DataFrame(result["response"]["groups"])
-
- # Safely attempt to get the policy name
- filtered = data.loc[data['groupid'] == str(groupid), 'name']
- if not filtered.empty:
- name = filtered.values[0]
- else:
- name = None # or "Unknown", depending on your preference
-
- return name
-
-def devicehistory(url, outputjson: bool):
- endpoint = url + '/v1/getexechistory'
- print("\n")
- print(ct.colorText("1. Today", "yellow"))
- print(ct.colorText("2. Last 24 Hours", "yellow"))
- print(ct.colorText("3. Past 7 Days", "yellow"))
- print(ct.colorText("4. Past 30 Days", "yellow"))
- print(ct.colorText("5. Custom Date Range","yellow"))
- choice = input(ct.colorText("\nSelect Date Range: ", "white"))
- today = datetime.date.today()
- today = today.strftime("%Y-%m-%d")
- date_selected = " "
- if choice == '1':
- date_selected = today
- elif choice == '2':
- date_selected = datetime.date.today() - datetime.timedelta(days=1)
- date_selected = date_selected.strftime('%Y-%m-%d')
- elif choice == '3':
- date_selected = datetime.date.today() - datetime.timedelta(days=7)
- date_selected = date_selected.strftime('%Y-%m-%d')
- elif choice == '4':
- date_selected = datetime.date.today() - datetime.timedelta(days=30)
- date_selected = date_selected.strftime('%Y-%m-%d')
- elif choice == "5":
- print(ct.colorText("Please Input Dates as YYYY-MM-DD", "cyan"))
- date_selected = input(ct.colorText("From: ", "white"))
- today = input(ct.colorText("Date To: ", "white"))
- print(ct.colorText("WARNING: Device Name is Case Sensitive", "red"))
- device = input(ct.colorText("Enter Device Name: ", "white"))
- payload_dict = {
- "datefrom": date_selected,
- "dateto": today,
- "hostname": device
- }
- payload = json.dumps(payload_dict)
- print(ct.colorText(payload, "green"))
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
-
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
-
- if outputjson:
- return response
-
- parse_text = json.loads(response.text)
-
- # Safely get exechistory
- exechistory = parse_text.get('response', {}).get('exechistory')
-
- if isinstance(exechistory, list):
- for block in exechistory:
- print(ct.colorText(f"Command: {block.get('commandline', 'N/A')}", "green"))
- print(ct.colorText(f"Date: {block.get('datetime', 'N/A')}", "green"))
- print(ct.colorText(f"Filename: {block.get('filename', 'N/A')}", "green"))
- print(ct.colorText(f"Policy Name: {block.get('policyname', 'N/A')}", "green"))
- print(ct.colorText(f"Hostname: {block.get('hostname', 'N/A')}", "green"))
- print(ct.colorText(f"Hash: {block.get('sha256', 'N/A')}", "green"))
- print("\n")
- else:
- print(ct.colorText("No execution history found or data is not in expected format.", "red"))
-
-def findAllAgents(url):
- endpoint = url + '/v1/agent/find'
- payload = {}
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["agents"])
-
- group_ids = sorted(data['groupid'].unique().tolist())
- group_policy_map = {}
-
- for groupid in group_ids:
- policy_name = getPolicyName(url, groupid)
- group_policy_map[groupid] = policy_name
-
- data['policy_name'] = data['groupid'].map(group_policy_map)
-
- status_map = {
- 0: 'Offline',
- 1: 'Online',
- 2: 'Hidden',
- 3: 'Safemode'
- }
-
- data['status'] = data['status'].map(status_map)
- return(data)
-
-def findAgents(url, device_input_str, return_dataframe):
-
- os.makedirs("device_search", exist_ok=True)
- df = findAllAgents(url)
-
- #Parse the input string into device names (newline-separated only)
- device_names = device_input_str.strip().split('\n')
- device_names = [name.strip() for name in device_names if name.strip()]
-
- #Build a regex pattern for case-insensitive matching
- pattern = '|'.join([re.escape(name) for name in device_names])
- regex = re.compile(pattern, re.IGNORECASE)
-
- #Filter the DataFrame using regex
- matched_df = df[df['hostname'].apply(lambda x: bool(regex.search(str(x))))]
-
- if return_dataframe : return matched_df
- else:
- #Export to CSV with timestamp
- timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
- filename = f"agentsearch_{timestamp}.csv"
- matched_df.to_csv(f"device_search\\{filename}", index=False)
-
- print(ct.colorText(f"\n✅ Matched devices exported to: device_search\\{filename}","green"))
-
-def promptForDevices():
-
- print(ct.colorText("🔍 Device Search", "cyan"))
- print(ct.colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
- print(ct.colorText("When you're done, press Enter twice.\n", "cyan"))
- print(ct.colorText("Example:", "cyan"))
- print(ct.colorText("H00000", "cyan"))
- print(ct.colorText("UTN00000", "cyan"))
- print(ct.colorText("i-hSuperSecretServer", "cyan"))
- print(ct.colorText("u-hVenderBroke\n", "cyan"))
-
-
- # Collect multiline input from user
- print(ct.colorText("Paste or type your device names below:", "white"))
- device_input_lines = []
- while True:
- line = input()
- if line == "":
- break
- device_input_lines.append(line)
-
- device_input_str = "\n".join(device_input_lines)
-
- return device_input_str
-
-def findQuietAgents(url):
- # Get policy selection and agent list
- choice, policynames, policyid = policyf.choosePolicies(url)
- policy = policynames[choice]
- groupid = policyid[choice]
- agents = findGroupAgents(url, groupid)
-
- # Prompt user for history range
- while True:
- try:
- history_days = int(input("Enter how many days of history to pull (1–150): "))
- if 1 <= history_days <= 150:
- break
- else:
- print("Invalid input. Please enter a number between 1 and 150.")
- except ValueError:
- print("Invalid input. Please enter a valid integer.")
-
- while True:
- try:
- required_quiet = int(input("Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): "))
- if 1 <= required_quiet <= 365:
- break
- else:
- print("Invalid input. Please enter a number between 1 and 365.")
- except ValueError:
- print("Invalid input. Please enter a valid integer.")
-
- # Get policy execution history
- policy_exec_history = policyf.getPolicyInfo(url, policy, [1, 2, 6, 7], history_days, False)
-
- # Convert 'datetime' column to timezone-aware datetime objects
- policy_exec_history['datetime'] = pd.to_datetime(policy_exec_history['datetime'], format='%Y-%m-%dT%H:%M:%SZ', utc=True)
-
- # Get current UTC time
- now = datetime.datetime.now(datetime.timezone.utc)
-
- # Calculate days ago
- policy_exec_history['days_ago'] = policy_exec_history['datetime'].apply(lambda dt: (now - dt).days)
-
- # Count total executions per hostname
- hostname_counts = policy_exec_history['hostname'].value_counts()
-
- # Map execution counts to agents
- agents['execution_count'] = agents['hostname'].map(hostname_counts).fillna(0).astype(int)
-
- # Find most recent execution per hostname
- most_recent_exec = policy_exec_history.sort_values(by='days_ago').drop_duplicates(subset='hostname', keep='first')
-
- # Map most recent execution age to agents
- agents['days_since'] = agents['hostname'].map(most_recent_exec.set_index('hostname')['days_ago'])
-
- #Check for enforcement readyness
- agents['required_quiet'] = required_quiet
- agents['enforce_ready'] = agents['days_since'].apply(
- lambda x: True if pd.isna(x) or x > required_quiet else False
- )
-
- # Sort agents by execution count and hostname
- agents = agents.sort_values(by=['execution_count', 'hostname'], ascending=[True, True])
-
- # Save to CSV
- print(ct.colorText(f"Saving CSV to {policy}_agents_last_{history_days}_days.csv", "green"))
- agents.to_csv(f"{policy}_agents_last_{history_days}_days.csv", index=False)
-
-
- # Summary statistics
- total_agents = len(agents)
- ready_agents = agents['enforce_ready'].sum()
- not_ready_agents = total_agents - ready_agents
- ready_percentage = (ready_agents / total_agents) * 100
-
- # Print results
- print(f"Total agents: {total_agents}")
- print(f"Agents marked as 'enforce_ready': {ready_agents}")
- print(f"Agents not ready: {not_ready_agents}")
- print(f"Percentage ready for enforcement: {ready_percentage:.2f}%")
-
-def findGroupAgents(url, groupid):
- endpoint = url + '/v1/agent/find'
- payload = {
- "groupid" : f"{groupid}"
- }
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["agents"])
- group_ids = sorted(data['groupid'].unique().tolist())
- group_policy_map = {}
-
- for groupid in group_ids:
- policy_name = getPolicyName(url, groupid)
- group_policy_map[groupid] = policy_name
-
- data['policy_name'] = data['groupid'].map(group_policy_map)
-
- status_map = {
- 0: 'Offline',
- 1: 'Online',
- 2: 'Hidden',
- 3: 'Safemode'
- }
-
- data['status'] = data['status'].map(status_map)
- return(data)
-
-def moveAgentToAudit(url, agentid, policy_relationship_map):
- policy_name, policy_id = getPolicyFromClientID(url, agentid)
-
- print(policy_id)
-
-
- if policy_id in policy_relationship_map:
- target_policy = policy_relationship_map[policy_id]
- elif policy_id in policy_relationship_map.values():
- print(f"Agent {agentid} is already in an audit group. No action needed.")
- return
- else:
- print(f"Error: No corresponding audit policy found for policy: {policy_name} - {policy_id}.")
- return
-
- moveAgent(url, agentid, target_policy, "audit")
-
-def moveAgentToEnforcement(url, agentid, policy_relationship_map):
- policy_name, policy_id = getPolicyFromClientID(url, agentid)
-
- # Invert the map for audit → enforcement
- inverse_map = {v: k for k, v in policy_relationship_map.items()}
-
- if policy_id in inverse_map:
- target_policy = inverse_map[policy_id]
- elif policy_id in inverse_map.values():
- print(f"Agent {agentid} is already in an enforcement group. No action needed.")
- return
- else:
- print(f"Error: No corresponding enforcement policy found for policy: {policy_name} - {policy_id}.")
- return
-
- moveAgent(url, agentid, target_policy, "enforcement")
-
-def moveAgent(url, agentid, target_policy, direction):
- endpoint = f"{url}/v1/agent/move"
- payload = {
- "groupid": target_policy,
- "agentid": agentid
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- response = None # Initialize to avoid unbound errors
-
- try:
- response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)
- response.raise_for_status() # Raises HTTPError for bad status codes
-
- result = response.json()
-
- # Check if 'error' key exists and if it's not a success message
- if "error" in result and result["error"].lower() != "success":
- print(f"API returned an error: {result['error']}")
- else:
- print(f"✅ Agent {agentid} successfully moved to {direction} group {target_policy}.")
-
- except requests.exceptions.HTTPError as http_err:
- print(f"HTTP error occurred: {http_err}")
- if response is not None:
- print("Raw response:", response.text)
- except requests.exceptions.RequestException as req_err:
- print(f"Request error occurred: {req_err}")
- except ValueError:
- print("Failed to parse JSON response.")
- if response is not None:
- print("Raw response:", response.text)
- except Exception as e:
- print(f"Unexpected error: {e}")
- if response is not None:
- print("Raw response:", response.text)
\ No newline at end of file
diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py
deleted file mode 100644
index 831330f..0000000
--- a/utils/hashfunctions.py
+++ /dev/null
@@ -1,296 +0,0 @@
-# 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 .
-
-#Local Imports
-import utils.pathfunctions as pathf
-import utils.utils as ct
-
-#Standard Libary Imports:
-import gc
-import json
-import os
-
-#3rd Party Imports:
-import pandas as pd
-import requests
-
-def aggregateHashes(executions_json) -> pd.DataFrame:
- """
- Takes the executions, aggregates all the data with sha256 as primary, then returns aggregated dataframe
- """
- data = json.loads(executions_json)
- df = pd.DataFrame(data["response"]["exechistories"])
-
- if df.empty:
- return df
- print(df)
- # Aggregate by sha256, deduplicate lists, and preserve order
- agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()
-
- # Add a column for the number of unique hostnames
- agg_df["num_devices"] = agg_df["hostname"].apply(len)
-
- # Sort by num_devices in descending order
- agg_df = agg_df.sort_values("num_devices", ascending=False)
-
- return agg_df
-
-def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
- """
- Takes output of aggregatedHashes, queries API for those hashes, flattens response while keeping one row per hash,
- aggregate applications and baselines into lists, then merges results back into agg_df to create a
- """
- if 'sha256' not in agg_df.columns or agg_df.empty:
- print("⚠️ 'sha256' column missing or DataFrame is empty. Skipping API query.")
- return agg_df.copy() # Return as-is to avoid breaking downstream logic
-
- endpoint = url + '/v1/hash/query'
- payload = {
- "hashes": agg_df['sha256'].tolist()
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- data = response.json()
- results = data.get("response", {}).get("results", [])
-
- rows = []
- for res in results:
- row = {"sha256": res.get("sha256"), "result": res.get("result")}
-
- if "data" in res:
- d = res["data"]
- for key in ["filename", "filepath", "description", "filesize", "md5",
- "productname", "productversion", "publisher", "createtime", "modtime",
- "sha128", "sha384", "sha512", "datetime"]:
- row[key] = d.get(key)
-
- row["applications"] = d.get("applications", [])
- row["baselines"] = d.get("baselines", [])
-
- reputation = d.get("reputation", {})
- for k, v in reputation.items():
- row[f"reputation_{k}"] = v
-
- rows.append(row)
-
- df_api = pd.DataFrame(rows)
-
- if 'sha256' not in df_api.columns:
- print("⚠️ API response missing 'sha256'. Skipping merge.")
- return agg_df.copy()
-
- df = agg_df.merge(df_api, on="sha256", how="left")
-
- # Only include columns that exist to avoid KeyErrors
- expected_columns = ['policy','sha256', 'filename_x', 'description', 'productname', 'productversion',
- 'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username',
- 'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
- 'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
- 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']
-
- available_columns = [col for col in expected_columns if col in df.columns]
- aug_df = df[available_columns]
-
- return aug_df
-
-def categorizeHashes(df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
- if untrusted_publishers is None: untrusted_publishers = []
- if pups is None: pups = []
-
- def reputationtool(row):
- val = row["reputation_scannermatch"]
- if pd.isna(val) or val == "N/A":
- return row["publisher"] == "Not Signed"
- try:
- return int(val) > threat_tolerance
- except (ValueError, TypeError):
- return row["publisher"] == "Not Signed"
-
- df["reputation_flag"] = df.apply(reputationtool, axis=1)
-
- mask_needsreview = (
- ((df["publisher"] == "Not Signed") & df["reputation_flag"]) |
- (df["reputation_status"] == "UNKNOWN")
- )
-
- mask_approved = (
- (
- (df["publisher"] != "Not Signed") &
- ~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
- ~df["reputation_status"].isna() &
- ~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
- ) |
- (
- (df["publisher"] == "Not Signed") &
- ~df["reputation_flag"] &
- ~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) &
- ~df["reputation_status"].isna() &
- ~df["description"].str.contains(pathf.regulator(pups), case=False, na=False)
- )
- )
-
- needsreview_df = df[mask_needsreview]
- approved_df = df[mask_approved]
- unapproved_df = df[~(mask_needsreview | mask_approved)]
- return needsreview_df, approved_df, unapproved_df
-
-def combineHashAndHist(hash_path, condensed_path):
- # Load both datasets
- condensed_combo = pd.read_parquet(condensed_path)
- df = pd.read_parquet(hash_path)
- # Merge on sha256
- df = pd.merge(condensed_combo, df, on='sha256', how='inner')
- # Rename and reorder columns
- df = df.rename(columns={'publisher_x': 'publisher'})
- df = df.rename(columns={'policy_x': 'policy'})
- df = df[['policy','sha256', 'publisher', 'description', 'filename', 'hostname', 'username',
- 'productname', 'productversion', 'reputation_lastseen', 'reputation_scannermatch',
- 'reputation_scannercount', 'reputation_status', 'reputation_threatlevel',
- 'reputation_threatname', 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline']]
- df = df.sort_values(by='filename')
-
- # Overwrite the original hash file
- df.to_parquet(hash_path, index=False)
-
- # Cleanup
- del df
- del condensed_combo
- gc.collect()
-
-def combineHashes(url, parquet_files) -> pd.DataFrame:
- combined_hashes = pd.DataFrame()
- hashes = []
-
- for file_path in parquet_files:
- try:
- hash_df = pd.read_parquet(file_path)
- pathf.inspect_parquet(file_path)
-
- if not hash_df.empty:
- hashes.append(hash_df)
- else:
- print(f"⚠️ Dataframe is empty: {file_path}")
- except Exception as e:
- print(f"❌ Error reading Parquet file '{file_path}': {e}")
-
- if hashes:
- combined_hashes = pd.concat(hashes, ignore_index=True)
- print(f"✅ Combined {len(combined_hashes)} hashes from {len(hashes)} files.")
- else:
- print("⚠️ No valid dataframes to combine.")
-
- combined_hashes = combined_hashes.drop_duplicates(subset=['sha256'])
- augmented_combo = augmentAggregatedHashes(url, combined_hashes)
-
- numeric_reputation_cols = [
- 'reputation_scannermatch',
- 'reputation_scannercount',
- 'reputation_threatlevel'
- ]
-
- for col in numeric_reputation_cols:
- if col in augmented_combo.columns:
- augmented_combo[col] = pd.to_numeric(augmented_combo[col].replace('N/A', pd.NA), errors='coerce')
-
- augmented_combo = augmented_combo.rename(columns={'publisher_x': 'publisher'})
- augmented_combo = augmented_combo[['sha256', 'publisher', 'description', 'productname', 'productversion',
- 'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
- 'reputation_status', 'reputation_threatlevel', 'reputation_threatname',
- 'reputation_timestamp']]
- augmented_combo = augmented_combo.sort_values(by=['publisher', 'description', 'productname'])
- del combined_hashes
- gc.collect()
- print(ct.colorText("Hash reputation info added to dataframe", "green"))
- return augmented_combo
-
-def condenseExecutions(parquet_paths):
- combined_df = pd.DataFrame()
- valid_files = []
-
- for file_path in parquet_paths:
- try:
- df = pd.read_parquet(file_path)
- # Optional: pathf.inspect_parquet(file_path)
- if not df.empty:
- combined_df = pd.concat([combined_df, df], ignore_index=True)
- valid_files.append(file_path)
- print(f"✅ Loaded {len(df)} rows from {file_path}")
- else:
- print(f"⚠️ DataFrame from '{file_path}' is empty.")
- except Exception as e:
- print(f"❌ Error reading Parquet file '{file_path}': {e}")
-
- if not combined_df.empty:
- print(f"✅ Combined {len(combined_df)} rows from {len(valid_files)} files.")
- else:
- print("⚠️ No valid dataframes to combine.")
-
- return combined_df
-
-def divideSortedHashExecutions(unknown_parq, good_parq, bad_parq, condensed_parq, pups) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
- # Run combineHashAndHist on each file
- combineHashAndHist(unknown_parq, condensed_parq)
- combineHashAndHist(good_parq, condensed_parq)
- combineHashAndHist(bad_parq, condensed_parq)
-
- # Load data
- unknown = pd.read_parquet(unknown_parq)
- good = pd.read_parquet(good_parq)
- bad = pd.read_parquet(bad_parq)
-
- # Build regex pattern once
- pattern = pathf.regulator(pups)
-
- # Move matching rows from unknown and good to bad
- bad = pd.concat([
- bad,
- unknown[unknown["filename"].str.contains(pattern, na=False)],
- good[good["filename"].str.contains(pattern, na=False)]
- ], ignore_index=True)
-
- # Remove matching rows from unknown and good
- unknown = unknown[~unknown["filename"].str.contains(pattern, na=False)]
- good = good[~good["filename"].str.contains(pattern, na=False)]
-
- return unknown, good, bad
-
-def generatePreflights(hashes, primary_path, secondary_path):
- all_hashes = pd.read_parquet(hashes)
-
- primarypathexclusions = ct.tryToReadCSV(primary_path)
- secondarypathexclusions = ct.tryToReadCSV(secondary_path)
-
- pathexclusions = pd.concat([primarypathexclusions, secondarypathexclusions], ignore_index=True)
-
- allowbyhash = all_hashes[~all_hashes['sha256'].isin(pathexclusions['sha256'])]
-
- allowbyhash.sort_values(by=["filename"])
-
- return pathexclusions, allowbyhash
-
-def generatePublist(all_hashes, bad_publisher_list):
- all_approved_hashes = ct.tryToReadParquet(all_hashes)
-
- #Drop all not signed, only keep unique values
- publist = all_approved_hashes[all_approved_hashes['publisher'] != "Not Signed"].drop_duplicates(subset=['publisher'])
- #Remove Bad publisher if somehow they made it this far
- pattern = pathf.regulator(bad_publisher_list)
- publist = publist[~publist["publisher"].str.contains(pattern, na=False)]
-
- return publist
-
\ No newline at end of file
diff --git a/utils/localapproval.py b/utils/localapproval.py
deleted file mode 100644
index da4ce75..0000000
--- a/utils/localapproval.py
+++ /dev/null
@@ -1,326 +0,0 @@
-# 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 .
-
-#Local Imports
-import utils.clientfunctions as clientf
-import utils.pathfunctions as pathf
-import utils.policyfunctions as policyf
-import utils.utils as ct
-from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid
-
-#Standard Libary Imports:
-import json
-import os
-import re
-import time
-import datetime
-
-#3rd Party Imports:
-import pandas as pd
-import numpy as np
-import requests
-
-def getLocalApprovals(url):
-
- endpoint = url + f'/v1/otp/usage'
- payload = {
- "status" : "0"
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- local_approval = pd.DataFrame(result["response"]["otpusage"])
- if os.path.exists("Local_Approval\\PARQ\\newest_local_approval.parquet"):
- previous_run = pd.read_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet")
- previous_run.to_parquet("Local_Approval\\PARQ\\last_local_approval.parquet", index=False)
- os.remove("Local_Approval\\PARQ\\newest_local_approval.parquet")
-
- #Only keep rows presumably created by the generate local approval function
- local_approval = local_approval[local_approval['purpose'].str.startswith('🎫 Local Approval 🎫')]
-
- local_approval['batchid'] = local_approval['purpose'].apply(lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1))
-
- if not local_approval.empty:
- ct.style_dataframe_dark(local_approval, f"Local_Approval\\HTML\\newest_local_approval.html")
- local_approval.to_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet", index=False)
-
- return local_approval
-
-def scheduleAddingLAHashes(url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant):
- try:
- register_function("add_hash", returnFromLocalApproval)
- register_function("move_device", clientf.moveAgentToEnforcement)
- except Exception as e:
- print(f"[ERROR] Failed to register functions: {e}")
- return
-
- try:
- approvals_df = getNewLocalApprovals(url)
- if approvals_df.empty:
- print("[INFO] No new local approvals found. Nothing to schedule.")
- return
- batches = approvals_df.groupby('batchid')
- except Exception as e:
- print(f"[ERROR] Failed to retrieve or group local approvals: {e}")
- return
-
- for batchid, batch_df in batches:
- try:
- duration_minutes = int(batch_df['duration'].iloc[0])
- start_time = datetime.datetime.now()
- run_time = start_time + datetime.timedelta(minutes=duration_minutes)
- early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))
-
- early_timestamp = early_time.timestamp()
- run_timestamp = run_time.timestamp()
-
- # Schedule add_hash job
- try:
- run_once_job(
- f"add_hash_{batchid}",
- "add_hash",
- early_timestamp,
- [url, batch_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant],
- None
- )
- print(f"[INFO] Scheduled add_hash for batch {batchid} at {early_time}")
- except Exception as e:
- print(f"[ERROR] Failed to schedule add_hash for batch {batchid}: {e}")
-
- # Schedule move_device jobs
- devices = batch_df['agentid'].drop_duplicates().tolist()
- for device in devices:
- try:
- run_once_job(
- f"move_device_{device}_{batchid}",
- "move_device",
- run_timestamp,
- [url, device, policy_relationship_map],
- None
- )
- print(f"[INFO] Scheduled move_device for device {device} in batch {batchid} at {run_time}")
- except Exception as e:
- print(f"[ERROR] Failed to schedule move_device for device {device} in batch {batchid}: {e}")
-
- except Exception as e:
- print(f"[ERROR] Failed to process batch {batchid}: {e}")
-
-def returnFromLocalApproval(url, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant):
- # Get unique policy names from device list
- policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())
-
- # Create inverse map to go from Audit to Enforcement
- inverse_map = {v: k for k, v in policy_relationship_map.items()}
-
- # Fetch all policies
- all_policies = policyf.getPolicyDataframe(url)
-
- # Define policy types
- policy_types = [1, 2, 6, 7]
-
- # Define output directories
- parq_base_dir = "Local_Approval\\PARQ\\"
- needappr_base_dir = "Local_Approval\\PARQ\\"
-
- # Ensure output directories exist
- os.makedirs(parq_base_dir, exist_ok=True)
- os.makedirs(needappr_base_dir, exist_ok=True)
-
- # Build execution history
- policyf.buildExecHistory(
- url,
- policies_in_devicelist,
- parq_base_dir,
- needappr_base_dir,
- policy_types,
- threat_tolerance_constant,
- bad_publisher_list,
- pups,
- )
-
- # Read hash data
- unknown_hashes = ct.tryToReadCSV(f"{needappr_base_dir}unknown_hashes.csv")
- good_hashes = ct.tryToReadCSV(f"{needappr_base_dir}good_hashes.csv")
- hashes = pd.concat([unknown_hashes, good_hashes], ignore_index=True)
-
- # Process each policy
- for policy in policies_in_devicelist:
- # Filter for matching policy
- matching_rows = all_policies[all_policies['name'] == policy]
-
- if matching_rows.empty:
- print(f"Warning: No group ID found for policy '{policy}'. Skipping.")
- continue
-
- # Extract group ID
- policy_id = matching_rows['groupid'].values[0]
-
- # Map to destination ID
- destination_id = inverse_map.get(policy_id)
- if destination_id is None:
- print(f"Warning: No corresponding enforcement policy found for group ID '{policy_id}'. Skipping.")
- continue
-
- allowlist = policyf.getDestAllowlist(url, destination_id)
-
- policyf.addHash(url, allowlist, hashes[hashes['group'] == policy_id])
-
-def moveToLocalApproval(url, policy_relationship_map):
-
- possible_durations = [15, 60, 360, 1440, 10080]
- duration_selected = None
-
- print(ct.colorText("Please select a duration:", "white"))
- for i, option in enumerate(possible_durations, start=1):
- print(f"{i}. {option}")
-
- try:
- choice = int(input("Enter the number of your choice: "))
- if 1 <= choice <= len(possible_durations):
- duration_selected = possible_durations[choice - 1]
- print(ct.colorText(f"You selected: {duration_selected}", "yellow"))
- else:
- print(ct.colorText("❌ Invalid choice.", "red"))
- return
- except ValueError:
- print(ct.colorText("❌ Invalid input. Please enter a number.", "red"))
- return
-
- devicelist = clientf.promptForDevices()
- device_df = clientf.findAgents(url, devicelist, True)
-
- batch = int(time.time())
-
- if device_df is None or device_df.empty:
- print(ct.colorText("❌ No agents found or error retrieving agents.", "red"))
- return
-
- for row in device_df.itertuples(index=False):
- try:
- addLocalApproval(url, batch, duration_selected, row.agentid)
- clientf.moveAgentToAudit(url, row.agentid, policy_relationship_map)
- except Exception as e:
- print(ct.colorText(f"❌ Error processing agent {row.agentid}: {e}", "red"))
-
-def addLocalApproval(url, batchid, duration_selected, agentid):
- purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
- endpoint = url + '/v1/otp/retrieve'
- payload = {
- "duration": str(duration_selected),
- "agentid": str(agentid),
- "purpose": purpose
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)
-
- try:
- result = response.json()
- otpcode = result["response"]["otpcode"]
- print(ct.colorText(f"The OTP code is: {otpcode}", "yellow"))
- except Exception as e:
- print(ct.colorText(f"An unexpected error occurred for agent {agentid}: {str(e)}", "red"))
-
-def monitorAuditStatus(url: str, policy_relationship_map: dict):
-
- # Simulated current agent list
- current_agent_list = clientf.findAllAgents(url)
-
- # Load old agent list
- old_agent_path = "Local_Approval\\PARQ\\last_agent_list.parquet"
- if os.path.exists(old_agent_path):
- old_agent_list = pd.read_parquet(old_agent_path)
- else:
- old_agent_list = pd.DataFrame(columns=current_agent_list.columns)
-
- # Merge on hostname
- merged = pd.merge(
- old_agent_list[['hostname', 'groupid']],
- current_agent_list[['hostname', 'groupid']],
- on='hostname',
- how='outer',
- suffixes=('_old', '_current'),
- indicator=True
- )
-
- # Reverse map for enforcement
- reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
- known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())
-
- # 1. Newly added
- newly_added = merged[merged['_merge'] == 'right_only']
-
- # 2. Same policy
- same_policy = merged[
- (merged['_merge'] == 'both') &
- (merged['groupid_old'] == merged['groupid_current'])
- ]
-
- # 3. Moved to audit
- moved_to_audit = merged[
- (merged['_merge'] == 'both') &
- (merged['groupid_old'].isin(policy_relationship_map)) &
- (merged['groupid_current'] == merged['groupid_old'].map(policy_relationship_map))
- ]
-
- # 4. Moved to enforcement
- moved_to_enforcement = merged[
- (merged['_merge'] == 'both') &
- (merged['groupid_old'].isin(reverse_policy_map)) &
- (merged['groupid_current'] == merged['groupid_old'].map(reverse_policy_map))
- ]
-
- # 5. Unusual moves
- unusual_move = merged[
- (merged['_merge'] == 'both') &
- (merged['groupid_old'] != merged['groupid_current']) &
- merged.apply(lambda row: (row['groupid_old'], row['groupid_current']) not in known_transitions, axis=1)
- ]
-
- current_agent_list.to_parquet("Local_Approval\\PARQ\\last_agent_list.parquet", index=False)
-
- # Return all five DataFrames
- return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move
-
-def getNewLocalApprovals(url):
- current_la = getLocalApprovals(url)
-
- # Load old approval list
- old_la_path = "Local_Approval\\PARQ\\last_la.parquet"
- if os.path.exists(old_la_path):
- old_la = pd.read_parquet(old_la_path)
- else:
- old_la = pd.DataFrame(columns=current_la.columns)
-
- # Create composite keys
- current_la['key'] = current_la['clientid'].astype(str) + "_" + current_la['granted'].astype(str)
- old_la['key'] = old_la['clientid'].astype(str) + "_" + old_la['granted'].astype(str)
-
- # Find new entries
- new_entries = current_la[~current_la['key'].isin(old_la['key'])]
-
- # Convert 'granted' to datetime and filter by last 10 minutes
- new_entries['granted'] = pd.to_datetime(new_entries['granted'], utc=True, errors='coerce')
- ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
- recent_entries = new_entries[new_entries['granted'] > ten_minutes_ago]
-
- # Save current approvals for next run
- current_la.drop(columns=['key'], inplace=True)
- current_la.to_parquet(old_la_path, index=False)
-
- return recent_entries
\ No newline at end of file
diff --git a/utils/menus.py b/utils/menus.py
new file mode 100644
index 0000000..6edfa76
--- /dev/null
+++ b/utils/menus.py
@@ -0,0 +1,258 @@
+# 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 .
+
+import logging
+import os
+import re
+
+import dotenv
+import pandas as pd
+
+import flows.localApproval as la
+import services.policyhandler as policyh
+from flows.prepPolicy import (
+ buildPathsandPublishers,
+ buildPreflights,
+ selectAllowlists,
+ selectPolicies,
+ sortHashes,
+)
+from flows.quietAgent import findQuietAgents
+from services.agenthandler import findAgents
+from services.API import AirlockAPIWrapper
+from utils.utils import (
+ areYouSure,
+ colorText,
+ displayIntro,
+ load_env,
+ open_directory,
+ printEnforceChecklist,
+)
+
+logger = logging.getLogger(__name__)
+
+dotenv.load_dotenv()
+
+def menu_main(api: AirlockAPIWrapper):
+ working_dir = load_env("WORKING_DIR")
+ while True:
+ displayIntro()
+ # Add Settings, and give option to change working dir
+ print(colorText("1. ➡️ - Move Device(s) to local approval", "yellow"))
+ print(colorText("2. 🎫 - OTP", "yellow"))
+ print(colorText("3. 🔍 - Device Search", "yellow"))
+ print(colorText("4. 🔇 - Find Quiet Hosts", "yellow"))
+ print(colorText("5. 🔒 - Prepare Policy For Enforcement", "yellow"))
+ print(colorText("6. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
+ print(colorText("F. 📂 - Open Working Directory", "yellow"))
+ print(colorText("S. 🛠️ - Settings", "yellow"))
+ print(colorText("Q. 🔚 - Quit", "yellow"))
+
+ choice = input(colorText("\nEnter Menu Item: ", "white"))
+ if choice == "1":
+ la.moveToLocalApproval(api)
+ elif choice == "2":
+ menu_otp(api)
+ elif choice == "3":
+
+ findAgents(api,False)
+ elif choice == "4":
+ findQuietAgents(api)
+ elif choice == "5":
+ menu_policy_enforce(api)
+ elif choice == "6":
+ areYouSure()
+ confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
+ if confirmation.strip().upper() == "I AGREE":
+ policyh.updateAuditPoliciesFromEnforcementPolices(api)
+
+ elif choice == "F":
+ open_directory(working_dir)
+ elif choice == "S":
+ menu_settings()
+ elif choice == "Q":
+ break
+ else:
+ print(colorText("Invalid choice. Please try again.", "red"))
+
+
+def menu_policy_enforce(api: AirlockAPIWrapper):
+ selected_policies = []
+ destination_policy = []
+ destination_allowlist = []
+ processed_paths = []
+ processed_hashes = []
+ processed_publishers = []
+ tested = False
+ working_dir = load_env("WORKING_DIR")
+
+ while True:
+ printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
+
+ choice = input(colorText("\nEnter your choice: ", "white"))
+
+ if choice == "1":
+ selected_policies = selectPolicies(api,True)
+
+ elif choice == "2":
+ print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
+
+ destination_policy = selectPolicies(api, False)
+
+ print(colorText("Please choose Allowlist for Hashes", "white"))
+
+ destination_allowlist = selectAllowlists(api, False)
+
+ elif choice == "3":
+ sortHashes(
+ api,
+ selected_policies,
+ type=[1, 2, 6, 7],
+ )
+
+ elif choice == "4":
+ if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
+ buildPathsandPublishers(False)
+ else:
+ print("File not found. Please make sure it's saved correctly and try again.")
+
+ elif choice == "5":
+ if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
+ f"{working_dir}\\Approved\\primary_Paths.csv"
+ ):
+ buildPreflights()
+ else:
+ print("File not found. Please make sure it's saved correctly and try again.")
+
+ elif choice == "6":
+ if (
+ os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv")
+ and os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv")
+ and destination_policy
+ and destination_allowlist
+ ):
+ print(colorText("These path exclusions would be added to:", "yellow"))
+ print(destination_policy)
+
+ pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\approved_paths.csv")
+ hashes = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
+
+ # Get unique combinations of longestcfp and file_extension
+ unique_combinations = pathexclusions[
+ ["longestcfp", "file_extension"]
+ ].drop_duplicates()
+
+ # Regex to match a Windows drive letter at the start (e.g., C:\)
+ drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
+
+ # Build processed paths like \\path\\**.exe or C:\path\**.jar
+ processed_paths = [
+ (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
+ for path, ext in unique_combinations.itertuples(index=False, name=None)
+ ]
+
+ print(processed_paths)
+
+ print(colorText("These publishers would added", "yellow"))
+
+ if os.path.exists(f"{working_dir}\\Preflight\\approved_publishers.csv"):
+ publishers = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
+ if publishers.empty:
+ print(colorText("The publishers list is empty.", "red"))
+ else:
+ processed_publishers = (
+ publishers[publishers["publisher_hash"] != "Not Signed"]
+ ["publisher_hash"]
+ .drop_duplicates()
+ .tolist()
+ )
+
+ print(processed_publishers)
+
+ print(colorText("These hashes would be added to:", "yellow"))
+ print(destination_allowlist)
+
+ processed_hashes = hashes["sha256"].unique().tolist()
+ print(processed_hashes)
+
+ if processed_paths and processed_hashes:
+ tested = True
+
+ elif choice == "7":
+ areYouSure()
+ confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
+
+ if (
+ tested
+ and destination_policy
+ and destination_allowlist
+ and confirmation.strip().upper() == "I AGREE"
+ ):
+ print(colorText("Proceeding with the code...", "yellow"))
+ api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
+ api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
+ if processed_publishers:
+ api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
+
+
+ elif choice == "F":
+ open_directory(working_dir)
+ elif choice == "S":
+ menu_settings()
+ elif choice == "Q":
+ break
+
+
+ else:
+ print(colorText("Invalid choice. Please try again.", "red"))
+
+
+def menu_otp(api: AirlockAPIWrapper):
+ while True:
+ print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
+ print(colorText("1. Generate OTP", "cyan"))
+ # print(colorText("2. Sub-option B","cyan"))
+ print(colorText("Q. Return to Main Menu", "cyan"))
+ choice = input("Enter your choice: ")
+
+ if choice == "1":
+ # TODO generateOTP(api,findAgents()
+ break
+
+ elif choice == "2":
+ print("You selected Sub-option B")
+ elif choice == "Q":
+ print("Returning to Main Menu...")
+ break
+ else:
+ print("Invalid choice. Please try again.")
+
+
+def menu_settings():
+ while True:
+ print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
+ print(colorText("1. Change Working Dir", "cyan"))
+ # print(colorText("2. Sub-option B","cyan"))
+ print(colorText("Q. Return to Main Menu", "cyan"))
+ choice = input("Enter your choice: ")
+
+ if choice == "1":
+ pass #TODO ADD CHANGE WORKDIR CODE
+
+ elif choice == "Q":
+ print("Returning to Main Menu...")
+ break
+ else:
+ print("Invalid choice. Please try again.")
diff --git a/utils/pathfunctions.py b/utils/pathfunctions.py
deleted file mode 100644
index e347bcb..0000000
--- a/utils/pathfunctions.py
+++ /dev/null
@@ -1,213 +0,0 @@
-# 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 .
-
-#Local Imports
-import utils.hashfunctions as hashf
-import utils.pathfunctions as pathf
-import utils.utils as ct
-
-#Standard Libary Imports:
-
-import os
-import re
-
-#3rd Party Imports:
-import pandas as pd
-
-
-def split_filepaths_grouped(df, col="filename", group_parts=4, min_parts=4):
- def clean_split(path):
- parts = os.path.normpath(path).split(os.sep)
- # Remove leading empty strings caused by UNC paths
- parts = [p for p in parts if p]
- return parts
-
- df = df.copy()
- split_paths = df[col].apply(clean_split)
-
- # Filter out paths with fewer than `min_parts` components
- df = df[split_paths.apply(lambda parts: len(parts) >= min_parts)].copy()
- split_paths = split_paths[df.index] # Update split_paths to match filtered df
-
- df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:group_parts]))
- grouped = df.groupby("group_key")
- new_rows = []
-
- for _, group_df in grouped:
- paths = group_df[col].tolist()
- split_parts = [clean_split(p) for p in paths]
-
- def longest_common_prefix(paths):
- if not paths:
- return []
- prefix = paths[0]
- for path in paths[1:]:
- prefix = [a for a, b in zip(prefix, path) if a == b]
- if not prefix:
- break
- return prefix
-
- common_prefix = longest_common_prefix(split_parts)
- prefix_str = os.sep.join(common_prefix)
-
- for i, parts in enumerate(split_parts):
- filename = parts[-1]
- middle = os.sep.join(parts[len(common_prefix):-1]) if len(parts) > len(common_prefix) + 1 else ""
- row = group_df.iloc[i].copy()
- row["longestcfp"] = prefix_str
- row["middle"] = middle
- row["filename_only"] = filename
- row["file_extension"] = os.path.splitext(filename)[1].lower()
- new_rows.append(row)
-
- return pd.DataFrame(new_rows).drop(columns=["group_key"])
-
-def inspect_parquet(path):
- try:
- df = pd.read_parquet(path)
- print(f"✅ Successfully read: {path}")
- print(f"📄 Columns: {df.columns.tolist()}")
- print(f"🔢 Rows: {len(df)}")
- return df
- except Exception as e:
- print(f"❌ Error reading {path}: {e}")
- return pd.DataFrame()
-
-def regulator(paths, case_insensitive=True):
- """
- Build a regex pattern that matches any of the given Windows path fragments.
- """
- escaped = [re.escape(p) for p in paths]
- pattern = "(?:" + "|".join(escaped) + ")"
- if case_insensitive:
- pattern = "(?i)" + pattern # Add inline case-insensitive flag
- print(f"Regulator is providing: {pattern}")
- return pattern
-
-def calculatePath(approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split):
-
- if split : dfs_by_policy = [group for _, group in approved_hashes.groupby('policy')]
- else : dfs_by_policy = [approved_hashes]
-
- processed_dfs = []
-
- for df in dfs_by_policy:
- haslcp = pathf.split_filepaths_grouped(df, "filename", path_exclusion_constant, min_files_for_path)
- haslcp = haslcp.drop_duplicates()
-
- forbidden = pathf.regulator(badpathparts, True)
- forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)
-
- print(ct.colorText("Removing forbidden filepaths for path exceptions", "green"))
- lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()
-
- lcp_not_forbidden_review = lcp_not_forbidden[['policy', 'longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']]
-
- unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index()
- unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count']
-
- lcp_not_forbidden_review = lcp_not_forbidden_review.merge(unique_sha_counts, on='longestcfp', how='left')
- lcp_not_forbidden_review = lcp_not_forbidden_review[lcp_not_forbidden_review['unique_sha256_count'] >= min_files_for_path]
- processed_dfs.append(lcp_not_forbidden_review)
-
- pathExclusions = pd.concat(processed_dfs, ignore_index=True)
-
- return pathExclusions
-
-def generatePathReview(unknown, good, badpathparts, path_exclusion_constant, min_files_for_path, split = False):
-
- df1 = ct.tryToReadCSV(unknown)
- df2 = ct.tryToReadCSV(good)
-
- all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename'])
-
- primary_path_exclusions = calculatePath(all_approved_hashes, badpathparts, path_exclusion_constant, min_files_for_path, split)
-
- remaining_hashes = all_approved_hashes[~all_approved_hashes['sha256'].isin(primary_path_exclusions['sha256'])]
-
- secondary_path_exclusions = calculatePath(remaining_hashes, badpathparts, 3, min_files_for_path, split)
-
- remaining_hashes = remaining_hashes[~remaining_hashes['sha256'].isin(secondary_path_exclusions['sha256'])]
-
- return all_approved_hashes, primary_path_exclusions, secondary_path_exclusions, remaining_hashes
-
-def clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir):
- """
- Prompts user to choose whether to delete all .parquet files or preserve execution_history ones.
- Then deletes .csv, .html, and .parquet files accordingly from specified folders.
- """
- # Prompt user
- user_input = input("Do you want to delete *all* .parquet files including execution_history ones? (yes/y or no/n): ").strip().lower()
- delete_execution_hist = user_input in ["yes", "y"]
-
- folders = [parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir]
-
- for folder in folders:
- folder_path = os.path.abspath(folder)
-
- if not os.path.isdir(folder_path):
- print(f"Folder not found: {folder_path}")
- continue
-
- for filename in os.listdir(folder_path):
- file_path = os.path.join(folder_path, filename)
-
- if not os.path.isfile(file_path):
- continue
-
- _, ext = os.path.splitext(filename)
-
- # Delete .csv and .html files
- if ext in [".csv", ".html"]:
- os.remove(file_path)
- print(f"Deleted: {file_path}")
-
- # Delete .parquet files based on user choice
- elif ext == ".parquet":
- if delete_execution_hist or not filename.startswith("execution_history"):
- os.remove(file_path)
-
-def generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path, split= False ):
-
- if os.path.exists(f"{needappr_base_dir}unknown_hashes.csv") and os.path.exists(f"{needappr_base_dir}good_hashes.csv"):
-
- all_hashes, primary_paths, secondary_paths, remaining = pathf.generatePathReview(f"{appr_base_dir}unknown_hashes.csv", f"{appr_base_dir}good_hashes.csv", badpathparts,path_exclusion_constant, min_files_for_path, split)
-
- all_hashes.to_parquet(f"{parq_base_dir}all_hashes.parquet", index=False)
-
- dataframes = {
- "all_hashes" : all_hashes,
- "primary_Paths": primary_paths,
- "secondary_Paths": secondary_paths,
- "remaining": remaining
-
- }
-
- for name, df in dataframes.items():
- df.to_csv(f"{needappr_base_dir}{name}.csv", index=False)
- df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
- ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html")
-
-
-
- publishers = hashf.generatePublist(f"{parq_base_dir}all_hashes.parquet",bad_publisher_list)
-
- publishers.to_csv(f"{needappr_base_dir}publishers.csv", index=False)
- publishers.to_parquet(f"{parq_base_dir}publishers.parquet", index=False)
- ct.style_dataframe_dark(publishers, f"{needappr_base_dir}publishers.html")
-
- else:
- print(ct.colorText(f"Please manually approve hashes prior to this step","red"))
-
\ No newline at end of file
diff --git a/utils/policyfunctions.py b/utils/policyfunctions.py
deleted file mode 100644
index b1d9e52..0000000
--- a/utils/policyfunctions.py
+++ /dev/null
@@ -1,657 +0,0 @@
-# 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 .
-
-#Local Imports
-import utils.utils as ct
-import utils.hashfunctions as hashf
-import utils.pathfunctions as pathf
-import utils.policyfunctions as policyf
-
-#Standard Libary Imports:
-import datetime
-import gc
-import json
-import os
-import re
-import sys
-#3rd Party Imports:
-import ijson
-import pandas as pd
-import requests
-import tqdm
-from bson import ObjectId
-
-def addHash(url, policy, hash):
- print(f"Adding the following: {hash} \n to {policy}:")
- for p in hash:
- pass
- # print(p)
-
-def addPath(url, policy, hash):
- print(f"Adding the following Path Exclusions to {policy}:")
- for p in hash:
- print(p)
-
-def addPub(url, policy, publist):
- print(f"Adding the following Publishers to {policy}:")
- for p in publist:
- print(p)
-
-def addHashReal(url, allowlistID, hashlist):
- endpoint = url + '/v1/hash/application/add'
-
- payload = {
- "applicationid" : allowlistID,
- "hashes" : hashlist
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- payload = json.dumps(payload)
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- response.raise_for_status() # Raise an error for bad status codes
- parse_text = json.loads(response.text)
- print(parse_text)
-
-def addPathReal(url, grouplistID, pathlist):
- endpoint = url + '/v1/group/path/add'
- payload = {
- "groupid" : grouplistID,
- "path" : pathlist
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- print(payload)
- payload = json.dumps(payload)
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- print(response.text)
-
-def addPubReal(url, grouplistID, publist):
- endpoint = url + '/v1/group/publisher/add'
- payload = {
- "groupid" : grouplistID,
- "publisher" : publist
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- print(payload)
- payload = json.dumps(payload)
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- print(response.text)
-
-def getPolicyInfo(url, policy, type, days, parquet=True):
- executionhist_policy = pd.DataFrame()
- exehist = pullPolicyExechistories(url, 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'])
- if parquet:
- executionhist_policy.to_parquet(f"prepare_policy\\parquet\\execution_history_{policy}.parquet", index=False)
- print(ct.colorText(f"Staging of Execution history for policy: {policy} is complete", "green"))
- del data
- del exehist
- gc.collect()
- return executionhist_policy
-
-def sendToPolicy(url, paths, hashes, publishers, destination_name, destination_id, allowlist_name, allowlist_id):
- pathexclusions = pd.read_parquet(paths)
- allowbyhash = pd.read_parquet(hashes)
- publishers = ct.tryToReadCSV(publishers)
-
- ct.areYouSure()
- confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white"))
-
-
- if confirmation.strip().upper() == "I AGREE":
- print(ct.colorText("Proceeding with the code...", "yellow"))
- print(ct.colorText(f"Adding path exclusions to {destination_name}", "yellow"))
-
- # Get unique combinations of longestcfp and file_extension
- unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
-
- # Regex to match a Windows drive letter at the start (e.g., C:\)
- drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\')
-
- # Build processed paths like \\path\\**.exe or C:\path\**.jar
- processed_paths = [
- (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
- for path, ext in unique_combinations.itertuples(index=False, name=None)
- ]
-
- addPathReal(url, destination_id,processed_paths)
-
- print(ct.colorText(f"Adding publishers to {destination_name}", "yellow"))
-
- if publishers.empty:
- print(ct.colorText("The publishers list is empty.", "red"))
- else:
- publisher_list = publishers['publisher'].tolist()
- addPubReal(url, destination_id, publisher_list)
-
-
- print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow"))
-
- allowlist = allowbyhash['sha256'].unique().tolist()
- addHash(url, allowlist_id,allowlist)
-
- ct.locked()
-
- exit()
-
- else:
- print(ct.colorText("Operation aborted. You MUST EXPLICITLY AGREE to proceed.", "red"))
-
-def pullPolicyExechistories(url, policiesnames, type, days, outputjson: bool):
- file_path = 'chunkinator.json'
- if not os.path.exists(file_path):
- with open(file_path, 'w') as file:
- json.dump({'error': 'Success', 'response': {'exechistories': []}}, file)
- print(f"File '{file_path}' has been created.")
- else:
- print(f"File '{file_path}' already exists.")
- headers = {"X-APIKey": os.getenv('APIKEY')}
- checkpoint = str(skipback(days))
- json_output = {'error': 'Success', 'response': {'exechistories': []}}
- with tqdm.tqdm(file=sys.stdout, leave=True, total=10000, desc=f"Checkpoint Progess: {checkpoint}", colour="blue", initial=1) as filebar:
- with tqdm.tqdm(file=sys.stdout, leave=True, total=100, desc=f"Total of {policiesnames} Complete: ") as pbar:
- while True:
- item = {}
- json_response_data = checkpoint_stomper(checkpoint, url, type, policiesnames, headers)
- histories = json_response_data['response']['exechistories']
- filebar.total=len(histories)
- if not histories:
- break
- match_found = True
- if match_found == True:
- for index, item in enumerate(histories):
- if index == len(histories) - 1:
- checkpoint = item['checkpoint']
- filebar.desc = f"Checkpoint Progress: {checkpoint}"
- break
- else:
- if (datetime.date.today() - datetime.timedelta(days=days) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()):
- pass
- else: json_output['response']['exechistories'].append(item)
- filebar.update(1)
- filebar.refresh()
- 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 item in combined:
- key = (item.get('sha256'), item.get('filename'), item.get('hostname'))
- seen[key] = item
- deduplicated = list(seen.values())
- with open(file_path, 'w') as file:
- json.dump({'error': 'Success', 'response': {'exechistories': deduplicated}}, file)
- json_output['response']['exechistories'].clear()
- date_diff = datetime.date.today() - datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()
- percentage_diff = (((days + 10) - date_diff.days) / (days + 10)) * 100
- pbar.n = round(percentage_diff)
- pbar.set_description_str(f"Total of {policiesnames} Complete: ")
- pbar.refresh()
- filebar.n = 1
- 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 checkpoint_stomper(checkpoint, url, type, policy, headers):
- json_output = {'error': 'Success', 'response': {'exechistories': []}}
- endpoint = url + '/v1/logging/exechistories'
- payload_dict = {
- "type":[type],
- "checkpoint": checkpoint,
- "policy": [policy]
- }
- payload = json.dumps(payload_dict)
- with requests.request("POST", endpoint, headers=headers, data=payload, verify=False, stream=True) as response:
- parser = ijson.items(response.raw, 'response.exechistories.item')
- for item in parser:
- key = (item.get('sha256'), item.get('hostname'))
- if key not in json_output:
- json_output['response']['exechistories'].append(item)
- parse_text = json.loads(json.dumps(json_output))
- return parse_text
-
-def listPolicies(url):
- endpoint = url + '/v1/group'
- print(ct.colorText("[+] Grabbing All Policies", "cyan"))
- payload = {}
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- parse_text = json.loads(response.text)
- return parse_text
-
-def choosePolicies(url):
- allpolicies = listPolicies(url)
- policiesnames = []
- policyids = []
- for index, list in enumerate(allpolicies['response']['groups'], start=1):
- print(ct.colorText(f"{index}. {list['name']}", "yellow"))
- policiesnames.append(list['name'])
- policyids.append(list['groupid'])
- choice = input(ct.colorText("Select Policy Group: ", "white"))
- choice = int(choice) - 1
- return choice, policiesnames, policyids
-
-def getPolicyDataframe(url) -> pd.DataFrame:
- return pd.DataFrame(listPolicies(url)['response']['groups'])
-
-def listATPolicies(url):
- endpoint = url + '/v1/group'
- print(ct.colorText("[+] Grabbing All Policies", "cyan"))
-
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
-
- try:
- response = requests.post(endpoint, headers=headers, json={}, verify=False)
- response.raise_for_status()
- parse_text = response.json()
-
- at_policies = {}
-
- for index, group in enumerate(parse_text.get('response', {}).get('groups', []), start=1):
- name = group.get('name', '')
- if "AT" in name:
- print(ct.colorText(f"{index}. {name}", "yellow"))
- at_policies[name] = group.get('groupid')
-
- return at_policies
-
- except requests.exceptions.RequestException as e:
- print(ct.colorText(f"[!] Request failed: {e}", "red"))
- return {}
- except (KeyError, json.JSONDecodeError) as e:
- print(ct.colorText(f"[!] Failed to parse response: {e}", "red"))
- return {}
-
-def listAllowlists(url: str) -> tuple[int, list, list]:
- endpoint = url + '/v1/application'
- print(ct.colorText("[+] Grabbing All Allowlists", "cyan"))
- payload = {}
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- parse_text = json.loads(response.text)
- policiesnames = []
- policyids = []
-
- for index, item in enumerate(parse_text['response']['applications'], start=1):
- if index >= 38:
- print(ct.colorText(f"{index}. {item['name']}", "yellow"))
- policiesnames.append(item['name'])
- policyids.append(item['applicationid'])
-
- while True:
- try:
- choice = int(input(ct.colorText("Select allowlist: ", "white")))
- if choice < 38 or choice > len(parse_text['response']['applications']):
- print(ct.colorText("Please only choose an allowlist designed for this use - '38+'", "red"))
- else:
- adjusted_choice = choice - 38
- return adjusted_choice, policiesnames, policyids
- except ValueError:
- print(ct.colorText("Invalid input. Please enter a number.", "red"))
-
-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 sendToPolicyTest(url, paths, hashes, publishers, destination_name, destination_id, allowlist_id, allowlist_name):
- pathexclusions = pd.read_parquet(paths)
- allowbyhash = pd.read_parquet(hashes)
- publishers = ct.tryToReadCSV(publishers)
-
-
- print(ct.colorText(f"These path exclusions would be added to {destination_name}", "yellow"))
-
- # Get unique combinations of longestcfp and file_extension
- unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
-
- # Regex to match a Windows drive letter at the start (e.g., C:\)
- drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\')
-
- # Build processed paths like \\path\\**.exe or C:\path\**.jar
- processed_paths = [
- (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
- for path, ext in unique_combinations.itertuples(index=False, name=None)
- ]
-
- addPath(url, destination_id,processed_paths)
-
- print(ct.colorText(f"These publishers would added to {destination_name}", "yellow"))
-
- if publishers.empty:
- print(ct.colorText("The publishers list is empty.", "red"))
- else:
- publisher_list = publishers['publisher'].tolist()
- addPub(url, destination_id, publisher_list)
-
- print(ct.colorText(f"These hashes would be added to {allowlist_name}", "yellow"))
-
- allowlist = allowbyhash['sha256'].unique().tolist()
- addHash(url, allowlist_id,allowlist)
-
-def updateAuditPoliciesFromEnforcementPolices(url, policy_relationship_map):
- for enforcement_policy, audit_policy in policy_relationship_map.items():
- assignPoliciesfromGroup(url, enforcement_policy, audit_policy)
- turnOnAudit(url, audit_policy)
-
-def assignPoliciesfromGroup(url, source_policy_id, target_policy_id):
-
- endpoint = url + '/v1/group/assign'
-
- payload = {
- "groupid" : {source_policy_id},
- "targetgroupid" : {target_policy_id}
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- parse_text = json.loads(response.text)
- print(parse_text)
-
-def turnOnAudit(url, policyid):
-
- endpoint = url + '/v1/group/settings/auditmode'
-
- payload = {
- "groupid" : {policyid},
- "auditmode" : "1"
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- parse_text = json.loads(response.text)
- print(parse_text)
-
-def agentsInPolicy(url, policyid):
- endpoint = url + '/v1/group/agents'
-
- payload = {
- "groupid" : {policyid}
- }
- headers = {
- "X-APIKey": os.getenv('APIKEY')
- }
- response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
- parse_text = json.loads(response.text)
- print(parse_text)
-
-def prepare_to_enforce(url, bad_publisher_list, pups, badpathparts, threat_tolerance_constant = 4, path_exclusion_constant = 4, min_files_for_path = 4):
-
- destination_name = " "
- destination_id = " "
- allowlist_name = " "
- allowlist_id = " "
- policylist = []
- type = [1, 2, 6, 7]
-
- parq_base_dir = "prepare_policy\\parquet\\"
- appr_base_dir = "prepare_policy\\approved\\"
- needappr_base_dir = "prepare_policy\\needs_approved\\"
- pflight_base_dir = "prepare_policy\\preflight\\"
-
- #If the directorys where we're going to store our output dont exist, make them.
- os.makedirs(parq_base_dir, exist_ok=True)
- os.makedirs(needappr_base_dir, exist_ok=True)
- os.makedirs(appr_base_dir, exist_ok=True)
- os.makedirs(pflight_base_dir, exist_ok=True)
-
- while True:
-
- ct.printEnforceChecklist(parq_base_dir,appr_base_dir, needappr_base_dir, pflight_base_dir, policylist, allowlist_name, destination_name)
-
- choice = input(ct.colorText("\nEnter your choice: ", "white"))
-
- if choice == "1":
-
- policylist, policyids = getMultiplePolicySelections(url)
-
- elif choice == "2":
-
- print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white"))
- choice, policynames, policyid = choosePolicies(url)
- #print(allowlist_parent_tuple)
- destination_name = policynames[choice]
- destination_id = policyid[choice]
-
- print(ct.colorText(f"Please choose Allowlist for Hashes","white"))
- choice, allowlists,allowid = listAllowlists(url)
- #print(allowlist_parent_tuple)
- allowlist_name = allowlists[choice]
- allowlist_id = allowid[choice]
-
- print(destination_name, allowlist_name)
-
-
- elif choice == "3":
-
- policyf.buildExecHistory(url,
- policylist,
- parq_base_dir,
- needappr_base_dir,
- type,
- threat_tolerance_constant,
- bad_publisher_list,
- pups,
- )
- csvs = [f"{needappr_base_dir}unknown_hashes.csv", f"{needappr_base_dir}good_hashes.csv"]
- #Since we want to build paths as if they were all in the same policy to begin with, lets group them that way
- for csv in csvs:
- df = ct.tryToReadCSV(csv)
- df['policy'] = destination_name
- df.to_csv(csv)
-
-
- elif choice == "4":
- pathf.generateApprovables(needappr_base_dir, appr_base_dir, parq_base_dir, badpathparts, bad_publisher_list, path_exclusion_constant, min_files_for_path)
-
- elif choice == "5":
- savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir)
-
-
- elif choice == "6":
- if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ":
- sendToPolicyTest(
- url,
- f"{parq_base_dir}final_path_exclusions.parquet",
- f"{parq_base_dir}final_hash_approvals.parquet",
- f"{appr_base_dir}publishers.parquet",
- destination_name,
- destination_id,
- allowlist_name,
- allowlist_id
- )
-
- elif choice == "7":
- if os.path.exists(f"{pflight_base_dir}final_path_exclusions.html") and os.path.exists(f"{pflight_base_dir}\\final_hash_approvals.html") and allowlist_name != " " and destination_name != " ":
- sendToPolicy(
- url,
- f"{parq_base_dir}final_path_exclusions.parquet",
- f"{parq_base_dir}final_hash_approvals.parquet",
- f"{appr_base_dir}publishers.parquet",
- destination_name,
- destination_id,
- allowlist_name,
- allowlist_id,
- )
- elif choice == "R":
-
- pathf.clean_folders(parq_base_dir, appr_base_dir, needappr_base_dir, pflight_base_dir)
-
- elif choice == "Q":
- break
- else:
- print(ct.colorText("Invalid choice. Please try again.", "red"))
-
-def buildExecHistory(url,
- policylist,
- parq_base_dir,
- needappr_base_dir,
- type,
- threat_tolerance_constant,
- bad_publisher_list,
- pups
- ):
- exe_hist_parq_list = []
- while True:
- try:
- history_days = int(input("Enter the how many days in of history do you want to pull - select a number between 1 and 150: "))
- if 1 <= history_days <= 150:
- break
- else:
- print("Invalid input. Please enter a number between 1 and 150.")
- except ValueError:
- print("Invalid input. Please enter a valid integer.")
- for policy in policylist:
- policy_exec_history = policyf.getPolicyInfo(url, policy, type, history_days)
- policy_exec_history['policy'] = policy
- policy_exec_history.to_parquet(f"{parq_base_dir}Exec_Hist_{policy}.parquet")
- exe_hist_parq_list.append(f"{parq_base_dir}Exec_Hist_{policy}.parquet")
-
- augmented_hashlist = hashf.combineHashes(url, exe_hist_parq_list)
- augmented_hashlist.to_parquet(f"{parq_base_dir}augmentedHashlist.parquet",index=False)
-
- needsreview_df, approved_df, unapproved_df = hashf.categorizeHashes(augmented_hashlist, threat_tolerance_constant, bad_publisher_list, pups)
- needsreview_df.to_parquet(f"{parq_base_dir}needsreview.parquet",index=False)
- approved_df.to_parquet(f"{parq_base_dir}approved.parquet",index=False)
- unapproved_df.to_parquet(f"{parq_base_dir}unapproved.parquet",index=False)
-
-
- condensed_executions = hashf.condenseExecutions(exe_hist_parq_list)
- condensed_executions.to_parquet(f"{parq_base_dir}condensed_executions.parquet", index=False)
-
- unknown, good, bad = hashf.divideSortedHashExecutions(
- f"{parq_base_dir}needsreview.parquet",
- f"{parq_base_dir}approved.parquet",
- f"{parq_base_dir}unapproved.parquet",
- f"{parq_base_dir}condensed_executions.parquet",
- pups
- )
-
- dataframes = {
- "unknown_hashes" : unknown,
- "good_hashes": good,
- "bad_hashes": bad
- }
-
- for name, df in dataframes.items():
- df.to_csv(f"{needappr_base_dir}{name}.csv", index=False)
- df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
- ct.style_dataframe_dark(df, f"{needappr_base_dir}{name}.html")
-
-def savePreFlights(appr_base_dir, parq_base_dir, pflight_base_dir):
- if os.path.exists(f"{appr_base_dir}primary_Paths.csv"):
- if not os.path.exists(f"{parq_base_dir}final_hash_approvals.parquet") and not os.path.exists(f"{parq_base_dir}final_path_exclusions.parquet"):
-
- pathexclusions, allowbyhash = hashf.generatePreflights(
- f"{parq_base_dir}all_hashes.parquet",
- f"{appr_base_dir}primary_Paths.csv",
- f"{appr_base_dir}secondary_Paths.csv")
-
- dataframes = {
- "final_path_exclusions" : pathexclusions,
- "final_hash_approvals": allowbyhash
- }
-
- for name, df in dataframes.items():
- df.to_csv(f"{pflight_base_dir}{name}.csv", index=False)
- df.to_parquet(f"{parq_base_dir}{name}.parquet", index=False)
- ct.style_dataframe_dark(df, f"{pflight_base_dir}{name}.html")
-
-def getMultiplePolicySelections(url):
-
- policynameslist = []
- policyidlist = []
- while True:
- choice, policynames, policyid = listPolicies(url)
- selected_policy = policynames[choice]
- selected_policyid = policyid[choice]
-
- if selected_policy not in policynameslist:
- policynameslist.append(selected_policy)
-
- if selected_policyid not in policyidlist:
- policyidlist.append(selected_policyid)
-
-
- while True:
- answer = input(ct.colorText("Do you want to load another policy? (yes/no): ", "white")).strip().lower()
- if answer in ("no", "n"):
- break # Exit the inner loop and then the outer loop
- elif answer in ("yes", "y"):
- break # Exit the inner loop and continue the outer loop
- else:
- print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))
-
- if answer in ("no", "n"):
- break
-
-
- return policynameslist, policyidlist
-
-def getDestAllowlist(url, groupid):
-
- allowlists = getPolicyAllowlists(url, groupid)
-
- matches = allowlists[
- allowlists['name'].str.contains('local', case=False, na=False) &
- allowlists['name'].str.contains('approval', case=False, na=False),
- 'applicationid'
- ].values
- app_id = matches[0] if len(matches) > 0 else None
-
- return app_id
-
-def getPolicyAllowlists(url, groupid):
-
- endpoint = url + '/v1/group/policies'
- payload = {
- "groupid" : f"{groupid}"
- }
-
- headers = {"X-APIKey": os.getenv('APIKEY')}
- payload = json.dumps(payload)
-
- response = requests.post(endpoint, headers=headers, data=payload, verify=False)
- result = json.loads(response.text)
- data = pd.DataFrame(result["response"]["applications"])
- return data
\ No newline at end of file
diff --git a/utils/utils.py b/utils/utils.py
index d43fe5a..10a7368 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -13,11 +13,597 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
-#Standard Libary Imports:
+import json
+import logging
import os
+import platform
+import re
+import subprocess
+import tempfile
+import tkinter as tk
+from tkinter import filedialog, messagebox, simpledialog
+from typing import Callable, Optional, TypeVar
+
import pandas as pd
-def colorText(text: str, color: str) -> str:
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+def load_env_json(key: str, default: str):
+ raw = os.getenv(key, default)
+ try:
+ return json.loads(raw)
+ except json.JSONDecodeError:
+ try:
+ escaped = raw.encode('unicode_escape').decode('utf-8')
+ return json.loads(escaped)
+ except Exception as e:
+ logging.error(f"Failed to parse {key}: {e}")
+ return json.loads(default)
+
+def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
+ """
+ Safely retrieves an environment variable and casts it to the desired type.
+
+ Parameters:
+ key (str): The name of the environment variable.
+ cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
+ default (Optional[T], optional): Default value if the variable is not set or invalid.
+
+ Returns:
+ Optional[T]: The casted value or the default.
+ """
+ value = os.getenv(key)
+ if value is None:
+ logger.warning(f"Environment variable '{key}' not set.")
+ return default
+ try:
+ value = value.strip("'\"") # Strip surrounding quotes
+ return cast_type(value)
+ except (ValueError, TypeError):
+ logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
+ return default
+
+
+def import_to_dataframe(file_path: str) -> pd.DataFrame:
+ df = pd.DataFrame()
+
+ try:
+ if not os.path.exists(file_path):
+ print(colorText(f"Error: File '{file_path}' does not exist.", "red"))
+ return df
+
+ ext = os.path.splitext(file_path)[1].lower()
+
+ if ext == ".csv":
+ df = pd.read_csv(file_path)
+ elif ext == ".parquet":
+ df = pd.read_parquet(file_path)
+ else:
+ print(colorText(f"Error: Unsupported file extension '{ext}'.", "red"))
+ return df
+
+ if df.empty:
+ print(colorText("Error: File has headers but no data rows.", "red"))
+ else:
+ print(colorText(f"Data loaded successfully from {file_path}", "green"))
+
+ return df
+
+ except pd.errors.EmptyDataError:
+ print(
+ colorText(
+ "Notice: CSV file is completely empty, falling back to empty frame",
+ "white",
+ )
+ )
+ return pd.DataFrame()
+
+ except Exception as e:
+ print(colorText(f"Error reading file: {e}", "red"))
+ return pd.DataFrame()
+
+
+def choose_directory():
+ root = tk.Tk()
+ root.withdraw() # Hide the main window
+ directory = filedialog.askdirectory(title="Select a Directory")
+ print("Selected directory:", directory)
+ return directory
+
+
+def choose_file(initial_directory=None, required_substring=None):
+ """Open a file dialog and ensure the selected file contains a required substring."""
+ while True:
+ root = tk.Tk()
+ root.withdraw() # Hide the main window
+ file_path = filedialog.askopenfilename(initialdir=initial_directory)
+
+ if not file_path:
+ print("No file selected.")
+ return None
+
+ if required_substring and required_substring not in file_path:
+ print(
+ f"The selected file must contain '{required_substring}' in its path or name. Please try again."
+ )
+ else:
+ return file_path
+
+
+def choose_save_location():
+ root = tk.Tk()
+ root.withdraw()
+ save_path = filedialog.asksaveasfilename(defaultextension=".txt")
+ return save_path
+
+
+def ask_user_input(message):
+ root = tk.Tk()
+ root.withdraw()
+ user_input = simpledialog.askstring("Input", "{message}]:")
+ return user_input
+
+
+def show_info_message_cli(title, message):
+ root = tk.Tk()
+ root.withdraw()
+ root.after(100, lambda: messagebox.showinfo(title, message))
+ root.mainloop()
+
+
+def show_confirm_question():
+ root = tk.Tk()
+ root.withdraw()
+ response = messagebox.askquestion("Confirm", "Do you want to continue?")
+ print("User response:", response)
+
+
+def regulator(paths, case_insensitive=True):
+ """
+ Build a regex pattern that matches any of the given Windows path fragments.
+ """
+ escaped = [re.escape(p) for p in paths]
+ pattern = "(?:" + "|".join(escaped) + ")"
+ if case_insensitive:
+ pattern = "(?i)" + pattern # Add inline case-insensitive flag
+ print(f"Regulator is providing: {pattern}")
+ return pattern
+
+
+def displayIntro():
+ print(
+ colorText(
+ r"""
+ ███
+ ████ ░████████
+ █████████████ ███████████████
+ █████████████████████ █████████████████████
+ ███████████████████ ██████████████████████▓
+ ███████████████████ ██████████████████████
+ █████████████████████ ███████████████████████
+ ████████████████████████████████████████████████████████
+ █████████ ██ ██ █████████
+ █████████ ██ ███ █ █████████
+ █████████ ██ ████ █████ █████████████
+ █████████ ██ ██████ █████████████
+ ████████ ██ ███████ ████████████░
+ ███████ ██ ██▓ ██████ ████████████
+ ██████ ██ ████ █████ ███████████
+ █████████████████████████████████████████████████
+ ▒████████████████████ ██████████████████
+ ███████████████████ ███████████████▒
+ ███████████████ █████████████
+ ██████████ ███████████
+ ████████
+ ████
+""",
+ "yellow",
+ )
+ )
+ print(
+ colorText(
+ r"""
+ _____ .__ .__ __ ___________ .__
+ / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
+ / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
+/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
+\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
+ \/ \/ \/ \/
+""",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "=================================================================================",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "======================== Welcome to the Airlock API Tool ========================",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "=================================================================================",
+ "cyan",
+ )
+ )
+
+
+def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist):
+ working_dir = load_env("WORKING_DIR")
+ print(
+ colorText(
+ "\n --------------------------------------------------------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " --------------------------------------------------------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "\nSequentually follow these steps to prepare a policy for enforcement:",
+ "white",
+ )
+ )
+
+ print(
+ colorText(
+ "\n1. Choose which originating policy or policies to move to enforcement",
+ "cyan",
+ )
+ )
+ if not selected_policies:
+ print(colorText(" [✗] No policies have been chosen", "red"))
+ else:
+ print(colorText("The following policies have been choosen:", "green"))
+ for policy in selected_policies:
+ print(colorText(f" [✓] {policy.name}", "green"))
+
+ print(colorText("2. Choose the destination policy and allowlist", "cyan"))
+
+ if not destination_policy:
+ print(colorText(" [✗] No destination policy has been chosen", "red"))
+ elif destination_policy:
+ print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green"))
+
+
+
+ if not destination_allowlist:
+ print(colorText(" [✗] No allowlist has been chosen", "red"))
+ elif destination_allowlist:
+ print(
+ colorText(
+ f" [✓] {destination_allowlist[0].name} has been selected as allowlist",
+ "green",
+ )
+ )
+
+
+
+ print(
+ colorText(
+ "3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
+ "cyan",
+ )
+ )
+ if not selected_policies:
+ print(colorText(" [✗] No policies have been chosen", "red"))
+ else:
+ if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
+ print(colorText(" [✓] Data has been fetched", "green"))
+ else:
+ print(colorText(" [✗] Data has not been fetched", "red"))
+
+ print(colorText("4. Manually review the files:", "cyan"))
+ print(
+ colorText(
+ " 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " If metarules need to be created, please make note of them, and remove the row from the csv.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " When complete, save both csv files to the directory 'approved' and choose this option.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
+ "cyan",
+ )
+ )
+
+ if os.path.exists(f"{working_dir}\\Approved\\approved_executions.csv"):
+ print(colorText(" [✓] Reviewed hashes have been loaded", "green"))
+ else:
+ print(colorText(" [✗] Reviewed hashes have not been loaded", "red"))
+
+ if os.path.exists(
+ f"{working_dir}\\Needs_Review\\Review_Second\\primary_Paths.csv",
+ ):
+ print(colorText(" [✓] Path review list created", "green"))
+ else:
+ print(colorText(" [✗] Path review list has not been created", "red"))
+
+ print(
+ colorText(
+ "5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful",
+ "cyan",
+ )
+ )
+ print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
+ print(
+ colorText(
+ " Do the same process with the list of publishers forthe same directories",
+ "cyan",
+ )
+ )
+ print(colorText(" Preflight Lists will be generated", "cyan"))
+
+ if os.path.exists(
+ f"{working_dir}\\Approved\\primary_Paths.csv",
+ ):
+ print(colorText(" [✓] Reviewed path list detected", "green"))
+ else:
+ print(colorText(" [✗] Path review list has not been detected", "red"))
+
+ if os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") and os.path.exists(
+ f"{working_dir}\\Preflight\\approved_hashes.csv"
+ ):
+ print(colorText(" [✓] Preflight Path Exclusion List has been generated", "green"))
+ else:
+ print(colorText(" [✗] Preflight Path Exclusion List has not been generated", "red"))
+
+ print(colorText("6. Test ------------------------------------------------------", "cyan"))
+ print(colorText(" Print rather than apply selected data.", "cyan"))
+
+ print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
+ print(
+ colorText(
+ " Apply path exclusions according to allowed and approved paths",
+ "cyan",
+ )
+ )
+ print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
+ print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
+
+ print(
+ colorText(
+ "R. Remove/Reset Generated data - will prompt to allow keeping execution history",
+ "cyan",
+ )
+ )
+ print(colorText("F. 📂 - Open Working Directory", "cyan"))
+ print(colorText("Q. 🔚 - Quit", "cyan"))
+
+
+def areYouSure():
+ print(
+ colorText(
+ "🛑****************************************************************************************************************************************🛑",
+ "red",
+ )
+ )
+ print(
+ colorText(
+ "⚠️=========================================================================================================================================⚠️",
+ "yellow",
+ )
+ )
+ print(
+ colorText(
+ "🛑========================================================================================================================================🛑",
+ "red",
+ )
+ )
+ print(
+ colorText(
+ "⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️",
+ "yellow",
+ )
+ )
+ print(
+ colorText(
+ "🛑========================================================================================================================================🛑",
+ "red",
+ )
+ )
+ print(
+ colorText(
+ "⚠️=========================================================================================================================================⚠️",
+ "yellow",
+ )
+ )
+ print(
+ colorText(
+ "🛑****************************************************************************************************************************************🛑",
+ "red",
+ )
+ )
+
+
+def locked():
+ print(
+ colorText(
+ r"""
+ ████████████████████████████████████████████████████████████████
+ ███ ██
+ ██ ██████ ███
+ ██ ████████████ ███
+ ██ ████ ███ ███
+ ██ ███ ███ ███
+ ██ ███ ███ ███
+ ██ ▒████████████████████ ███
+ ██ ██████████████████████ ███
+ ██ ██████████████████████ ███
+ ██ ██████████████████████ ███
+ ██ ██████████████████████ ███
+ ██ ██████████████████████ ███
+ ██ ███
+ ███ ███
+ ████████████████████████████████████████████████████████████████████
+ ▒██████████████████████████████████████████████████████████████████▒
+ ▒████
+ ▒████
+ ▓██████████████████████████████████████████
+ █████████████████████████████████████████████░
+""",
+ "yellow",
+ )
+ )
+
+
+def printDeviceEnforceChecklist():
+ print(
+ colorText(
+ "\n --------------------------------------------------------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " --------------------------------------------------------------------",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "\nSequentually follow these steps to prepare a policy for enforcement:",
+ "white",
+ )
+ )
+
+ print(
+ colorText(
+ "\n1. Choose which originating policy or policies to move to enforcement",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes",
+ "cyan",
+ )
+ )
+ print(colorText("3. Manually review the files:", "cyan"))
+ print(
+ colorText(
+ " 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " If metarules need to be created, please make note of them, and remove the row from the csv.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " When complete, save both csv files to the directory 'approved' and choose this option.",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ "4. Manually review the file 'needs_approved\\paths_needing_review.csv'",
+ "cyan",
+ )
+ )
+ print(
+ colorText(
+ " Remove the rows containing path exclusions you do not approve of",
+ "cyan",
+ )
+ )
+ print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan"))
+ print(
+ colorText(
+ " Do the same process with the list of publishers forthe same directories",
+ "cyan",
+ )
+ )
+ print(colorText(" Preflight Lists will be generated", "cyan"))
+
+ print(colorText("5. Choose the destination policy and parent and child allow list", "cyan"))
+
+ print(colorText("6. Test ------------------------------------------------------", "cyan"))
+ print(colorText(" Print rather than apply selected data.", "cyan"))
+
+ print(colorText("7. Liftoff ------------------------------------------------------", "cyan"))
+ print(
+ colorText(
+ " Apply path exclusions according to allowed and approved paths",
+ "cyan",
+ )
+ )
+ print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan"))
+ print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan"))
+
+ print(
+ colorText(
+ "R. Remove/Reset Generated data - will prompt to allow keeping execution history",
+ "cyan",
+ )
+ )
+
+ print(colorText("Q. Quit", "cyan"))
+
+
+def colorText(text, color):
colors = {
"red": "\033[91m",
"green": "\033[92m",
@@ -26,17 +612,17 @@ def colorText(text: str, color: str) -> str:
"magenta": "\033[95m",
"cyan": "\033[96m",
"white": "\033[97m",
- "reset": "\033[0m"
+ "reset": "\033[0m",
}
-
return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"
-def style_dataframe_dark(df, output_html_path=None, overwrite=True):
+
+def formatHTML(df, output_html_path=None, overwrite=True):
from datetime import datetime
# Get current date and filename for subtitle
today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year"
- filename = output_html_path.replace('.html', '') if output_html_path else "Report"
+ filename = output_html_path.replace(".html", "") if output_html_path else "Report"
dark_css = """