diff --git a/.gitignore b/.gitignore index 33dfcb9..7862b06 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,5 @@ jobs.json .exe securitytest.py *.toml -system_config.json \ No newline at end of file +system_config.json +Devel_unused/ \ No newline at end of file diff --git a/AirlockToolServer.py b/AirlockToolServer.py index 26d6064..f5e25bb 100644 --- a/AirlockToolServer.py +++ b/AirlockToolServer.py @@ -59,16 +59,16 @@ def main(): 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"), - ) + base_url=str(os.getenv("URL")), + api_key = getAPI(username, "AirlockTools"), + ) + logger.info("Running non-interactively to start monitoring Airlock Changes") diff --git a/Server/scheduler.py b/Server/scheduler.py deleted file mode 100644 index a653c3e..0000000 --- a/Server/scheduler.py +++ /dev/null @@ -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 . -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.") diff --git a/flows/localApproval.py b/flows/localApproval.py index 566ce5d..aae7202 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -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()} diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index b6eabc3..d79d242 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -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 = [] diff --git a/models/agent.py b/models/agent.py index 6376fbe..f21c47a 100644 --- a/models/agent.py +++ b/models/agent.py @@ -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) """ diff --git a/models/execution.py b/models/execution.py index 73d903d..a7466f6 100644 --- a/models/execution.py +++ b/models/execution.py @@ -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: diff --git a/services/agenthandler.py b/services/agenthandler.py index 694785a..00c53d8 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -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: diff --git a/services/policyhandler.py b/services/policyhandler.py index 23f8172..f76b969 100644 --- a/services/policyhandler.py +++ b/services/policyhandler.py @@ -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") diff --git a/utils/configmanager.py b/utils/configmanager.py new file mode 100644 index 0000000..56c8f3c --- /dev/null +++ b/utils/configmanager.py @@ -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 \ No newline at end of file diff --git a/utils/menus.py b/utils/menus.py index 455e88b..118f646 100644 --- a/utils/menus.py +++ b/utils/menus.py @@ -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")) diff --git a/utils/setup.py b/utils/setup.py index 68299aa..17b99ed 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -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 . - - +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . 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 diff --git a/utils/utils.py b/utils/utils.py index 10a7368..888ded7 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -13,7 +13,7 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -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(