Menu and Error Checking updated

This commit is contained in:
=
2025-08-25 09:13:16 -04:00
parent d7e6b43e39
commit 81e2f7f7d1
4 changed files with 65 additions and 55 deletions
+45 -41
View File
@@ -42,10 +42,10 @@ def apivalidation():
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""", "green"))
print(ct.colorText("=================================================================================", "magenta"))
print(ct.colorText("======================== Welcome to the Airlock API Tool ========================", "magenta"))
print(ct.colorText("=================================================================================", "magenta"))
""", "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"))
@@ -54,14 +54,14 @@ def apivalidation():
def menu_main():
while True:
print(ct.colorText("\n-----------------------------------", "white"))
print(ct.colorText("------------ Main Menu ------------", "white"))
print(ct.colorText("-----------------------------------", "white"))
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("11. Exit", "yellow"))
print(ct.colorText("Q. Quit", "yellow"))
choice = input(ct.colorText("\nEnter Menu Item: ", "white"))
if choice == '1':
@@ -72,7 +72,7 @@ def menu_main():
menu_feature2()
elif choice == "4":
menu_prepare_to_enforce()
elif choice == "11":
elif choice == "Q":
break
else:
print(ct.colorText("Invalid choice. Please try again.","red"))
@@ -121,78 +121,79 @@ def menu_prepare_to_enforce():
#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 --------------------------------------------------------------------", "magenta"))
print(ct.colorText(" -------------------- Prepare to Enforce Policy ---------------------", "magenta"))
print(ct.colorText(" --------------------------------------------------------------------", "magenta"))
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"))
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"))
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(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"))
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"))
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"))
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(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"))
print(ct.colorText(" [✓] This step has been completed","green"))
else:
print(ct.colorText(" This step has not been completed","red"))
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"))
print(ct.colorText(" [✓] This step has been completed","green"))
else:
print(ct.colorText(" This step has not been completed","red"))
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"))
print(ct.colorText(" [✓] This step has been completed","green"))
else:
print(ct.colorText(" This step has not been completed", "red"))
print(ct.colorText(" [✗] This step has not been completed", "red"))
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(ct.colorText(" This step has been completed","green"))
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_remaining_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(" [✗] 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"))
print(ct.colorText(" [✓] This step has been completed","green"))
else:
print(ct.colorText(" This step has not been completed","red"))
print(ct.colorText(" [✗] This step has not been completed","red"))
print(ct.colorText("11. Exit", "cyan"))
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): ", "magenta").strip().lower())
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]]
@@ -220,12 +221,12 @@ def menu_prepare_to_enforce():
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"):
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("dataframe_csv\\df_aggregated_{first_policy}.csv") and os.path.exists("dataframe_csv\\df_aggregated_{second_policy}.csv"):
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)
@@ -235,10 +236,11 @@ def menu_prepare_to_enforce():
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"):
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"))
@@ -249,11 +251,12 @@ def menu_prepare_to_enforce():
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_path_ineligible_{first_policy}_{second_policy}.csv") == True and os.path.exists(f"dataframe_csv\\df_path_eligible_{first_policy}_{second_policy}.csv") == True:
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)
@@ -261,19 +264,20 @@ def menu_prepare_to_enforce():
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)
print(ct.colorText(f"Hashes have been categorized","green"))
else:
print(ct.colorText(f"Please complete step 5 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 == "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:
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_remaining_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)
allowed_paths_determined = True
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 == "11":
elif choice == "Q":
break
else:
+2 -2
View File
@@ -58,7 +58,7 @@ def checkpoint_stomper(checkpoint, url, policy, headers):
def listPolicies(url):
endpoint = url + '/v1/group'
print(ct.colorText("[+] Grabbing All Policies", "magenta"))
print(ct.colorText("[+] Grabbing All Policies", "cyan"))
payload = {}
headers = {
"X-APIKey": os.getenv('APIKEY')
@@ -71,6 +71,6 @@ def listPolicies(url):
print(ct.colorText(f"{index}. {list['name']}", "yellow"))
policiesnames.append(list['name'])
policyids.append(list['groupid'])
choice = input(ct.colorText("Select Policy Group: ", "yellow"))
choice = input(ct.colorText("Select Policy Group: ", "white"))
choice = int(choice) - 1
return choice, policiesnames
+16 -10
View File
@@ -58,17 +58,23 @@ def devicehistory(url, outputjson: bool):
}
response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False)
if outputjson == True:
if outputjson:
return response
parse_text = json.loads(response.text)
for block in parse_text['response']['exechistory']:
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")
# Safely get exechistory
exechistory = parse_text.get('response', {}).get('exechistory')
if isinstance(exechistory, list):
for block in exechistory:
print(ct.colorText(f"Command: {block.get('commandline', 'N/A')}", "green"))
print(ct.colorText(f"Date: {block.get('datetime', 'N/A')}", "green"))
print(ct.colorText(f"Filename: {block.get('filename', 'N/A')}", "green"))
print(ct.colorText(f"Policy Name: {block.get('policyname', 'N/A')}", "green"))
print(ct.colorText(f"Hostname: {block.get('hostname', 'N/A')}", "green"))
print(ct.colorText(f"Hash: {block.get('sha256', 'N/A')}", "green"))
print("\n")
else:
print(ct.colorText("No execution history found or data is not in expected format.", "red"))
+2 -2
View File
@@ -16,7 +16,7 @@ import pandas as pd
import requests
import os
import json
import utils.colortext as ct
def aggregateHashes(executions_json) -> pd.DataFrame:
@@ -28,7 +28,7 @@ def aggregateHashes(executions_json) -> pd.DataFrame:
if df.empty:
return df
print(color_text(df,"green"))
print(df)
# Aggregate by sha256, deduplicate lists, and preserve order
agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()