#    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("approvals"): os.makedirs("approvals")

	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("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.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"))
		
		elif choice == "Q":
			break
		
		else:
			print(ct.colorText("Invalid choice. Please try again.", "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()