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:
2025-10-05 22:20:42 -04:00
parent b1e6af01c6
commit 89db386ffe
26 changed files with 3530 additions and 2389 deletions
+77 -107
View File
@@ -12,57 +12,80 @@
#
# 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/>.
#TODO Continue implementing logger
#TODO Add input sanitation and CSV injection prevention
#TODO Continue OTP and Local approval rewrites
#TODO Explore pywin32
#TODO Fix Requirements.txt
#TODO Create Generic system_config.json for gitea
import argparse
import dotenv
import json
import logging
import os
import dotenv
import urllib3
import utils.clientfunctions
import utils.localapproval as la
import utils.otpfunctions
import utils.policyfunctions
import utils.utils as ct
import pandas as pd
import utils.menus as menus
from services.API import AirlockAPIWrapper
from services.security import getAPI
from services.setup import setup
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler
urllib3.disable_warnings(
urllib3.exceptions.InsecureRequestWarning
)
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
dotenv.load_dotenv()
#Constants
url = os.getenv('url')
bad_publisher_list = ["Brave","Zoom", "GlavSoft", "VNC"]
pups = ["logmein", "invalid", "nmap", "LTSvc", "VNC", "Kaseya", "Solarwinds", "mRemoteNG"]
badpathparts = ["users", "wwwroot", "windows\\temp", "windows\\task", "windows\\system32", "startup", "windows\\fonts", "Recycle.Bin", "AppData", "programdata", "Solarwinds", "kaseya"]
path_exclusion_constant = 4
min_files_for_path = 4
threat_tolerance_constant = 4
policy_relationship_map ={ #Enforcement : Audit
"bf0b1f9b-bfea-4f44-97c0-80e27ff61712" : "538d3218-92f4-4943-a6ee-db9267ab62d8", #AT Servers General, AT Servers General Audit
"31ababac-65de-4c6a-86dd-6691d7e3ee3b" : "fc05b42a-b846-4e72-88ca-c35d416e699f", #AT Epic, #AT Epic Audit
"d1f58960-f866-49e0-848a-a5b09fffd4cd" : "d55c03a6-c376-4391-8626-4f843b882a7c", #AT DMZ Enforced, #AT DMZ Audit
"504dd011-86b6-489a-b78f-eff589cef8aa" : "88a1cfdc-3b30-448b-b309-be16fe437ca3", #AT Workstations BCA, #AT Workstations BCA Audit
"d126db36-72ed-4937-adc7-d88b7509a5b5" : "5aebf6a0-1d67-47b4-9c5f-2866ffca5671" #AT Testing, AT Testing
}
def main():
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument('--monitorOTP', action='store_true', help='Run in non-interactive mode')
#Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored
working_dir = setup()
logger = logging.getLogger(__name__)
dotenv.load_dotenv(dotenv_path=working_dir / ".env")
try:
url = os.getenv("URL")
username = os.getenv("USERNAME")
if not url:
raise ValueError("Missing URL in environment variables.")
if not username:
raise ValueError("Missing USERNAME in environment variables.")
logger.debug(f"Retrieved URL: {url}")
logger.debug(f"Retrieved Username: {username}")
except ValueError as e:
logger.error(f"Configuration error: {e}", exc_info=True)
raise
api = AirlockAPIWrapper(
base_url=str(os.getenv("URL")),
api_key = getAPI(username, "AirlockTools"),
)
parser = argparse.ArgumentParser(description="Program for Managing Airlock via API & CMD")
parser.add_argument("--monitor", action="store_true", help="Run in non-interactive mode")
# Add other arguments as needed
args = parser.parse_args()
if args.monitorOTP:
if args.monitor:
# Non-interactive logic
print(f"Running non-interactively to start monitoring OTP")
logger.info("Running non-interactively to start monitoring Airlock Changes")
"""
os.makedirs("scheduling", exist_ok=True)
os.makedirs("OTP/HTML", exist_ok=True)
os.makedirs("OTP/PARQ", exist_ok=True)
os.makedirs("Local_Approval/HTML", exist_ok=True)
os.makedirs("Local_Approval/PARQ", exist_ok=True)
ct.apivalidation()
register_function("monitorOTP", utils.otpfunctions.monitorOTP)
register_function("monitorLA", la.scheduleAddingLAHashes)
register_function("updateAuditPolicies", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices)
@@ -75,80 +98,27 @@ def main():
else:
reload_jobs()
start_scheduler()
"""
else:
# Interactive logic
ct.apivalidation()
menu_main()
def menu_main():
while True:
ct.displayIntro();
print(ct.colorText("1. 🖥️ - Get All Events for Single Device", "yellow"))
print(ct.colorText("2. 🎫 - OTP", "yellow"))
print(ct.colorText("3. 🔇 - Find Quiet Hosts", "yellow"))
print(ct.colorText("4. 🔒 - Prepare Policy For Enforcement", "yellow"))
print(ct.colorText("5. 🔄 - Update Audit Policies from Enforcement Policies", "yellow"))
print(ct.colorText("6. 🔍 - Device Search", "yellow"))
print(ct.colorText("7. ➡️ - Move Devices to Local Approval", "yellow"))
print(ct.colorText("Q. 🔚 - Quit", "yellow"))
choice = input(ct.colorText("\nEnter Menu Item: ", "white"))
if choice == '1':
utils.clientfunctions.devicehistory(url,False)
elif choice == "2":
menu_otp()
elif choice == "3":
utils.clientfunctions.findQuietAgents(url)
elif choice == "4":
utils.policyfunctions.prepare_to_enforce(url, bad_publisher_list, pups, badpathparts, threat_tolerance_constant, path_exclusion_constant, min_files_for_path)
elif choice == "5":
ct.areYouSure()
confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white"))
if confirmation.strip().upper() == "I AGREE": utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices(url, policy_relationship_map)
elif choice == "6":
devicelist = utils.clientfunctions.promptForDevices()
utils.clientfunctions.findAgents(url, devicelist, False)
elif choice == "7":
la.moveToLocalApproval(url, policy_relationship_map)
elif choice == "8":
las = la.getLocalApprovals(url)
las.to_csv("la.csv", index=False)
elif choice == "9":
pass
elif choice == "10":
utils.clientfunctions.moveAgentToAudit(url,"6a221ece-0c10-4eb8-b1e5-06a1000a5696",policy_relationship_map)
elif choice == "11":
utils.clientfunctions.moveAgentToEnforcement(url,"6a221ece-0c10-4eb8-b1e5-06a1000a5696",policy_relationship_map)
elif choice == "Q":
break
else:
print(ct.colorText("Invalid choice. Please try again.","red"))
def menu_otp():
while True:
print(ct.colorText("\n--- 🎫 OTP Submenu 🎫 ---","cyan"))
print(ct.colorText("1. Generate OTP","cyan"))
#print(ct.colorText("2. Sub-option B","cyan"))
print(ct.colorText("Q. Return to Main Menu","cyan"))
choice = input("Enter your choice: ")
if choice == "1":
utils.otpfunctions.generateOTP(url, utils.clientfunctions.findAgentID(url))
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.")
# Interactive logic
raw = os.getenv("POLICY_MAP_ENF_AUD", "{}")
try:
# Escape backslashes before parsing
escaped = raw.encode('unicode_escape').decode('utf-8')
badpathparts = json.loads(escaped)
except Exception as e:
logging.error(f"Failed to parse BAD_PATH_PARTS: {e}")
badpathparts = []
print(badpathparts)
menus.menu_main(api)
if __name__ == "__main__":
main()
main()