Major refactor: security enhancements, modularization, config integration, reduced Parquet reliance
- Migrated codebase to class-based architecture for better modularity and maintainability - Introduced system_config.json for centralized configuration (required for runtime) - Added structured working directories for improved file organization - Significantly reduced reliance on Parquet; replaced with alternative data handling - Implemented security improvements across modules - Several TODOs remain in the main script for future enhancements - Linter formatting affected readability in some files (e.g., utils); cleanup is on the agenda
This commit is contained in:
+258
@@ -0,0 +1,258 @@
|
||||
# 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, 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")
|
||||
|
||||
# Get unique combinations of longestcfp and file_extension
|
||||
unique_combinations = pathexclusions[
|
||||
["longestcfp", "file_extension"]
|
||||
].drop_duplicates()
|
||||
|
||||
# Regex to match a Windows drive letter at the start (e.g., C:\)
|
||||
drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\")
|
||||
|
||||
# Build processed paths like \\path\\**.exe or C:\path\**.jar
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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.")
|
||||
Reference in New Issue
Block a user