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
.gitignore
+1
View File
@@ -10,3 +10,4 @@ jobs.json
securitytest.py
*.toml
system_config.json
Devel_unused/
AirlockToolServer.py
+4 -4
View File
@@ -64,11 +64,11 @@ def main():
logger.error(f"Configuration error: {e}", exc_info=True)
raise
api = AirlockAPIWrapper(
base_url=str(os.getenv("URL")),
api_key = getAPI(username, "AirlockTools"),
)
base_url=str(os.getenv("URL")),
api_key = getAPI(username, "AirlockTools"),
)
logger.info("Running non-interactively to start monitoring Airlock Changes")
Server/scheduler.py
-352
View File
@@ -1,352 +0,0 @@
# 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 json
import logging
import os
import time
from typing import Any, Callable, Dict, List
import schedule
logger = logging.getLogger(__name__)
# TODO - move this File where all jobs are persisted
JOBS_FILE = f"{os.getenv('WORKING_DIR')}\\scheduling\\jobs.json"
# Ensure directory exists
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)
# Registry of functions that can be scheduled
FUNCTION_MAP: Dict[str, Callable] = {}
# -------------------------------
# Function Registration
# -------------------------------
def register_function(name: str, func: Callable):
"""
Register a function so it can be called by name later.
Example:
register_function("say_hello", say_hello)
"""
FUNCTION_MAP[name] = func
# -------------------------------
# Persistence Helpers
# -------------------------------
def load_jobs() -> List[Dict[str, Any]]:
"""Load jobs from the JSON file, or return [] if none exist."""
if not os.path.exists(JOBS_FILE):
return []
with open(JOBS_FILE, "r") as f:
return json.load(f)
def _atomic_save(path: str, data: Any):
"""Write JSON atomically to avoid partial writes."""
tmp = f"{path}.tmp"
with open(tmp, "w") as f:
json.dump(data, f, indent=4)
os.replace(tmp, path)
def save_jobs(jobs: List[Dict[str, Any]]):
"""Save jobs to the JSON file (overwrite)."""
_atomic_save(JOBS_FILE, jobs)
# -------------------------------
# Uniqueness Helpers
# -------------------------------
def job_in_store(job_id: str) -> bool:
"""Check if a job id exists in the persisted JSON file."""
return any(j.get("id") == job_id for j in load_jobs())
def job_in_scheduler(job_id: str) -> bool:
"""
Check if a job with this tag exists in the in-memory scheduler.
Uses schedule.get_jobs(tag=...) if available, otherwise scans tags.
"""
try:
jobs = schedule.get_jobs(tag=job_id) # schedule >= 1.2.0
return len(jobs) > 0
except TypeError:
# Fallback for older versions
return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)
def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
"""
Ensure the job_id is unique across persistence and in-memory schedule.
on_conflict:
- "error": raise ValueError if exists.
- "skip" : print and return False.
- "replace": remove existing (in-memory + JSON), then continue.
"""
exists = job_in_store(job_id) or job_in_scheduler(job_id)
if not exists:
return True
if on_conflict == "error":
raise ValueError(f"Job id '{job_id}' already exists.")
elif on_conflict == "skip":
logger.info(f"Job '{job_id}' already exists. Skipping creation.")
return False
elif on_conflict == "replace":
# Clear from scheduler
schedule.clear(job_id)
# Remove from persistence
jobs = [j for j in load_jobs() if j.get("id") != job_id]
save_jobs(jobs)
return True
else:
raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")
# -------------------------------
# Internal scheduling (no persistence)
# -------------------------------
def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
args = args or []
kwargs = kwargs or {}
def job_wrapper():
"""Executes the job once, then removes it."""
if func_name not in FUNCTION_MAP:
logger.error(f"Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
# Remove from persistence
jobs = load_jobs()
jobs = [j for j in jobs if j["id"] != job_id]
save_jobs(jobs)
# Clear from in-memory schedule
schedule.clear(job_id)
delay_seconds = run_at_timestamp - time.time()
if delay_seconds <= 0:
logger.info(f"Job {job_id} scheduled in the past. Skipping.")
return
# Schedule via schedule library
schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)
def _schedule_recurring(
job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None
):
args = args or []
kwargs = kwargs or {}
def job_wrapper():
if func_name not in FUNCTION_MAP:
logger.error(f"Function '{func_name}' is not registered.")
return
FUNCTION_MAP[func_name](*args, **kwargs)
if unit == "seconds":
schedule.every(interval).seconds.do(job_wrapper).tag(job_id)
elif unit == "minutes":
schedule.every(interval).minutes.do(job_wrapper).tag(job_id)
elif unit == "hours":
schedule.every(interval).hours.do(job_wrapper).tag(job_id)
elif unit == "days":
schedule.every(interval).days.do(job_wrapper).tag(job_id)
else:
raise ValueError(f"Unsupported unit: {unit}")
# -------------------------------
# Public APIs (with uniqueness + persistence)
# -------------------------------
def run_once_job(
job_id: str,
func_name: str,
run_at_timestamp: float,
args=None,
kwargs=None,
*,
replace: bool = False,
persist: bool = True,
):
"""
Schedule a job to run once at a specific timestamp.
replace: if True, replace existing job with same id; otherwise print and skip.
persist: if False, do not write to JSON (used by reload_jobs()).
"""
if persist:
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
elif job_in_scheduler(job_id):
schedule.clear(job_id)
_schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
jobs.append(
{
"id": job_id,
"type": "once",
"run_at": run_at_timestamp,
"function": func_name,
"args": args or [],
"kwargs": kwargs or {},
}
)
save_jobs(jobs)
def recurring_job(
job_id: str,
func_name: str,
interval: int,
unit: str,
args=None,
kwargs=None,
*,
replace: bool = False,
persist: bool = True,
):
"""
Schedule a recurring job.
replace: if True, replace existing job with same id; otherwise print and skip.
persist: if False, do not write to JSON (used by reload_jobs()).
"""
if persist:
policy = "replace" if replace else "skip"
if not ensure_unique(job_id, on_conflict=policy):
return
elif job_in_scheduler(job_id):
schedule.clear(job_id)
_schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
if persist:
jobs = [j for j in load_jobs() if j["id"] != job_id]
jobs.append(
{
"id": job_id,
"type": "recurring",
"interval": interval,
"unit": unit,
"function": func_name,
"args": args or [],
"kwargs": kwargs or {},
}
)
save_jobs(jobs)
def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0):
"""
Find all jobs whose ID contains the given PID substring and reschedule them to run sooner.
"""
jobs = load_jobs()
matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]
if not matched_jobs:
logger.info(f"No jobs found containing PID substring '{pid_substring}'.")
return
logger.info(f"Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
for job in matched_jobs:
job_id = job["id"]
logger.debug(f" - Prioritizing job: {job_id}")
# Clear existing job from scheduler
schedule.clear(job_id)
# Reschedule based on job type
if job["type"] == "once":
run_once_job(
job_id,
job["function"],
time.time() + new_delay_seconds,
job.get("args"),
job.get("kwargs"),
replace=True,
persist=True,
)
elif job["type"] == "recurring":
recurring_job(
job_id,
job["function"],
job["interval"],
job["unit"],
job.get("args"),
job.get("kwargs"),
replace=True,
persist=True,
)
else:
logger.warning(f"Unknown job type for job '{job_id}'")
# -------------------------------
# Reload Saved Jobs
# -------------------------------
def reload_jobs():
"""Reload jobs from JSON and reschedule them (no re-persist)."""
jobs = load_jobs()
for job in jobs:
if job["type"] == "once":
if job["run_at"] > time.time():
run_once_job(
job["id"],
job["function"],
job["run_at"],
job.get("args"),
job.get("kwargs"),
persist=False,
)
elif job["type"] == "recurring":
recurring_job(
job["id"],
job["function"],
job["interval"],
job["unit"],
job.get("args"),
job.get("kwargs"),
persist=False,
)
# -------------------------------
# Scheduler Loop
# -------------------------------
def start_scheduler():
"""
Start the scheduler loop (blocking).
Call this once in main to begin.
"""
try:
while True:
schedule.run_pending()
time.sleep(0.5)
except KeyboardInterrupt:
logger.critical("Scheduler stopped.")
flows/localApproval.py
+6 -8
View File
@@ -27,11 +27,9 @@ from utils.setup import get_base_directory
from models.agent import Agent
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from Server.scheduler import (
register_function,
run_once_job,
)
from utils.utils import colorText, load_env, load_env_json
from utils.utils import colorText
from utils.configmanager import get_protected_json, load_env, load_env_json
logger = logging.getLogger(__name__)
@@ -68,7 +66,7 @@ def getLocalApprovals(api: AirlockAPIWrapper):
def scheduleAddingLAHashes(api: AirlockAPIWrapper):
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int)
@@ -168,7 +166,7 @@ def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publish
#TODO finish logic for adding hashes
"""
working_dir = load_env("WORKING_DIR")
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
pups = load_env_json("PUPS", "[]")
threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
@@ -225,7 +223,7 @@ def monitorAuditStatus(api: AirlockAPIWrapper):
last_agents = []
if not last_agents:
last_agents = current_agents
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD","{}")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
# Reverse map for audit → enforcement
reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
flows/prepPolicy.py
+4 -5
View File
@@ -24,13 +24,12 @@ import pandas as pd
from models.execution import ExecutionHistoryRecord, Hash
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_value, load_env, load_env_json
from utils.selector import Selector
from utils.utils import (
colorText,
formatHTML,
import_to_dataframe,
load_env,
load_env_json,
regulator,
)
@@ -250,8 +249,8 @@ def buildPreflights():
formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")
def splitFilepathsGrouped(df, col="filename"):
path_exclusion_constant = load_env("PATH_EXCLUSION_CONST", cast_type= int)
min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type= int)
path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int)
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int)
def clean_split(path):
if not isinstance(path, (str, bytes, os.PathLike)):
@@ -317,7 +316,7 @@ def calculatePath(approved_hashes, split):
dfs_by_policy = [approved_hashes]
badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
min_files_for_path = load_env("MIN_FILES_FOR_PATH", cast_type = int)
min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int)
processed_dfs = []
models/agent.py
+1 -1
View File
@@ -69,7 +69,7 @@ groupid_to_name = {policy.groupid: policy.name for policy in policies}
# Step 3: Enrich agents
for agent in agents:
agent.enrich(groupid_to_name)
agent.enrich_with_policies(groupid_to_name)
"""
models/execution.py
+5 -4
View File
@@ -26,7 +26,8 @@ import dotenv
import pandas as pd
from services.policyhandler import pullPolicyExechistories
from utils.utils import colorText, load_env, load_env_json, regulator
from utils.configmanager import get_protected_value, load_env_json
from utils.utils import colorText, regulator
logger = logging.getLogger(__name__)
@@ -126,7 +127,7 @@ class Hash:
@classmethod
def categorize_hashes(cls, hashes):
threat_tolerance = load_env("VT_THREAT_TOLERANCE", cast_type=int)
threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int)
bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
pups_pattern = regulator(load_env_json("PUPS", "[]"))
@@ -162,9 +163,9 @@ class Hash:
# 3. Approved or Unapproved based on threat level
try:
score = int(scannermatch)
score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}")
if score > threat_tolerance:
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
logger.debug("Unapproved: Unsigned file with high threat score.")
unapproved.append(hash_obj)
else:
services/agenthandler.py
+3 -2
View File
@@ -27,8 +27,9 @@ import pandas as pd
from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import load_env, load_env_json, get_protected_json
from utils.selector import Selector
from utils.utils import colorText, load_env, load_env_json
from utils.utils import colorText
logger = logging.getLogger(__name__)
@@ -207,7 +208,7 @@ def moveAgentToRelatedPolicy(
policy_relationship_map: Dict mapping enforcement → audit.
mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
"""
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
if mode == "audit":
if agent.groupid in policy_relationship_map:
services/policyhandler.py
+3 -2
View File
@@ -27,8 +27,9 @@ from bson import ObjectId
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_json
from utils.setup import get_base_directory
from utils.utils import colorText, load_env_json
from utils.utils import colorText
logger = logging.getLogger(__name__)
@@ -230,7 +231,7 @@ def skipback(days):
def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
for enforcement_policy, audit_policy in policy_relationship_map.items():
api.policy_clone(enforcement_policy, audit_policy)
api.policy_set_auditmode(audit_policy, "1")
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
+46 -22
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(