AirlockTools/AirlockTools.py

#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
import argparse
import dotenv
import os
import re
import pandas as pd
import urllib3
import utils.clientfunctions
import utils.hashfunctions
import utils.otpfunctions
import utils.pathfunctions
import utils.policyfunctions
import utils.pretty as ct

from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler


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"]
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

#"bf0b1f9b-bfea-4f44-97c0-80e27ff61712" : "d126db36-72ed-4937-adc7-d88b7509a5b5" #AT Servers General, AT Testing
}



def main():
	parser = argparse.ArgumentParser(description="Your script description")
	parser.add_argument('--monitorOTP', action='store_true', help='Run in non-interactive mode')
	# Add other arguments as needed

	args = parser.parse_args()

	if args.monitorOTP:
        # Non-interactive logic
			print(f"Running non-interactively to start monitoring OTP")
			
			os.makedirs("scheduling", exist_ok=True)
			os.makedirs("OTP/HTML", exist_ok=True)
			os.makedirs("OTP/PARQ", exist_ok=True)

			apivalidation()
					
			register_function("monitorOTP", utils.otpfunctions.monitorOTP)
			register_function("updateAudit", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices)
			
			if not os.path.exists("scheduling\\jobs.json"): 
				recurring_job("monitorOTP", "monitorOTP", interval=60, unit="seconds", args=[url, pups])
				recurring_job("updateAudit", "updateAudit", interval=10, unit="minutes", args=[url, policy_relationship_map])
			else:
				reload_jobs()
			start_scheduler()

	else:
       # Interactive logic
		apivalidation()
		menu_main()


def apivalidation(): 
	match os.getenv('APIKEY'):
		case '':
			print(ct.colorText("Please add your API Key to the .env file", "red"))


def tryToReadCSV(csv):
	try:
		df =pd.read_csv(csv)
		if df.empty:
			print(ct.colorText("Error: CSV file has headers but no data rows.", "red"))
		else:
			print(ct.colorText(f"Data loaded successfully from {csv}", "green"))
	except pd.errors.EmptyDataError:
			print(ct.colorText("Notice : CSV file is completely empty (no headers, no data), falling back to empty frame", "white"))
			df = pd.DataFrame()  # Create an empty DataFrame as fallback
	return df

def tryToReadParquet(parquet):
	try:
		df = pd.read_parquet(parquet)
		if df.empty:
			print(ct.colorText("Error: Parquet file has headers but no data rows.", "red"))
		else:
			print(ct.colorText(f"Data loaded successfully from {parquet}", "green"))
	except pd.errors.EmptyDataError:
			print(ct.colorText("Notice : Parquet file is completely empty (no headers, no data), falling back to empty frame", "white"))
			df = pd.DataFrame()  # Create an empty DataFrame as fallback
	return df

def deduplicate_list(lst):
    seen = set()
    return [x for x in lst if not (x in seen or seen.add(x))]

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 Unexcluded from 24hr Execution", "yellow"))
		print(ct.colorText("4. Prepare Policy For Enforcement", "yellow"))
		print(ct.colorText("5. Update Audit Policies from Enforcement Policies", "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":
			find_unexcluded()
		elif choice == "4":
			menu_prepare_to_enforce()
		elif choice == "5":
			utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices(url, 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.")

def find_unexcluded():
	while True:
		print("\n--- Follow Steps Sequentially, Files will land in directory named 'exclusions'")
		print("1. Pull last 24-72 hours execution for all ATPolicys")
		print("2. Generate Paths")
		print("3.  Merge")
		print("4. Pull existing paths")
		print("5. Divide By Excluded or not excluded by path")
		print("Q. Exit")

		if not os.path.exists("exclusions"): os.makedirs("exclusions")
		choice = input("Enter your choice: ")

		if choice == "1":

			# Get the AT policies dictionary
			atpolicies = utils.policyfunctions.listATPolicies(url)

			# List to hold each policy's DataFrame
			all_dfs = []

			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 <= 3:
						break
					else:
						print("Invalid input. Please enter a number between 1 and 3.")
				except ValueError:
					print("Invalid input. Please enter a valid integer.")




			# Loop through each policy name
			for policy_name in atpolicies:
				try:
					# Get the policy info DataFrame
					df = utils.policyfunctions.getPolicyInfo(url, policy_name, [1, 2, 6, 7], history_days, False)
					
					# Add a column to indicate the policy name
					df['PolicyName'] = policy_name
					
					# Append to the list
					all_dfs.append(df)
				except Exception as e:
					print(f"Error processing policy '{policy_name}': {e}")

			# Combine all DataFrames into one
			if all_dfs:
				combined_df = pd.concat(all_dfs, ignore_index=True)
				print("✅ Combined DataFrame created.")
			else:
				combined_df = pd.DataFrame()
				print("⚠️ No data was retrieved.")
			
			combined_df.to_csv("exclusions\\allATPolicyExecs.csv", index=False)


		elif choice == "2":
			utils.pathfunctions.allATpaths(url,pups, bad_publisher_list, badpathparts, threat_tolerance_constant, path_exclusion_constant, min_files_for_path)

		elif choice == "3":
			utils.pathfunctions.mergeTesting()

		elif choice == "4":


			# Get the AT policies dictionary
			atpolicies = utils.policyfunctions.listATPolicies(url)

			# List to hold each policy's path data
			all_paths = []

			# Loop through each policy name and group ID
			for policy_name, group_id in atpolicies.items():
				try:
					# Get the list of paths using the group ID
					paths = utils.pathfunctions.listPaths(url, group_id)

					# Ensure each path is a string and append properly
					for path in paths:
						if isinstance(path, str):
							all_paths.append({
								'PolicyName': policy_name,
								'GroupID': group_id,
								'Path': path
							})
				except Exception as e:
					print(f"Error retrieving paths for policy '{policy_name}': {e}")

			# Convert to DataFrame
			if all_paths:
				paths_df = pd.DataFrame(all_paths)
				print(paths_df.head())  # Optional: preview first few rows
				print("✅ Paths DataFrame created.")
			else:
				paths_df = pd.DataFrame()
				print("⚠️ No paths were retrieved.")

			# Save to CSV
			paths_df.to_csv("exclusions\\paths.csv", index=False)
		elif choice == "5":

			# Load and clean data
			filenames_df = pd.read_csv('exclusions\\merged_output.csv')  # Contains 'PolicyName' and 'filename'
			exclusions_df = pd.read_csv('exclusions\\paths.csv')  # Contains 'PolicyName' and 'Path'

			# Clean and normalize columns
			filenames_df['filename'] = filenames_df['filename'].fillna('').astype(str).str.strip()
			filenames_df['PolicyName'] = filenames_df['PolicyName'].fillna('').astype(str).str.strip()
			exclusions_df['Path'] = exclusions_df['Path'].fillna('').astype(str).str.strip()
			exclusions_df['PolicyName'] = exclusions_df['PolicyName'].fillna('').astype(str).str.strip()

			# Decode escaped backslashes in exclusion patterns
			exclusions_df['Path'] = exclusions_df['Path'].apply(lambda p: p.encode('utf-8').decode('unicode_escape'))

			# Build exclusion map: {PolicyName: [compiled regex patterns]}
			exclusion_map = {}
			for _, row in exclusions_df.iterrows():
				policy = row['PolicyName']
				raw_pattern = row['Path']
				try:
					regex = utils.pathfunctions.wildcardRegex(raw_pattern)
					print(f"[EXCLUSION MAP] Policy: {policy}, Pattern: {raw_pattern} → Regex: {regex.pattern}")
					exclusion_map.setdefault(policy, []).append(regex)
				except Exception as e:
					print(f"[ERROR] Failed to compile pattern for Policy: {policy}, Path: {raw_pattern}, Error: {e}")

			# Function to check if a filename matches any exclusion pattern for its PolicyName
			def is_excluded(row):
				policy = row['PolicyName'].strip()
				path = os.path.normpath(row['filename'].strip())
				patterns = exclusion_map.get(policy, [])
				for r in patterns:
					if r.match(path):
						print(f"[MATCH] {policy}: {path} matches {r.pattern}")
						return True
				print(f"[NO MATCH] {policy}: {path}")
				return False

			# Apply matching
			excluded = filenames_df[filenames_df.apply(is_excluded, axis=1)]
			not_excluded = filenames_df[~filenames_df.apply(is_excluded, axis=1)]

			
			# Desired column order
			column_order = [
				'PolicyName', 'filename', 'longestcfp',
				'pprocess', 'gprocess',  'sha256', 'publisher', 'description', 'productname','commandline',  'middle', 'filename_only',
				'file_extension', 'unique_sha256_count', 'hostname', 'username', 'productversion',
				'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount',
				'reputation_status', 'reputation_threatlevel', 'reputation_threatname', 'reputation_timestamp'
			]

			# Reorder columns (ignore missing ones)
			excluded = excluded[[col for col in column_order if col in excluded.columns]]
			not_excluded = not_excluded[[col for col in column_order if col in not_excluded.columns]]


			# Save results
			excluded.to_csv('exclusions\\excluded_filenames.csv', index=False)
			not_excluded.to_csv('exclusions\\non_excluded_filenames.csv', index=False)


		elif choice == "Q":
			break
		else:
			print("Invalid choice. Please try again.")

def menu_prepare_to_enforce():

	first_policy = " "
	second_policy = " "
	destination_name = " "
	destination_id = " "
	allowlist_parent_name = " "
	allowlist_parent_id = " "
	allowlist_child_name = " "
	allowlist_child_id = " "
	history_days = 1
	
	#If the directorys where we're going to store our output dont exist, make them. 
	if not os.path.exists("parquet"): os.makedirs("parquet")
	if not os.path.exists("needs_approved"): os.makedirs("needs_approved")
	if not os.path.exists("approved"): os.makedirs("approved")
	if not os.path.exists("preflight"): os.makedirs("preflight")

	while True:
		
		ct.printEnforceChecklist(first_policy, second_policy, allowlist_child_name, allowlist_parent_name, destination_name)
		
		choice = input(ct.colorText("\nEnter your choice: ", "white"))
		
		if choice == "1":
			
			choice, policynames, policyid = utils.policyfunctions.listPolicies(url)
			first_policy = policynames[choice]
			while True:
				answer = input(ct.colorText(f"{"Do you want to load a second policy?"} (yes/no): ", "white").strip().lower())
				if answer in ("yes", "y"):
					choice, policynames, policyid = utils.policyfunctions.listPolicies(url)
					second_policy = policynames[choice]
			
					break 
				elif answer in ("no", "n"):
					second_policy = first_policy
					break
				else:
					print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))

		elif choice == "2":

			
			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.")


			if not os.path.exists(f"parquet\\execution_history_{first_policy}.parquet"):
				utils.policyfunctions.getPolicyInfo(url, first_policy, [1,2,6,7], history_days)

			if not os.path.exists(f"parquet\\execution_history_{second_policy}.parquet"):
				utils.policyfunctions.getPolicyInfo(url, second_policy, [1,2,6,7], history_days)

			if not os.path.exists(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"):
				utils.hashfunctions.combineHashes(url, first_policy, second_policy)

			if not os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"):
				utils.hashfunctions.categorizeHashes(
					first_policy,
					second_policy,
					pd.read_parquet(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"),
					threat_tolerance_constant,
					bad_publisher_list,
					pups
				)
						
			if not os.path.exists(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet"):
				utils.hashfunctions.condenseExecutions(first_policy,second_policy)
				
			if (
				os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") and
				os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") and
				os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet") and
				not os.path.exists(f"needs_approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv")
			):utils.hashfunctions.divideSortedHashExecutions(first_policy,second_policy,pups)
				
		elif choice == "3":
			
			if os.path.exists(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv") and os.path.exists(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv"):
				utils.pathfunctions.generatePathReview(first_policy, second_policy, badpathparts,path_exclusion_constant, min_files_for_path)
				utils.hashfunctions.generatePublist(first_policy,second_policy,bad_publisher_list)
			else:
				print(ct.colorText(f"Please manually approve hashes prior to this step","red"))
					
		elif choice == "4":
			
			if  os.path.exists(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv"):
				if not os.path.exists(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet"):
					utils.hashfunctions.generatePreflights(first_policy, second_policy)

		elif choice == "5":
	
			print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white"))
			choice, policynames, policyid = utils.policyfunctions.listPolicies(url)
			#print(allowlist_parent_tuple)
			destination_name = policynames[choice]
			destination_id = policyid[choice]
			
			print(ct.colorText(f"Please choose Parent Allowlist for Known Hashes","white"))
			choice, allowlists,allowid = utils.policyfunctions.listAllowlists(url)
			#print(allowlist_parent_tuple)
			allowlist_parent_name = allowlists[choice]
			allowlist_parent_id = allowid[choice]

			print(ct.colorText(f"Please choose Child Allowlist for Less-Known Hashes","white"))
			choice, allowlists, allowid = utils.policyfunctions.listAllowlists(url)
			#print(allowlist_child_tuple)
			allowlist_child_name = allowlists[choice]
			allowlist_child_id = allowid[choice]


		elif choice == "6":
			if os.path.exists(f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html") and os.path.exists(f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html")	and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ":
				utils.policyfunctions.sendToPolicyTest(
					url,
					first_policy,
					second_policy,
					destination_name,
					destination_id,
					allowlist_parent_name,
					allowlist_parent_id, 
					allowlist_child_name,
					allowlist_child_id
					)

		elif choice == "7":
			if os.path.exists(f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html") and os.path.exists(f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html")	and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ":
				utils.policyfunctions.sendToPolicy(
					url,
					first_policy,
					second_policy,
					destination_name,
					destination_id,
					allowlist_parent_name,
					allowlist_parent_id, 
					allowlist_child_name,
					allowlist_child_id
					)

		elif choice == "Q":
			break
		else:
			print(ct.colorText("Invalid choice. Please try again.", "red"))    
		

if __name__ == "__main__":
    main()