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

UI Improvements

Browse Source
This commit is contained in:
Zarithas
2025-10-13 15:02:13 -04:00
parent 21370898a4
commit 864ae04db5
730 changed files with 10917 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/menus.py
+51 -52
View File
@@ -20,8 +20,8 @@ import re
import dotenv
import pandas as pd
import flows.localApproval as la
import services.policyhandler as policyh
from flows.otp import generate, otp_activities_by_agent, revoke
from flows.prepPolicy import (
buildPathsandPublishers,
buildPreflights,
@@ -30,32 +30,26 @@ from flows.prepPolicy import (
sortHashes,
)
from flows.quietAgent import findQuietAgents
from flows.otp import (
generate,
otp_activities_by_agent,
revoke
)
from utils.selector import Selector
from services.agenthandler import findAgents, selectAgents, moveAgentToRelatedPolicy
from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from utils.configmanager import load_env
from utils.selector import Selector
from utils.utils import (
areYouSure,
colorText,
displayIntro,
get_sanitized_input,
open_directory,
printEnforceChecklist,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def menu_main(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
extras = load_env("EXTRAS")
while True:
displayIntro()
# Add Settings, and give option to change working dir
@@ -63,31 +57,41 @@ def menu_main(api: AirlockAPIWrapper):
print(colorText("2. 🎫 - OTP", "yellow"))
print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow"))
print(colorText("4. 🔍 - Device Search", "yellow"))
print(colorText("5. 🛡️ - Policy Enforcement Tools", "yellow"))
print(colorText("5. 🔇 - Find Quiet Hosts", "yellow"))
if extras == "POLICYPREP" : print(colorText("6. 🛡️ - Policy Enforcement Tools", "yellow"))
print(colorText("F. 📂 - Open Working Directory", "yellow"))
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("Q. 🔚 - Quit", "yellow"))
choice = input(colorText("\nEnter Menu Item: ", "white"))
choice = get_sanitized_input("\nEnter Menu Item: ")
if choice == "1":
print("This Feature is still in development")
input("Press enter to continue")
get_sanitized_input("Press enter to continue")
elif choice == "2":
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)
if direction and devices:
print(colorText("Would you like to continue with these devices?","white"))
for device in devices:
print(device.hostname)
confirm = Selector.confirm()
if direction and devices and confirm:
for device in devices:
moveAgentToRelatedPolicy(api,device, direction[0])
elif choice == "4":
findAgents(api,False)
elif choice == "5":
menu_policymanagment(api)
elif choice == "F":
findQuietAgents(api)
elif choice == "6":
if extras == "POLICYPREP": menu_policymanagment(api)
elif choice.upper() == "F":
open_directory(working_dir)
elif choice == "S":
elif choice.upper() == "S":
menu_settings()
elif choice == "Q":
elif choice.upper() == "Q":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
@@ -106,8 +110,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
while True:
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
choice = input(colorText("\nEnter your choice: ", "white"))
choice = get_sanitized_input("\nEnter your choice: ")
if choice == "1":
selected_policies = selectPolicies(api,True)
@@ -206,13 +209,12 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
elif choice == "7":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if (
tested
and destination_policy
and destination_allowlist
and confirmation.strip().upper() == "I AGREE"
and confirmation.strip() == "I AGREE"
):
print(colorText("Proceeding with the code...", "yellow"))
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
@@ -227,14 +229,14 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
logger.error(" - `destination_policy` is missing or invalid.")
if not destination_allowlist:
logger.error(" - `destination_allowlist` is missing or invalid.")
if confirmation.strip().upper() != "I AGREE":
if confirmation.strip() != "I AGREE":
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())
elif choice == "F":
elif choice.upper() == "F":
open_directory(working_dir)
elif choice == "S":
elif choice.upper() == "S":
menu_settings()
elif choice == "Q":
elif choice.upper == "B":
break
@@ -252,9 +254,9 @@ def menu_otp(api: AirlockAPIWrapper):
print(colorText("3. ❌ -Revoke OTPs", "cyan"))
print(colorText("F. 📂 - Open Working Directory", "yellow"))
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("Q. 🔚 - Quit", "yellow"))
print(colorText("B. 🔙 - Back", "yellow"))
choice = input("Enter your choice: ")
choice = get_sanitized_input("Enter your choice: ")
if choice == "1":
otp_list = generate(api)
@@ -263,54 +265,51 @@ def menu_otp(api: AirlockAPIWrapper):
otp_activities_by_agent(api)
elif choice == "3":
revoke(api)
elif choice == "F":
elif choice.upper() == "F":
open_directory(working_dir)
elif choice == "S":
elif choice.upper() == "S":
menu_settings()
elif choice == "Q":
elif choice.upper() == "B":
break
def menu_policymanagment(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
while True:
print(colorText("1. 🔇 - Find Quiet Hosts", "yellow"))
print(colorText("2. 🔒 - Prepare Policy For Enforcement", "yellow"))
print(colorText("3. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow"))
print(colorText("2. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
print(colorText("F. 📂 - Open Working Directory", "yellow"))
print(colorText("S. 🛠️ - Settings", "yellow"))
print(colorText("Q. 🔚 - Quit", "yellow"))
choice = input(colorText("\nEnter Menu Item: ", "white"))
print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("\n Enter Menu Item: ")
if choice == "1":
findQuietAgents(api)
elif choice == "2":
menu_policy_enforce(api)
elif choice == "3":
elif choice == "2":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
if confirmation.strip().upper() == "I AGREE":
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
if confirmation.strip() == "I AGREE":
policyh.updateAuditPoliciesFromEnforcementPolices(api)
elif choice == "F":
elif choice.upper() == "F":
open_directory(working_dir)
elif choice == "S":
elif choice.upper() == "S":
menu_settings()
elif choice == "Q":
elif choice.upper() == "B":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_settings():
while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
print(colorText("1. Change Working Dir", "cyan"))
print(colorText("This Feature is still in development", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("Q. Return to Main Menu", "cyan"))
choice = input("Enter your choice: ")
print(colorText("B. 🔙 - Back", "yellow"))
choice = get_sanitized_input("Enter your choice: ")
if choice == "1":
pass #TODO ADD CHANGE WORKDIR CODE
elif choice == "Q":
elif choice.upper() == "B":
print("Returning to Main Menu...")
break
else:
utils/selector.py
+7 -5
View File
@@ -17,11 +17,13 @@
import logging
from typing import Any, List, Optional, Union
from utils.utils import get_sanitized_input
logger = logging.getLogger(__name__)
from typing import List, Optional, Union, Any, Callable
import logging
from typing import Callable
logger = logging.getLogger(__name__)
@@ -63,7 +65,7 @@ class Selector:
if allow_multiple:
while True:
choice = input("Select an item by number (or Q to finish): ").strip().lower()
choice = get_sanitized_input("Select an item by number (or Q to finish): ").strip().lower()
if choice == "q":
break
try:
@@ -83,7 +85,7 @@ class Selector:
return selected if selected else None
else:
try:
choice = int(input("Select one item by number: "))
choice = int(get_sanitized_input("Select one item by number: "))
if 1 <= choice <= len(sorted_items):
selected_item = sorted_items[choice - 1]
logger.info(f"Selected: {label_func(selected_item)}")
@@ -144,7 +146,7 @@ class Selector:
allow_quit: bool = False
) -> Optional[Any]:
while True:
user_input = input(prompt).strip().lower()
user_input = get_sanitized_input(prompt).strip().lower()
if allow_quit and user_input == "q":
logger.info("User opted to quit value selection.")
return None
@@ -163,7 +165,7 @@ class Selector:
@staticmethod
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
while True:
response = input(prompt).strip().lower()
response = get_sanitized_input(prompt).strip().lower()
if response in ["y", "yes"]:
logger.info("User confirmed action.")
return True
utils/setup.py
+2 -1
View File
@@ -20,10 +20,12 @@ import os
import platform
import sys
from pathlib import Path
from dotenv import load_dotenv, set_key
from utils.configmanager import PROTECTED_KEYS, load_protected_config
def get_base_directory() -> Path:
system = platform.system()
home = Path.home()
@@ -100,7 +102,6 @@ def load_user_config(config_dir: Path) -> dict:
return json.load(f)
def write_config_to_env(config: dict, env_path: Path):
from utils.configmanager import PROTECTED_KEYS
for key, value in config.items():
if key in PROTECTED_KEYS:
continue # Skip protected keys
utils/utils.py
+13 -28
View File
@@ -21,12 +21,12 @@ import re
import subprocess
import tempfile
import tkinter as tk
from tkinter import filedialog, messagebox, simpledialog
from utils.configmanager import load_env
from tkinter import filedialog
import pandas as pd
from utils.configmanager import load_env
logger = logging.getLogger(__name__)
@@ -98,32 +98,17 @@ def choose_file(initial_directory=None, required_substring=None):
return file_path
def choose_save_location():
root = tk.Tk()
root.withdraw()
save_path = filedialog.asksaveasfilename(defaultextension=".txt")
return save_path
def ask_user_input(message):
root = tk.Tk()
root.withdraw()
user_input = simpledialog.askstring("Input", "{message}]:")
return user_input
def show_info_message_cli(title, message):
root = tk.Tk()
root.withdraw()
root.after(100, lambda: messagebox.showinfo(title, message))
root.mainloop()
def show_confirm_question():
root = tk.Tk()
root.withdraw()
response = messagebox.askquestion("Confirm", "Do you want to continue?")
print("User response:", response)
def get_sanitized_input(prompt: str) -> str:
while True:
user_input = input(prompt)
if user_input.strip() == "":
return user_input # Allow blank lines
if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()):
return user_input
else:
print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.")
def regulator(paths, case_insensitive=True):
@@ -565,7 +550,7 @@ def printDeviceEnforceChecklist():
)
)
print(colorText("Q. Quit", "cyan"))
print(colorText("B. Back", "cyan"))
def colorText(text, color):