Files
AirlockTools/utils/menus.py
T

276 lines
10 KiB
Python

# Copyright (C) 2025 James Brotosky, Brandon Wickline
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import logging
import os
import re
import dotenv
import pandas as pd
import flows.localApproval as la
import services.policyhandler as policyh
from flows.prepPolicy import (
buildPathsandPublishers,
buildPreflights,
selectAllowlists,
selectPolicies,
sortHashes,
)
from flows.quietAgent import findQuietAgents
from services.agenthandler import findAgents
from services.API import AirlockAPIWrapper
from utils.utils import (
areYouSure,
colorText,
displayIntro,
load_env,
open_directory,
printEnforceChecklist,
)
logger = logging.getLogger(__name__)
dotenv.load_dotenv()
def menu_main(api: AirlockAPIWrapper):
working_dir = load_env("WORKING_DIR")
while True:
displayIntro()
# Add Settings, and give option to change working dir
print(colorText("1. ➡️ - Move Device(s) to local approval", "yellow"))
print(colorText("2. 🎫 - OTP", "yellow"))
print(colorText("3. 🔍 - Device Search", "yellow"))
print(colorText("4. 🔇 - Find Quiet Hosts", "yellow"))
print(colorText("5. 🔒 - Prepare Policy For Enforcement", "yellow"))
print(colorText("6. 🔄 - 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"))
if choice == "1":
la.moveToLocalApproval(api)
elif choice == "2":
menu_otp(api)
elif choice == "3":
findAgents(api,False)
elif choice == "4":
findQuietAgents(api)
elif choice == "5":
menu_policy_enforce(api)
elif choice == "6":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
if confirmation.strip().upper() == "I AGREE":
policyh.updateAuditPoliciesFromEnforcementPolices(api)
elif choice == "F":
open_directory(working_dir)
elif choice == "S":
menu_settings()
elif choice == "Q":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_policy_enforce(api: AirlockAPIWrapper):
selected_policies = []
destination_policy = []
destination_allowlist = []
processed_paths = []
processed_hashes = []
processed_publishers = []
tested = False
working_dir = load_env("WORKING_DIR")
while True:
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
choice = input(colorText("\nEnter your choice: ", "white"))
if choice == "1":
selected_policies = selectPolicies(api,True)
elif choice == "2":
print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
destination_policy = selectPolicies(api, False)
print(colorText("Please choose Allowlist for Hashes", "white"))
destination_allowlist = selectAllowlists(api, destination_policy, False)
elif choice == "3":
sortHashes(
api,
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)
else:
print("File not found. Please make sure it's saved correctly and try again.")
elif choice == "5":
if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
f"{working_dir}\\Approved\\primary_Paths.csv"
):
buildPreflights()
else:
print("File not found. Please make sure it's saved correctly and try again.")
elif choice == "6":
if (
os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv")
and os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv")
and destination_policy
and destination_allowlist
):
print(colorText("These path exclusions would be added to:", "yellow"))
print(destination_policy)
pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\approved_paths.csv")
hashes = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates()
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
processed_paths = [
(path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}"
for path, ext in unique_combinations.itertuples(index=False, name=None)
]
print(processed_paths)
print(colorText("These publishers would added", "yellow"))
if os.path.exists(f"{working_dir}\\Preflight\\approved_publishers.csv"):
publishers = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv")
if publishers.empty:
print(colorText("The publishers list is empty.", "red"))
else:
processed_publishers = (
publishers[publishers["publisher_hash"] != "Not Signed"]
["publisher_hash"]
.drop_duplicates()
.tolist()
)
print(processed_publishers)
print(colorText("These hashes would be added to:", "yellow"))
print(destination_allowlist)
processed_hashes = hashes["sha256"].unique().tolist()
print(processed_hashes)
if processed_paths and processed_hashes:
tested = True
else:
# Log which condition(s) failed
missing_items = []
if not os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv"):
missing_items.append("approved_paths.csv not found")
if not os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv"):
missing_items.append("approved_hashes.csv not found")
if not destination_policy:
missing_items.append("destination_policy is empty or None")
if not destination_allowlist:
missing_items.append("destination_allowlist is empty or None")
logger.error("Preflight check failed due to the following:")
for item in missing_items:
logger.error(f" - {item}")
elif choice == "7":
areYouSure()
confirmation = input(colorText("Type 'I AGREE' to continue: ", "white"))
if (
tested
and destination_policy
and destination_allowlist
and confirmation.strip().upper() == "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)
if processed_publishers:
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
else:
logger.error("Confirmation block failed. Reasons:")
if not tested:
logger.error(" - Preflight checks were not completed successfully (`tested` is False).")
if not destination_policy:
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":
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())
elif choice == "F":
open_directory(working_dir)
elif choice == "S":
menu_settings()
elif choice == "Q":
break
else:
print(colorText("Invalid choice. Please try again.", "red"))
def menu_otp(api: AirlockAPIWrapper):
while True:
print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
print(colorText("1. Generate OTP", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("Q. Return to Main Menu", "cyan"))
choice = input("Enter your choice: ")
if choice == "1":
# TODO generateOTP(api,findAgents()
break
elif choice == "2":
print("You selected Sub-option B")
elif choice == "Q":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")
def menu_settings():
while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
print(colorText("1. Change Working Dir", "cyan"))
# print(colorText("2. Sub-option B","cyan"))
print(colorText("Q. Return to Main Menu", "cyan"))
choice = input("Enter your choice: ")
if choice == "1":
pass #TODO ADD CHANGE WORKDIR CODE
elif choice == "Q":
print("Returning to Main Menu...")
break
else:
print("Invalid choice. Please try again.")