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

Cleaning up menu, added Move to Audit/Enforcement

Browse Source
This commit is contained in:
Zarithas
2025-10-10 17:20:07 -04:00
parent b74f77a9db
commit 21370898a4
13 changed files with 213 additions and 468 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/configmanager.py
+115
View File
@@ -0,0 +1,115 @@
import json
import logging
import os
import sys
from pathlib import Path
from typing import Callable, Optional, TypeVar
T = TypeVar("T")
logger = logging.getLogger(__name__)
PROTECTED_KEYS = [
"APPNAME",
"LOG_LEVEL",
"PATH_EXCLUSION_CONST",
"MIN_FILES_FOR_PATH",
"VT_THREAT_TOLERANCE",
"POLICY_MAP_ENF_AUD"
]
_protected_config = {}
def get_system_config_path() -> Path:
# Check inside bundled EXE directory first
bundled_dir = Path(getattr(sys, '_MEIPASS', ''))
bundled_path = bundled_dir / "system_config.json"
if bundled_path.exists():
return bundled_path
# Fallback to external location
return Path(__file__).parent.parent / "system_config.json"
def load_protected_config() -> dict:
global _protected_config
try:
with open(get_system_config_path(), "r") as f:
system_config = json.load(f)
except FileNotFoundError:
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
system_config = {
"APPNAME": "AirlockTools",
"PATH_EXCLUSION_CONST": 4,
"MIN_FILES_FOR_PATH": 4,
"VT_THREAT_TOLERANCE": 4,
"POLICY_MAP_ENF_AUD": {
"enforced_id": "audit_id"
}
}
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
return _protected_config
def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
value = _protected_config.get(key)
if value is None:
logging.warning(f"Protected config key '{key}' not found.")
return default
try:
if isinstance(value, str):
value = value.strip("'\"")
return cast_type(value)
except (ValueError, TypeError):
logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.")
return default
def get_protected_json(key: str, default: str = "{}") -> dict:
raw = _protected_config.get(key, default)
if isinstance(raw, dict):
return raw
try:
return json.loads(raw)
except json.JSONDecodeError:
try:
escaped = raw.encode('unicode_escape').decode('utf-8')
return json.loads(escaped)
except Exception as e:
logging.error(f"Failed to parse protected JSON key '{key}': {e}")
return json.loads(default)
def load_env_json(key: str, default: str):
raw = os.getenv(key, default)
try:
return json.loads(raw)
except json.JSONDecodeError:
try:
escaped = raw.encode('unicode_escape').decode('utf-8')
return json.loads(escaped)
except Exception as e:
logging.error(f"Failed to parse {key}: {e}")
return json.loads(default)
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
"""
Safely retrieves an environment variable and casts it to the desired type.
Parameters:
key (str): The name of the environment variable.
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
default (Optional[T], optional): Default value if the variable is not set or invalid.
Returns:
Optional[T]: The casted value or the default.
"""
value = os.getenv(key)
if value is None:
logger.warning(f"Environment variable '{key}' not set.")
return default
try:
value = value.strip("'\"") # Strip surrounding quotes
return cast_type(value)
except (ValueError, TypeError):
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
return default
utils/menus.py
+47 -23
View File
@@ -36,14 +36,14 @@ from flows.otp import (
otp_activities_by_agent,
revoke
)
from services.agenthandler import findAgents
from utils.selector import Selector
from services.agenthandler import findAgents, selectAgents, moveAgentToRelatedPolicy
from services.API import AirlockAPIWrapper
from utils.configmanager import load_env
from utils.utils import (
areYouSure,
colorText,
displayIntro,
load_env,
open_directory,
printEnforceChecklist,
)
@@ -59,34 +59,30 @@ def menu_main(api: AirlockAPIWrapper):
while True:
displayIntro()
# Add Settings, and give option to change working dir
print(colorText("1. ➡️ - Move Device(s) to local approval", "yellow"))
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"))
print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow"))
print(colorText("4. 🔍 - Device Search", "yellow"))
print(colorText("5. 🛡️ - Policy Enforcement Tools", "yellow"))
choice = input(colorText("\nEnter Menu Item: ", "white"))
if choice == "1":
la.moveToLocalApproval(api)
print("This Feature is still in development")
input("Press enter to continue")
elif choice == "2":
menu_otp(api)
elif choice == "3":
findAgents(api,False)
choices = ["audit", "enforcement"]
direction = Selector.select_string(choices, False, False)
devices = selectAgents(api)
if direction and devices:
for device in devices:
moveAgentToRelatedPolicy(api,device, direction[0])
elif choice == "4":
findQuietAgents(api)
findAgents(api,False)
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)
menu_policymanagment(api)
elif choice == "F":
open_directory(working_dir)
elif choice == "S":
@@ -97,6 +93,7 @@ def menu_main(api: AirlockAPIWrapper):
print(colorText("Invalid choice. Please try again.", "red"))
def menu_policy_enforce(api: AirlockAPIWrapper):
selected_policies = []
destination_policy = []
@@ -273,8 +270,35 @@ def menu_otp(api: AirlockAPIWrapper):
elif choice == "Q":
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("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":
findQuietAgents(api)
elif choice == "2":
menu_policy_enforce(api)
elif choice == "3":
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_settings():
while True:
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
utils/setup.py
+18 -26
View File
@@ -1,19 +1,17 @@
# Copyright (C) 2025 James Brotosky, Brandon Wickline
# 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 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.
# 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/>.
# 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 json
import logging
@@ -22,16 +20,9 @@ import os
import platform
import sys
from pathlib import Path
from dotenv import load_dotenv, set_key
PROTECTED_KEYS = [
"APPNAME",
"PATH_EXCLUSION_CONST",
"MIN_FILES_FOR_PATH",
"VT_THREAT_TOLERANCE",
"POLICY_MAP_ENF_AUD"
]
from utils.configmanager import PROTECTED_KEYS, load_protected_config
def get_base_directory() -> Path:
system = platform.system()
@@ -73,7 +64,6 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
logger.debug("✅ Logging configured.")
def get_system_config_path() -> Path:
base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__))))
return base_path.parent / "system_config.json"
@@ -110,7 +100,10 @@ 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
try:
serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value)
set_key(env_path, key, serialized)
@@ -157,13 +150,13 @@ def setup() -> Path:
for subfolder in subfolders:
subfolder_path = folder_path / subfolder
subfolder_path.mkdir(parents=True, exist_ok=True)
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
logging.debug(f" └─ '{subfolder}' subfolder created at: {subfolder_path}")
user_config = load_user_config(dirs['config'])
merged_config = {**system_config, **user_config}
for key in PROTECTED_KEYS:
merged_config[key] = system_config.get(key, "")
protected_config = load_protected_config()
merged_config.update(protected_config)
# ✅ URL resolution order: system_config → .env → user prompt
url = system_config.get("URL")
@@ -171,7 +164,6 @@ def setup() -> Path:
url = os.getenv("URL")
if not url:
url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip()
merged_config["URL"] = url
set_key(env_path, "URL", url)
os.environ["URL"] = url
utils/utils.py
+3 -38
View File
@@ -13,7 +13,7 @@
# 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 json
import logging
import os
import platform
@@ -22,48 +22,14 @@ import subprocess
import tempfile
import tkinter as tk
from tkinter import filedialog, messagebox, simpledialog
from typing import Callable, Optional, TypeVar
from utils.configmanager import load_env
import pandas as pd
logger = logging.getLogger(__name__)
T = TypeVar("T")
def load_env_json(key: str, default: str):
raw = os.getenv(key, default)
try:
return json.loads(raw)
except json.JSONDecodeError:
try:
escaped = raw.encode('unicode_escape').decode('utf-8')
return json.loads(escaped)
except Exception as e:
logging.error(f"Failed to parse {key}: {e}")
return json.loads(default)
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
"""
Safely retrieves an environment variable and casts it to the desired type.
Parameters:
key (str): The name of the environment variable.
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
default (Optional[T], optional): Default value if the variable is not set or invalid.
Returns:
Optional[T]: The casted value or the default.
"""
value = os.getenv(key)
if value is None:
logger.warning(f"Environment variable '{key}' not set.")
return default
try:
value = value.strip("'\"") # Strip surrounding quotes
return cast_type(value)
except (ValueError, TypeError):
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
return default
def import_to_dataframe(file_path: str) -> pd.DataFrame:
@@ -171,7 +137,6 @@ def regulator(paths, case_insensitive=True):
print(f"Regulator is providing: {pattern}")
return pattern
def displayIntro():
print(
colorText(