Files
AirlockTools/AirlockTools.py
T
2025-08-22 11:38:24 -04:00

180 lines
7.0 KiB
Python

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
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 = 5
treat_tolerance_constant = 4
def apivalidation():
print(r"""
_____ .__ .__ __ ___________ .__
/ _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______
/ /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/
/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \
\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ >
\/ \/ \/ \/
""")
print("=== Welcome to the Airlock API Tool ===")
match os.getenv('APIKEY'):
case '':
print("Please add your API Key to the .env file")
case _:
menu_main()
def menu_main():
while True:
print("\n--- Main Menu ---")
print("1. Get All Events for Single Device")
print("2. Placeholder for Local Approval")
print("3. Placeholder for Another Tool" )
print("4. Prepare Policy For Enforcement")
print("11. Exit")
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("Invalid choice. Please try again.")
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 = [""]
history_staged = False
hashes_threat_pulled = False
hashes_categorized = False
path_exclusions_calculated = False
allowed_paths_determined = False
df_aggregated_combo = pd.DataFrame()
while True:
print("\n--- Prepare to Enforce Policy ---")
print("Sequentually follow steps to prepare for policy enforcement")
print("1. Choose first policy, Typically the audit version of the policy - Currently selected first policy is: " + str([first_policy[0]]))
print("2. Choose second policy, If an enforcement policy of that type exists, include it here - Currently selected second policy is: " + str(second_policy[0]))
print("3. Pull and stage event history for the selected policies - This step has been done - " + str(history_staged))
print("4. Pull Hash info - This step has been done - " + str(hashes_threat_pulled))
print("5. Determine if path exclusions are possible: - This step has been done - " + str(path_exclusions_calculated))
print("6. Categorize your hashes - This step has been done - " + str(hashes_categorized))
print("7. Combine Potential path exclusions and allowed hashes - This step has been done" + str(allowed_paths_determined))
print("11. Exit")
choice = input("Enter your choice: ")
if choice == "1":
first_policy = utils.allowlist.listPolicies(url)
elif choice == "2":
second_policy = utils.allowlist.listPolicies(url)
elif choice == "3":
executionhist_policy1 = utils.allowlist.pullPolicyExechistories(url,first_policy[0], first_policy[1],True)
executionhist_policy2 = utils.allowlist.pullPolicyExechistories(url,second_policy[0], second_policy[1],True)
df_aggregated_policy1 = utils.hashfunctions.aggregateHashes(executionhist_policy1)
df_aggregated_policy2 = utils.hashfunctions.aggregateHashes(executionhist_policy2)
df_aggregated_combo= pd.concat([df_aggregated_policy1, df_aggregated_policy2], ignore_index=True)
df_aggregated_combo.to_html(f"df_aggregated_combo_{first_policy}_{second_policy}.html", index=False)
history_staged = True
elif choice == "4":
if history_staged == True:
df_augmented = utils.hashfunctions.augmentAggregatedHashes(url, pd.read_html(f"df_aggregated_combo_{first_policy}_{second_policy}.html")[0])
df_augmented.to_html(f"df_augmented_combo_{first_policy}_{second_policy}.html", index=False)
hashes_threat_pulled = True
else:
print("Data not yet staged, please complete steps 1-3")
elif choice == "5":
if hashes_threat_pulled == True:
path_eligible, path_ineligible = utils.pathfunctions.filepathInitialGroup(pd.read_html(f"df_augmented_combo_{first_policy}_{second_policy}.html")[0])
path_eligible.to_html(f"df_path_eligible_{first_policy}_{second_policy}.html", index=False)
path_ineligible.to_html(f"df_path_ineligible_{first_policy}_{second_policy}.html", index=False)
path_exclusions_calculated = True
else:
print("Step 4 not complete, Please complete step 4")
elif choice == "6":
if path_exclusions_calculated == True:
categorized = utils.hashfunctions.categorizeHashes(pd.read_html(f"df_augmented_combo_{first_policy}_{second_policy}.html")[0], treat_tolerance_constant, badpublisherlist)
categorized[0].to_html(f"df_hashes_needing_approval_{first_policy}_{second_policy}.html", index=False)
categorized[1].to_html(f"df_automatically_approved_hashes_{first_policy}_{second_policy}.html", index=False)
categorized[2].to_html(f"df_remaining_hashes__{first_policy}_{second_policy}.html", index=False)
hashes_categorized = True
else:
print("Step 5 not complete, Please complete step 5")
elif choice == "7":
if hashes_categorized == True:
allowpaths = utils.allowfunctions.filter_and_drop(pd.read_html(f"df_automatically_approved_hashes_{first_policy}_{second_policy}.html")[0], pd.read_html(f"df_path_eligible_{first_policy}_{second_policy}.html")[0], path_exclusion_constant)
allowpaths.to_html(f"df_allowed_paths_{first_policy}_{second_policy}.html", index=False)
allowed_paths_determined = True
else:
print("Step 6 not complete, Please complete step 6")
elif choice == "11":
break
else:
print("Invalid choice. Please try again.")
if __name__ == "__main__":
apivalidation()