#    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 dotenv
import os
import utils.getdeviceevents
import utils.allowlist
import utils.hashfunctions
import utils.pathfunctions
import utils.allowfunctions
import utils.colortext as ct
import urllib3
import pandas as pd

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
dotenv.load_dotenv()

url = "https://172.17.22.240:3129"
badpublisherlist = ["Brave Software, Inc.",	"Zoom Video Communications, Inc."]
path_exclusion_constant = 3
threat_tolerance_constant = 4



def apivalidation(): 
	print(ct.colorText(r"""
   _____  .__       .__                 __     ___________           .__          
  /  _  \ |__|______|  |   ____   ____ |  | __ \__    ___/___   ____ |  |   ______
 /  /_\  \|  \_  __ \  |  /  _ \_/ ___\|  |/ /   |    | /  _ \ /  _ \|  |  /  ___/
/    |    \  ||  | \/  |_(  <_> )  \___|    <    |    |(  <_> |  <_> )  |__\___ \ 
\____|__  /__||__|  |____/\____/ \___  >__|_ \   |____| \____/ \____/|____/____  >
        \/                           \/     \/                                 \/ 
""", "cyan"))
	print(ct.colorText("=================================================================================", "cyan"))
	print(ct.colorText("======================== Welcome to the Airlock API Tool ========================", "cyan"))
	print(ct.colorText("=================================================================================", "cyan"))
	match os.getenv('APIKEY'):
		case '':
			print(ct.colorText("Please add your API Key to the .env file", "red"))
		case _:
			menu_main()

def menu_main():
	while True:
		print(ct.colorText("\n-----------------------------------", "magenta"))
		print(ct.colorText("------------ Main Menu ------------", "magenta"))
		print(ct.colorText("-----------------------------------", "magenta"))
		print(ct.colorText("1. Get All Events for Single Device", "yellow"))
		print(ct.colorText("2. Placeholder for Local Approval", "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.getdeviceevents.devicehistory(url,False)
		elif choice == "2":
			menu_local_approve()
		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_local_approve():
   while True:
        print("\n--- Submenu ---")
        print("1. Sub-option A")
        print("2. Sub-option B")
        print("3. Return to Main Menu")
        choice = input("Enter your choice: ")

        if choice == "1":
            print("You selected Sub-option A")
        elif choice == "2":
            print("You selected Sub-option B")
        elif choice == "3":
            print("Returning to Main Menu...")
            break
        else:
            print("Invalid choice. Please try again.")

def menu_feature2():
   while True:
        print("\n--- Submenu ---")
        print("1. Sub-option A")
        print("2. Sub-option B")
        print("3. Return to Main Menu")
        choice = input("Enter your choice: ")

        if choice == "1":
            print("You selected Sub-option A")
        elif choice == "2":
            print("You selected Sub-option B")
        elif choice == "3":
            print("Returning to Main Menu...")
            break
        else:
            print("Invalid choice. Please try again.")

def menu_prepare_to_enforce():

	first_policy = " "
	second_policy = " "

	#If the directorys where we're going to store our output dont exist, make them. 
	if not os.path.exists("dataframe_html"): os.makedirs("dataframe_html")
	if not os.path.exists("dataframe_csv"): os.makedirs("dataframe_csv")
	if not os.path.exists("manuallyapproved"): os.makedirs("manuallyapproved")

	df_aggregated_combo = pd.DataFrame()
	while True:
		print(ct.colorText("\n --------------------------------------------------------------------", "cyan"))
		print(ct.colorText(" -------------------- Prepare to Enforce Policy ---------------------", "cyan"))
		print(ct.colorText(" --------------------------------------------------------------------", "cyan"))
		print(ct.colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
		print(ct.colorText("\n1. Choose which policy or policies to work with - : ", "cyan"))
		
		if first_policy == " " and second_policy == " ": 
			print(ct.colorText(f"  [✗]  No policies have been chosen","red"))
		elif first_policy != " " and second_policy is first_policy:
			print(ct.colorText(f"  [✓]  {first_policy} has been selected,", "green"))
		elif first_policy != " " and second_policy != " ":
			print(ct.colorText(f"  [✓]  {first_policy} has been selected as Policy 1","green"))
			print(ct.colorText(f"  [✓]  {second_policy} has been selected as Policy 2","green"))

		print(ct.colorText("2. Pull and stage event history", "cyan"))
		
		if os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == True: 
			print(ct.colorText(f"  [✓]  This has been completed for {first_policy}","green"))
		elif os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == False: 
			print(ct.colorText(f"  [✗]  This step has not been completed","red"))
		elif second_policy is not first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv") == True:
			print(ct.colorText(f"  [✓]  This has been completed for {second_policy}","green"))
		elif second_policy is not first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv") == False:
			print(ct.colorText(f"  [✓]  This has not been completed for {second_policy}","red"))

		print(ct.colorText("3. Combine Staged policies", "cyan"))
		
		if os.path.exists(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv") == True:
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed","red"))

		print(ct.colorText("4. Add hash threat information to list of executions", "cyan"))
	
		if os.path.exists(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv") == True:
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed","red"))

		print(ct.colorText("5. Determine if path exclusions are possible", "cyan"))
		
		if os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True:
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed", "red"))

		print(ct.colorText("6. Categorize your hashes ", "cyan"))
				   
		if os.path.isfile(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv") and os.path.isfile(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv") and os.path.isfile(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv"):
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed","red"))

		print(ct.colorText(f"7. Manually review the files \\dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv and dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv", "cyan"))
		print(ct.colorText("     Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", "cyan"))
		print(ct.colorText("     If metarules need to be created, please make note of them, and remove the row from the csv.", "cyan"))
		print(ct.colorText("     When complete, move both csv files to the directory 'manuallyapproved' and choose this option to combine these approved hashes with the automatically approved hashes", "cyan"))

		if os.path.isfile(f"dataframe_csv\\df_all_approved_hashes_{first_policy}_{second_policy}.csv"):
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed","red"))

		"""	
		print(ct.colorText("7. Compare potential path exclusions with allowed hashes", "cyan"))
				   
		if os.path.exists(f"dataframe_csv\\df_allowed_paths_{first_policy}_{second_policy}.csv") == True:
			print(ct.colorText("  [✓]  This step has been completed","green"))
		else:
			print(ct.colorText("  [✗]  This step has not been completed","red"))
		"""

		print(ct.colorText("Q. Quit", "cyan"))

		choice = input(ct.colorText("\nEnter your choice: ", "white"))
		if choice == "1":
			first_policy_tuple = utils.allowlist.listPolicies(url)
			first_policy = first_policy_tuple[1][first_policy_tuple[0]]
			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"):
					second_policy_tuple = utils.allowlist.listPolicies(url)
					second_policy = second_policy_tuple[1][second_policy_tuple[0]]
					break 
				elif answer in ("no", "n"):
					second_policy_tuple = first_policy_tuple
					second_policy = first_policy
					break
				else:
					print(ct.colorText("Please answer with 'yes' or 'no'.", "red"))

		elif choice == "2":
			if  not os.path.exists("dataframe_csv\\df_aggregated_{first_policy}.csv"):
				executionhist_policy1 = utils.allowlist.pullPolicyExechistories(url,first_policy_tuple[0], first_policy_tuple[1],True)
				df_aggregated_policy1 = utils.hashfunctions.aggregateHashes(executionhist_policy1)
				df_aggregated_policy1.to_html(f"dataframe_html\\df_aggregated_{first_policy}.html", index=False)
				df_aggregated_policy1.to_csv(f"dataframe_csv\\df_aggregated_{first_policy}.csv", index=False)
				print(ct.colorText(f"Staging of Exection history for policy: {first_policy} is complete","green"))
			
			if not os.path.exists("dataframe_csv\\df_aggregated_{second_policy}.csv"):
				executionhist_policy2 = utils.allowlist.pullPolicyExechistories(url,second_policy_tuple[0], second_policy_tuple[1],True)
				df_aggregated_policy2 = utils.hashfunctions.aggregateHashes(executionhist_policy2)
				df_aggregated_policy2.to_html(f"dataframe_html\\df_aggregated_{second_policy}.html", index=False)
				df_aggregated_policy2.to_csv(f"dataframe_csv\\df_aggregated_{second_policy}.csv", index=False)
				print(ct.colorText(f"Staging of Exection history for policy: {second_policy} is complete","green"))

		elif choice == "3":					
			if  second_policy is first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv"):
				df1 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{first_policy}.csv")
				df_aggregated_combo = df1
				df_aggregated_combo.to_html(f"dataframe_html\\df_aggregated_combo_{first_policy}_{second_policy}.html", index=False)
				df_aggregated_combo.to_csv(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv", index=False)
			elif os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv"):
				df1 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{first_policy}.csv")
				df2 = tryToReadCSV(f"dataframe_csv\\df_aggregated_{second_policy}.csv")
				df_aggregated_combo = pd.concat([df1 , df2], ignore_index=True)
				df_aggregated_combo.to_html(f"dataframe_html\\df_aggregated_combo_{first_policy}_{second_policy}.html", index=False)
				df_aggregated_combo.to_csv(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv", index=False)
			else:
				print(ct.colorText(f"Please stage your data before attempting this step","red"))

		elif choice == "4":
			if os.path.exists(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv"):
				df_augmented = utils.hashfunctions.augmentAggregatedHashes(url,tryToReadCSV(f"dataframe_csv\\df_aggregated_combo_{first_policy}_{second_policy}.csv"))
				df_augmented.to_html(f"dataframe_html\\df_augmented_combo_{first_policy}_{second_policy}.html", index=False)
				df_augmented.to_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv", index=False)
				print(ct.colorText(f"Hash reputation info added to dataframe","green"))
			else:
				print(ct.colorText(f"Please combine your data with step 3 prior to attempting this step","red"))

		elif choice == "5":
			if  os.path.exists(f"dataframe_html\\df_augmented_combo_{first_policy}_{second_policy}.html"):
				path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"))
				path_eligible.to_html(f"dataframe_html\\df_path_eligible_{first_policy}_{second_policy}.html", index=False)
				path_eligible.to_csv(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv", index=False)
				path_ineligible.to_html(f"dataframe_html\\df_path_ineligible_{first_policy}_{second_policy}.html", index=False)
				path_ineligible.to_csv(f"dataframe_csv\\df_path_ineligible_{first_policy}_{second_policy}.csv", index=False)
				print(ct.colorText(f"Eligible paths determined","green"))
			else:
				print(ct.colorText(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red"))

		elif choice == "6":
			if os.path.exists(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"):
				categorized = utils.hashfunctions.categorizeHashes(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"), threat_tolerance_constant, badpublisherlist)
				categorized[0].to_html(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html", index=False)
				categorized[0].to_csv(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv", index=False)
				categorized[1].to_html(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html", index=False)
				categorized[1].to_csv(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv", index=False)
				categorized[2].to_html(f"dataframe_html\\df_unapproved_hashes__{first_policy}_{second_policy}.html", index=False)
				categorized[2].to_csv(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv", index=False)
				print(ct.colorText(f"Hashes have been categorized","green"))
			else:
				print(ct.colorText(f"Please Augment your data with hash threat info using step 4 prior to attempting this step","red"))

		elif choice == "7":
			if os.path.isfile(f"manuallyapproved\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv"):
				df1 = tryToReadCSV(f"manuallyapproved\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv")
				df2 = tryToReadCSV(f"manuallyapproved\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv")
				df_all_approved_hashes = pd.concat([df1 , df2], ignore_index=True)
				df_aggregated_combo.to_html(f"dataframe_html\\df_all_approved_hashes_{first_policy}_{second_policy}.html", index=False)
				df_aggregated_combo.to_csv(f"dataframe_csv\\df_all_approved_hashes_{first_policy}_{second_policy}.csv", index=False)


		elif choice == "Q":
			break
		else:
			print(ct.colorText("Invalid choice. Please try again.", "red"))    
		
		"""
		elif choice == "7":
			if os.path.exists(f"dataframe_csv\\df_hashes_needing_approval_{first_policy}_{second_policy}.csv") and os.path.exists(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv") and os.path.exists(f"dataframe_csv\\df_unapproved_hashes__{first_policy}_{second_policy}.csv"):
				allowpaths = utils.allowfunctions.filter_and_drop(pd.read_csv(f"dataframe_csv\\df_automatically_approved_hashes_{first_policy}_{second_policy}.csv"),tryToReadCSV(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv"), path_exclusion_constant)
				allowpaths.to_html(f"dataframe_html\\df_allowed_paths_{first_policy}_{second_policy}.html", index=False)
				allowpaths.to_csv(f"dataframe_csv\\df_allowed_paths_{first_policy}_{second_policy}.csv", index=False)
				print(ct.colorText(f"Allowable paths determined","green"))
			else:
				print(ct.colorText(f"Please complete step 6 prior to attempting this step","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("Data loaded successfully.", "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



if __name__ == "__main__":
	apivalidation()