#    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 urllib3
import pandas as pd
import sqlite3
import pathlib

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
treat_tolerance_constant = 4

def color_text(text: str, color: str) -> str:
    colors = {
        "red": "\033[91m",
        "green": "\033[92m",
        "yellow": "\033[93m",
        "blue": "\033[94m",
        "magenta": "\033[95m",
        "cyan": "\033[96m",
        "white": "\033[97m",
        "reset": "\033[0m",
        "bold": "\033[1m",
        "underline": "\033[4m",
    }
    
    return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"

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

def menu_main():
	while True:
		print(color_text("\n--- Main Menu ---", "yellow") + color_text("", "bold"))
		print(color_text("1. Get All Events for Single Device", "yellow"))
		print(color_text("2. Placeholder for Local Approval", "yellow"))
		print(color_text("3. Placeholder for Another Tool", "yellow"))
		print(color_text("4. Prepare Policy For Enforcement", "yellow"))
		print(color_text("11. Exit", "yellow"))

		choice = input("Enter Menu Item: ")
		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 == "11":
			break
		else:
			print(color_text("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 = " "

	df_aggregated_combo = pd.DataFrame()
	while True:
		print(color_text("\n---------Prepare to Enforce Policy ---------------------", "white") + color_text("", "bold"))
		print(color_text("Sequentually follow steps to prepare for policy enforcement", "white") + color_text("", "underline"))
		print(color_text("1. Choose which policy or policies to work with - : ", "cyan"))
		
		if first_policy == " " and second_policy == " ": 
			print(color_text(f"     No policies have been chosen","red"))
		elif first_policy != " " and second_policy is first_policy:
			print(color_text(f"     {first_policy} has been selected,", "green"))
		elif first_policy != " " and second_policy != " ":
			print(color_text(f"     {first_policy} has been selected as Policy 1","green"))
			print(color_text(f"     {second_policy} has been selected as Policy 2","green"))

		print(color_text("2. Pull and stage event history.", "cyan"))
		
		if os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == True: 
			print(color_text(f"     This has been completed for {first_policy}","green"))
		elif os.path.exists(f"dataframe_csv\\df_aggregated_{first_policy}.csv") == False: 
			print(color_text(f"     This has not been completed for {first_policy}","red"))
		elif second_policy is not first_policy and os.path.exists(f"dataframe_csv\\df_aggregated_{second_policy}.csv") == True:
			print(color_text(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(color_text(f"     This has not been completed for {second_policy}","red"))

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

		print(color_text("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(color_text("     This step has been completed","green"))
		else:
			print(color_text("     This step has not been completed","red"))

		print(color_text("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(color_text("     This step has been completed","green"))
		else:
			print(color_text("     This step has not been completed", "red"))

		print(color_text("6. Categorize your hashes ", "cyan"))
				   
		if os.path.exists(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_remaining_hashes_{first_policy}_{second_policy}.html") == True:
			print(color_text("     This step has been completed","green"))
		else:
			print(color_text("     This step has not been completed","red"))
			
		print(color_text("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(color_text("     This step has been completed","green"))
		else:
			print(color_text("     This step has not been completed","red"))
		

		print(color_text("11. Exit", "cyan"))

		choice = input(color_text("Enter your choice: ", "yellow"))
		if choice == "1":
			first_policy_tuple = utils.allowlist.listPolicies(url)
			first_policy = first_policy_tuple[1][first_policy_tuple[0]]
			while True:
				answer = input(color_text(f"{"Do you want to load a second policy?"} (yes/no): ", "magenta").strip().lower())
				if answer in ("yes", "y"):
					second_policy_tuple = utils.allowlist.listPolicies(url)
					second_policy = second_policy_tuple[1][second_policy_tuple[0]]
					return False
				elif answer in ("no", "n"):
					second_policy_tuple = first_policy_tuple
					second_policy = first_policy
					return False
				else:
					print(color_text("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(color_text(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(color_text(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("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("dataframe_csv\\df_aggregated_{first_policy}.csv") and  os.path.exists("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(color_text(f"Please stage your data before attempting this step","red"))

		elif choice == "4":
			if os.path.exists("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)
			else:
				print(color_text(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)
			else:
				print(color_text(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_path_ineligible_{first_policy}_{second_policy}.csv") == True and os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True:
				categorized = utils.hashfunctions.categorizeHashes(pd.read_csv(f"dataframe_csv\\df_augmented_combo_{first_policy}_{second_policy}.csv"), treat_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_remaining_hashes__{first_policy}_{second_policy}.html", index=False)
				categorized[2].to_csv(f"dataframe_csv\\df_remaining_hashes__{first_policy}_{second_policy}.csv", index=False)
			else:
				print(color_text(f"Please complete step 5 prior to attempting this step","red"))

		elif choice == "7":
			if os.path.exists(f"dataframe_html\\df_hashes_needing_approval_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_automatically_approved_hashes_{first_policy}_{second_policy}.html") == True and os.path.exists(f"dataframe_html\\df_remaining_hashes_{first_policy}_{second_policy}.html") == True:
				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)
				allowed_paths_determined = True
			else:
				print(color_text(f"Please complete step 6 prior to attempting this step","red"))
		
		elif choice == "11":
			break
		
		else:
			print("Invalid choice. Please try again.")    
		
def tryToReadCSV(csv):
	try:
		df =pd.read_csv(csv)
		if df.empty:
			print("CSV file has headers but no data rows.")
		else:
			print("Data loaded successfully.")
	except pd.errors.EmptyDataError:
			print("CSV file is completely empty (no headers, no data).")
			df = pd.DataFrame()  # Create an empty DataFrame as fallback
	return df



if __name__ == "__main__":
	apivalidation()