mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

First Round Async. Much work left to do, dont trust results of hash categorization presently.

Browse Source
This commit is contained in:
Zarithas
2025-10-16 16:43:56 -04:00
parent fa0c18ee02
commit 6f2355fea9
21 changed files with 903 additions and 1647 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/menus.py
+50 -48
View File
@@ -20,7 +20,7 @@ import re
import dotenv
import pandas as pd
import services.policyhandler as policyh
import services.PolicyHandler as policyh
from flows.otp import generate, otp_activities_by_agent, revoke
from flows.prepPolicy import (
buildPathsandPublishers,
@@ -32,8 +32,9 @@ from flows.prepPolicy import (
from flows.quietAgent import findQuietAgents
from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from services.TaskQueue import AsyncTaskQueue
from utils.configmanager import load_env
from utils.selector import Selector
from utils.Selector import Selector
from utils.utils import (
areYouSure,
colorText,
@@ -47,9 +48,9 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def menu_main(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
extras = load_env("EXTRAS")
async def menu_main(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
working_dir = await load_env("WORKING_DIR")
extras = await load_env("EXTRAS")
while True:
displayIntro()
# Add Settings, and give option to change working dir
@@ -63,34 +64,34 @@ def menu_main(api: AirlockAPIWrapper):
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("Q. 🔚 - Quit", "yellow"))
choice = get_sanitized_input("\nEnter Menu Item: ")
choice = await get_sanitized_input("\nEnter Menu Item: ")
if choice == "1":
print("This Feature is still in development")
get_sanitized_input("Press enter to continue")
await get_sanitized_input("Press enter to continue")
elif choice == "2":
menu_otp(api)
await menu_otp(api)
elif choice == "3":
choices = ["audit", "enforcement"]
print(colorText("Move devices to which state?:", "yellow"))
direction = Selector.select_string(choices, False, False)
devices = selectAgents(api)
direction = await Selector.select_string(choices, False, False)
devices = await selectAgents(api)
print(colorText("Would you like to continue with these devices?","white"))
for device in devices:
print(device.hostname)
confirm = Selector.confirm()
confirm =await Selector.confirm()
if direction and devices and confirm:
for device in devices:
moveAgentToRelatedPolicy(api,device, direction[0])
await moveAgentToRelatedPolicy(api,device, direction)
elif choice == "4":
findAgents(api,False)
await findAgents(api,False)
elif choice == "5":
findQuietAgents(api)
await findQuietAgents(api)
elif choice == "6":
if extras == "POLICYPREP": menu_policymanagment(api)
if extras == "POLICYPREP": await menu_policymanagment(api, queue)
elif choice.upper() == "F":
open_directory(working_dir)
await open_directory(working_dir)
elif choice.upper() == "S":
menu_settings()
await menu_settings()
elif choice.upper() == "Q":
break
else:
@@ -98,7 +99,7 @@ def menu_main(api: AirlockAPIWrapper):
def menu_policy_enforce(api: AirlockAPIWrapper):
async def menu_policy_enforce(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
selected_policies = []
destination_policy = []
destination_allowlist = []
@@ -106,34 +107,35 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
processed_hashes = []
processed_publishers = []
tested = False
working_dir = load_env("WORKING_DIR")
working_dir = await load_env("WORKING_DIR")
while True:
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
choice = get_sanitized_input("\nEnter your choice: ")
choice = await get_sanitized_input("\nEnter your choice: ")
if choice == "1":
selected_policies = selectPolicies(api,True)
selected_policies = await selectPolicies(api,True)
elif choice == "2":
print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
destination_policy = selectPolicies(api, False)
destination_policy = await selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, destination_policy, False)
destination_allowlist = await selectAllowlists(api, destination_policy, False) # pyright: ignore[reportArgumentType]
elif choice == "3":
sortHashes(
await sortHashes(
api,
queue,
selected_policies,
type=[1, 2, 6, 7],
)
elif choice == "4":
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
buildPathsandPublishers(False)
await buildPathsandPublishers(False)
else:
print("File not found. Please make sure it's saved correctly and try again.")
@@ -141,7 +143,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
f"{working_dir}\\Approved\\primary_Paths.csv"
):
buildPreflights()
await buildPreflights()
else:
print("File not found. Please make sure it's saved correctly and try again.")
@@ -209,7 +211,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
elif choice == "7":
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
confirmation = await get_sanitized_input("Type 'I AGREE' to continue: ")
if (
tested
and destination_policy
@@ -217,10 +219,10 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
and confirmation.strip() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
await api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
await api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
if processed_publishers:
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
await api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
else:
logger.error("Confirmation block failed. Reasons:")
if not tested:
@@ -233,9 +235,9 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())
elif choice.upper() == "F":
open_directory(working_dir)
await open_directory(working_dir)
elif choice.upper() == "S":
menu_settings()
await menu_settings()
elif choice.upper == "B":
break
@@ -244,7 +246,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
print(colorText("Invalid choice. Please try again.", "red"))
def menu_otp(api: AirlockAPIWrapper):
async def menu_otp(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
while True:
@@ -256,23 +258,23 @@ def menu_otp(api: AirlockAPIWrapper):
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("Enter your choice: ")
choice = await get_sanitized_input("Enter your choice: ")
if choice == "1":
otp_list = generate(api)
otp_list = await generate(api)
print(colorText(otp_list,"green"))
elif choice == "2":
otp_activities_by_agent(api)
await otp_activities_by_agent(api)
elif choice == "3":
revoke(api)
await revoke(api)
elif choice.upper() == "F":
open_directory(working_dir)
await open_directory(working_dir)
elif choice.upper() == "S":
menu_settings()
await menu_settings()
elif choice.upper() == "B":
break
def menu_policymanagment(api: AirlockAPIWrapper):
async def menu_policymanagment(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
working_dir = load_env("WORKING_DIR")
while True:
print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow"))
@@ -280,31 +282,31 @@ def menu_policymanagment(api: AirlockAPIWrapper):
print(colorText("F. 📂 - Open Working Directory", "yellow"))
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("\n Enter Menu Item: ")
choice = await get_sanitized_input("\n Enter Menu Item: ")
if choice == "1":
menu_policy_enforce(api)
await menu_policy_enforce(api, queue)
elif choice == "2":
areYouSure()
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
confirmation = await get_sanitized_input("Type 'I AGREE' to continue: ")
if confirmation.strip() == "I AGREE":
policyh.updateAuditPoliciesFromEnforcementPolices(api)
await policyh.updateAuditPoliciesFromEnforcementPolices(api)
elif choice.upper() == "F":
open_directory(working_dir)
await open_directory(working_dir)
elif choice.upper() == "S":
menu_settings()
await menu_settings()
elif choice.upper() == "B":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_settings():
async def menu_settings():
while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
print(colorText("This Feature is still in development", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("Enter your choice: ")
choice = await get_sanitized_input("Enter your choice: ")
if choice == "1":
pass #TODO ADD CHANGE WORKDIR CODE