diff --git a/AirlockTools.py b/AirlockTools.py index 9e70a80..967b873 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -24,6 +24,7 @@ import urllib3 import pandas as pd import sqlite3 import pathlib +import utils.colortext as ct urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) dotenv.load_dotenv() @@ -33,24 +34,10 @@ 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""" + print(ct.colorText(r""" _____ .__ .__ __ ___________ .__ / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ @@ -58,23 +45,27 @@ def apivalidation(): \____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > \/ \/ \/ \/ """, "green")) - print(color_text("============ Welcome to the Airlock API Tool ============", "magenta")) + print(ct.colorText("=================================================================================", "magenta")) + print(ct.colorText("======================== Welcome to the Airlock API Tool ========================", "magenta")) + print(ct.colorText("=================================================================================", "magenta")) match os.getenv('APIKEY'): case '': - print("Please add your API Key to the .env file") + print(ct.colorText("Please add your API Key to the .env file", "red")) 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")) + print(ct.colorText("\n-----------------------------------", "yellow")) + print(ct.colorText("------------ Main Menu ------------", "yellow")) + print(ct.colorText("-----------------------------------", "yellow")) + 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("11. Exit", "yellow")) - choice = input("Enter Menu Item: ") + choice = input(ct.colorText("Enter Menu Item: ", "white")) if choice == '1': utils.getdeviceevents.devicehistory(url,False) elif choice == "2": @@ -86,7 +77,7 @@ def menu_main(): elif choice == "11": break else: - print(color_text("Invalid choice. Please try again.","red")) + print(ct.colorText("Invalid choice. Please try again.","red")) def menu_local_approve(): while True: @@ -129,75 +120,81 @@ 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") + 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")) + print(ct.colorText("\n --------------------------------------------------------------------", "magenta")) + print(ct.colorText(" -------------------- Prepare to Enforce Policy ---------------------", "magenta")) + print(ct.colorText(" --------------------------------------------------------------------", "magenta")) + 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(color_text(f" No policies have been chosen","red")) + print(ct.colorText(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")) + print(ct.colorText(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(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(color_text("2. Pull and stage event history.", "cyan")) + print(ct.colorText("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")) + 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(color_text(f" This has not been completed for {first_policy}","red")) + 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(color_text(f" This has been completed for {second_policy}","green")) + 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(color_text(f" This has not been completed for {second_policy}","red")) + print(ct.colorText(f" This has not been completed for {second_policy}","red")) - print(color_text("3. Combine Staged policies", "cyan")) + 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(color_text(" This step has been completed","green")) + print(ct.colorText(" This step has been completed","green")) else: - print(color_text(" This step has not been completed","red")) + print(ct.colorText(" This step has not been completed","red")) - print(color_text("4. Add hash threat information to list of executions", "cyan")) + 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(color_text(" This step has been completed","green")) + print(ct.colorText(" This step has been completed","green")) else: - print(color_text(" This step has not been completed","red")) + print(ct.colorText(" This step has not been completed","red")) - print(color_text("5. Determine if path exclusions are possible", "cyan")) + 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(color_text(" This step has been completed","green")) + print(ct.colorText(" This step has been completed","green")) else: - print(color_text(" This step has not been completed", "red")) + print(ct.colorText(" This step has not been completed", "red")) - print(color_text("6. Categorize your hashes ", "cyan")) + print(ct.colorText("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")) + print(ct.colorText(" This step has been completed","green")) else: - print(color_text(" This step has not been completed","red")) + print(ct.colorText(" This step has not been completed","red")) - print(color_text("7. Compare potential path exclusions with allowed hashes", "cyan")) + 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(color_text(" This step has been completed","green")) + print(ct.colorText(" This step has been completed","green")) else: - print(color_text(" This step has not been completed","red")) + print(ct.colorText(" This step has not been completed","red")) - print(color_text("11. Exit", "cyan")) + print(ct.colorText("11. Exit", "cyan")) - choice = input(color_text("Enter your choice: ", "yellow")) + choice = input(ct.colorText("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()) + answer = input(ct.colorText(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]] @@ -207,7 +204,7 @@ def menu_prepare_to_enforce(): second_policy = first_policy return False else: - print(color_text("Please answer with 'yes' or 'no'.", "red")) + 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"): @@ -215,14 +212,14 @@ def menu_prepare_to_enforce(): 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")) + 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(color_text(f"Staging of Exection history for policy: {second_policy} is complete","green")) + 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("dataframe_csv\\df_aggregated_{first_policy}.csv"): @@ -237,7 +234,7 @@ def menu_prepare_to_enforce(): 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")) + print(ct.colorText(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"): @@ -245,7 +242,7 @@ def menu_prepare_to_enforce(): 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")) + 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"): @@ -255,7 +252,7 @@ def menu_prepare_to_enforce(): 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")) + 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_path_ineligible_{first_policy}_{second_policy}.csv") == True and os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True: @@ -267,7 +264,7 @@ def menu_prepare_to_enforce(): 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")) + print(ct.colorText(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: @@ -276,23 +273,23 @@ def menu_prepare_to_enforce(): 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")) + print(ct.colorText(f"Please complete step 6 prior to attempting this step","red")) elif choice == "11": break else: - print("Invalid choice. Please try again.") + print(ct.colorText("Invalid choice. Please try again.", "red")) def tryToReadCSV(csv): try: df =pd.read_csv(csv) if df.empty: - print("CSV file has headers but no data rows.") + print(ct.colorText("Error: CSV file has headers but no data rows.", "red")) else: - print("Data loaded successfully.") + print(ct.colorText("Data loaded successfully.", "green")) except pd.errors.EmptyDataError: - print("CSV file is completely empty (no headers, no data).") + 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 diff --git a/utils/allowlist.py b/utils/allowlist.py index fb24b58..5a4376a 100644 --- a/utils/allowlist.py +++ b/utils/allowlist.py @@ -17,6 +17,7 @@ import requests import json import os import time +import utils.colortext as ct def pullPolicyExechistories(url, choice, policiesnames, outputjson: bool): @@ -32,7 +33,7 @@ def pullPolicyExechistories(url, choice, policiesnames, outputjson: bool): for index, item in enumerate(json_response_data['response']['exechistories']): if index == len(json_response_data['response']['exechistories']) -1: checkpoint = item['checkpoint'] - print(f"Date Greater than 30 Days, Stepping to new Checkpoint. {item['checkpoint']}") + print(ct.colorText(f"Date Greater than 30 Days, Stepping to new Checkpoint. {item['checkpoint']}", "blue")) else: if (datetime.date.today() - datetime.timedelta(days=10) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()): pass @@ -57,7 +58,7 @@ def checkpoint_stomper(checkpoint, url, policy, headers): def listPolicies(url): endpoint = url + '/v1/group' - print("[+] Grabbing All Policies") + print(ct.colorText("[+] Grabbing All Policies", "magenta")) payload = {} headers = { "X-APIKey": os.getenv('APIKEY') @@ -67,9 +68,9 @@ def listPolicies(url): policiesnames = [] policyids = [] for index, list in enumerate(parse_text['response']['groups'], start=1): - print(f"{index}. {list['name']}") + print(ct.colorText(f"{index}. {list['name']}", "yellow")) policiesnames.append(list['name']) policyids.append(list['groupid']) - choice = input("Select Policy Group: ") + choice = input(ct.colorText("Select Policy Group: ", "yellow")) choice = int(choice) - 1 return choice, policiesnames \ No newline at end of file diff --git a/utils/colortext.py b/utils/colortext.py new file mode 100644 index 0000000..fcd101c --- /dev/null +++ b/utils/colortext.py @@ -0,0 +1,14 @@ + +def colorText(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" + } + + return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}" \ No newline at end of file diff --git a/utils/getdeviceevents.py b/utils/getdeviceevents.py index ef98e86..9e08c59 100644 --- a/utils/getdeviceevents.py +++ b/utils/getdeviceevents.py @@ -16,16 +16,17 @@ import datetime import requests import json import os +import utils.colortext as ct def devicehistory(url, outputjson: bool): endpoint = url + '/v1/getexechistory' print("\n") - print("1. Today") - print("2. Last 24 Hours") - print("3. Past 7 Days") - print("4. Past 30 Days") - print("5. Custom Date Range") - choice = input("\nSelect Date Range: ") + print(ct.colorText("1. Today", "yellow")) + print(ct.colorText("2. Last 24 Hours", "yellow")) + print(ct.colorText("3. Past 7 Days", "yellow")) + print(ct.colorText("4. Past 30 Days", "yellow")) + print(ct.colorText("5. Custom Date Range","yellow")) + choice = input(ct.colorText("\nSelect Date Range: ", "white")) today = datetime.date.today() today = today.strftime("%Y-%m-%d") if choice == '1': @@ -40,18 +41,18 @@ def devicehistory(url, outputjson: bool): date_selected = datetime.date.today() - datetime.timedelta(days=30) date_selected = date_selected.strftime('%Y-%m-%d') elif choice == "5": - print("Please Input Dates as YYYY-MM-DD") - date_selected = input("From: ") - today = input("Date To: ") - print("WARNING: Device Name is Case Sensitive") - device = input("Enter Device Name: ") + print(ct.colorText("Please Input Dates as YYYY-MM-DD", "cyan")) + date_selected = input(ct.colorText("From: ", "white")) + today = input(ct.colorText("Date To: ", "white")) + print(ct.colorText("WARNING: Device Name is Case Sensitive", "red")) + device = input(ct.colorText("Enter Device Name: ", "white")) payload_dict = { "datefrom": date_selected, "dateto": today, "hostname": device } payload = json.dumps(payload_dict) - print(payload) + print(ct.colorText(payload, "green")) headers = { "X-APIKey": os.getenv('APIKEY') } @@ -64,10 +65,10 @@ def devicehistory(url, outputjson: bool): parse_text = json.loads(response.text) for block in parse_text['response']['exechistory']: - print(f"Command: {block['commandline']}") - print(f"Date: {block['datetime']}") - print(f"Filename: {block['filename']}") - print(f"Policy Name: {block['policyname']}") - print(f"Hostname: {block['hostname']}") - print(f"Hash: {block['sha256']}") + print(ct.colorText(f"Command: {block['commandline']}","green")) + print(ct.colorText(f"Date: {block['datetime']}","green")) + print(ct.colorText(f"Filename: {block['filename']}","green")) + print(ct.colorText(f"Policy Name: {block['policyname']}","green")) + print(ct.colorText(f"Hostname: {block['hostname']}","green")) + print(ct.colorText(f"Hash: {block['sha256']}","green")) print("\n") \ No newline at end of file diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py index 35bd804..8805388 100644 --- a/utils/hashfunctions.py +++ b/utils/hashfunctions.py @@ -16,6 +16,7 @@ import pandas as pd import requests import os import json +import utils.colortext as ct def aggregateHashes(executions_json) -> pd.DataFrame: @@ -27,7 +28,7 @@ def aggregateHashes(executions_json) -> pd.DataFrame: if df.empty: return df - print(df) + print(color_text(df,"green")) # Aggregate by sha256, deduplicate lists, and preserve order agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()