#    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 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", "VNC", "Kaseya", "Solarwinds", "mRemoteNG"]
badpathparts = ["users", "wwwroot", "windows\\task", "windows\\system32", "startup", "windows\\fonts", "Recycle.Bin", "AppData", "programdata", "Solarwinds", "kaseya", ] #"Windows\\assembly", "WindowsPowerShell\\Modules", "windows\\temp"3
path_exclusion_constant = 4
min_files_for_path = 4
threat_tolerance_constant = 4

"""
Execution Types
 0 = "Trusted Execution",
 1 = "Blocked Execution", 
 2 = "Untrusted Execution [Audit]", 
 3 = "Untrusted Execution [OTP]", 
 4 = "Trusted Path Execution", 
 5 = "Trusted Publisher Execution", 
 6 = "Blocklist Execution", 
 7 = "Blocklist Execution [Audit]", 
 8 = "Trusted Process Execution", 
 9 = "Constrained Execution", 
 10 = "Trusted Metadata Execution",
 11 = "Trusted Browser Execution", 
 12 = "Blocked Browser Execution", 
 13 = "Untrusted Browser Execution [Audit]", 
 14 = "Untrusted Browser Execution [OTP]", 
 15 = "Blocklist Browser Execution [Audit]", 
 16 = "Blocklist Browser Execution", 
 17 = "Trusted Installer Execution", 
 18 = "Trusted Browser Metadata Execution"
"""


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)
			
			if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitor", "monitorOTP", interval=60, unit="seconds", args=[url, pups])
			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. Placeholder for Another Tool", "yellow"))
		print(ct.colorText("4. Prepare Policy For Enforcement", "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":
			menu_feature2()
		elif choice == "4":
			menu_prepare_to_enforce()
		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 menu_feature2():
	while True:
		print("\n--- Submenu ---")
		print("1. Pull last 24 horus execution for all ATPolicys")
		print("2. Generate Paths")
		print("3  Merge")
		print("Q. Exit")
		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 = []

			# 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], 1, 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("allATPolicyExecs.csv", index=False)


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

		elif choice == "3":
			utils.pathfunctions.mergeTesting()
		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()