diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py index e881a2d..e1bf1cb 100644 --- a/AirlockTools_Client.py +++ b/AirlockTools_Client.py @@ -36,7 +36,7 @@ import utils.menus as menus from services.API import AirlockAPIWrapper from services.security import getAPI from services.TaskQueue import AsyncTaskQueue -from utils.setup import setup +from utils.setup import setup, get_base_directory urllib3.disable_warnings( urllib3.exceptions.InsecureRequestWarning @@ -46,10 +46,10 @@ urllib3.disable_warnings( async def main(): #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - working_dir = await setup() + await setup() + base_dir = get_base_directory() logger = logging.getLogger(__name__) - logger.debug("🔍 Logging test: this should appear in both console and file.") - dotenv.load_dotenv(dotenv_path=working_dir / ".env") + dotenv.load_dotenv(dotenv_path=base_dir / ".env") queue = AsyncTaskQueue(worker_count = 3) await queue.start_workers() diff --git a/IRT_icon_32-512.ico b/IRT_icon_32-512.ico new file mode 100644 index 0000000..103b72e Binary files /dev/null and b/IRT_icon_32-512.ico differ diff --git a/flows/otp.py b/flows/otp.py index 1fe48b1..2defd05 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -3,7 +3,7 @@ import logging from services.agenthandler import selectAgents from services.API import AirlockAPIWrapper -from utils.Selector import Selector +from utils.selector import Selector from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index 2d4bab8..dadb6ce 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -27,7 +27,7 @@ from models.policy import Allowlist, Policy from services.API import AirlockAPIWrapper from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread from utils.configmanager import get_protected_value, load_env, load_env_json -from utils.Selector import Selector +from utils.selector import Selector from utils.utils import ( colorText, formatHTML, diff --git a/flows/quietAgent.py b/flows/quietAgent.py index 9a85bea..7b247e9 100644 --- a/flows/quietAgent.py +++ b/flows/quietAgent.py @@ -5,8 +5,8 @@ import pandas as pd from flows.prepPolicy import selectPolicies from services.API import AirlockAPIWrapper -from services.PolicyHandler import getPolicyInfo -from utils.Selector import Selector +from services.policyhandler import getPolicyInfo +from utils.selector import Selector from utils.utils import colorText, load_env logger = logging.getLogger(__name__) diff --git a/models/execution.py b/models/execution.py index 27f1130..cbb61cd 100644 --- a/models/execution.py +++ b/models/execution.py @@ -11,7 +11,7 @@ from typing import List, Optional import aiofiles import pandas as pd -from services.PolicyHandler import pullPolicyExechistories +from services.policyhandler import pullPolicyExechistories from utils.configmanager import get_protected_value, load_env_json from utils.utils import regulator diff --git a/services/agenthandler.py b/services/agenthandler.py index 4682413..83365d1 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -14,7 +14,7 @@ from models.policy import Policy from services.API import AirlockAPIWrapper from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread from utils.configmanager import get_protected_json, load_env -from utils.Selector import Selector +from utils.selector import Selector from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) diff --git a/utils/menus.py b/utils/menus.py index 53f8c82..d4a89b6 100644 --- a/utils/menus.py +++ b/utils/menus.py @@ -20,7 +20,7 @@ import re import dotenv import pandas as pd -import services.PolicyHandler as policyh +import services.policyhandler as policyh from flows.otp import generate, otp_activities_by_agent, revoke from flows.prepPolicy import ( buildPathsandPublishers, @@ -34,7 +34,7 @@ from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAg from services.API import AirlockAPIWrapper from services.TaskQueue import AsyncTaskQueue from utils.configmanager import load_env -from utils.Selector import Selector +from utils.selector import Selector from utils.utils import ( areYouSure, colorText, diff --git a/utils/setup.py b/utils/setup.py index 5d3e153..799ce3f 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -22,29 +22,46 @@ def get_base_directory() -> Path: else: return home / '.local' / 'share' / "AirlockTools" + + + +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" + + def configure_logging(log_dir: Path, log_level: str = "DEBUG"): log_file = log_dir / "airlocktools.log" logger = logging.getLogger() - logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG)) - + # Always allow all messages to propagate to handlers + logger.setLevel(logging.DEBUG) + + # Set httpx logger to DEBUG as well httpx_logger = logging.getLogger("httpx") - httpx_logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG)) - + httpx_logger.setLevel(logging.DEBUG) + # Remove existing handlers for handler in logger.handlers[:]: logger.removeHandler(handler) + # File handler always logs DEBUG and above file_handler = logging.handlers.RotatingFileHandler( log_file, maxBytes=5_000_000, backupCount=5, encoding='utf-8' ) - file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(logging.Formatter( + '%(asctime)s - %(name)s - %(levelname)s - %(message)s' + )) logger.addHandler(file_handler) + # Console handler respects the configured log level console_handler = logging.StreamHandler() + console_handler.setLevel(getattr(logging, log_level.upper(), logging.INFO)) console_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s')) logger.addHandler(console_handler) + # Optional Windows Event Log handler if platform.system() == "Windows": try: event_handler = logging.handlers.NTEventLogHandler("AirlockTools") @@ -56,40 +73,50 @@ 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" def load_system_config() -> dict: + default_config = { + "APPNAME": "AirlockTools", + "LOG_LEVEL": "DEBUG", + "PATH_EXCLUSION_CONST": 4, + "MIN_FILES_FOR_PATH": 4, + "VT_THREAT_TOLERANCE": 4, + "EXTRAS": "NOTTODAY", + "POLICY_MAP_ENF_AUD": { + "enforced_id": "audit_id" + } + } + try: config_path = get_system_config_path() with open(config_path, "r") as f: return json.load(f) except FileNotFoundError: logging.warning("⚠️ system_config.json not found. Using built-in defaults.") - return { - "APPNAME": "AirlockTools", - "LOG_LEVEL": "DEBUG", - "PATH_EXCLUSION_CONST": 4, - "MIN_FILES_FOR_PATH": 4, - "VT_THREAT_TOLERANCE": 4, - "POLICY_MAP_ENF_AUD": { - "enforced_id": "audit_id" - } - } + except json.JSONDecodeError as e: + logging.error(f"❌ Failed to parse system_config.json: {e}") + + return default_config + def load_user_config(config_dir: Path) -> dict: user_config_path = config_dir / "user_config.json" if not user_config_path.exists(): default_user_config = { "URL": "", - "LOG_LEVEL": "INFO" + "LOG_LEVEL": "INFO", + "EXTRAS" : "NOTTODAY" } with open(user_config_path, "w") as f: json.dump(default_user_config, f, indent=4) logging.debug(f"Created user config at {user_config_path}") - with open(user_config_path, "r") as f: - return json.load(f) + try: + with open(user_config_path, "r") as f: + return json.load(f) + except json.JSONDecodeError as e: + logging.error(f"❌ Failed to parse user_config.json: {e}") + return {} + def write_config_to_env(config: dict, env_path: Path): for key, value in config.items(): @@ -101,7 +128,7 @@ def write_config_to_env(config: dict, env_path: Path): except Exception as e: logging.warning(f"Failed to write {key} to .env: {e}") -async def setup() -> Path: +async def setup(): base_dir = get_base_directory() dirs = { 'config': base_dir / 'config', @@ -158,4 +185,3 @@ async def setup() -> Path: logging.debug(f"Service URL set to: {url}") write_config_to_env(merged_config, env_path) - return working_dir \ No newline at end of file