From f1f33948e3b5ec8acb9164067a85a62e4c4bf109 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 20 Oct 2025 14:25:03 -0400 Subject: [PATCH] Bugfix, plus some QoL upgrades from the Async branch --- AirlockTools_Client.py | 17 +- flows/otp.py | 6 +- flows/prepPolicy.py | 64 ++++--- models/execution.py | 378 +++++++++++++++++++++++---------------- services/agenthandler.py | 10 +- utils/menus.py | 22 +-- utils/setup.py | 20 ++- 7 files changed, 299 insertions(+), 218 deletions(-) diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py index 08a46d9..53e3f04 100644 --- a/AirlockTools_Client.py +++ b/AirlockTools_Client.py @@ -30,7 +30,7 @@ import urllib3 import utils.menus as menus from services.API import AirlockAPIWrapper from services.security import getAPI -from utils.setup import setup +from utils.setup import get_base_directory, setup urllib3.disable_warnings( urllib3.exceptions.InsecureRequestWarning @@ -40,10 +40,10 @@ def main(): #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - working_dir = setup() + 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") try: url = os.getenv("URL") @@ -61,13 +61,16 @@ def main(): logger.error(f"Configuration error: {e}", exc_info=True) raise + + api_key = getAPI(username, "AirlockTools") + if api_key is None: + raise ValueError("API key for AirlockTools is missing.") api = AirlockAPIWrapper( base_url=str(os.getenv("URL")), - api_key = getAPI(username, "AirlockTools"), - ) + api_key=api_key, + ) - menus.menu_main(api) diff --git a/flows/otp.py b/flows/otp.py index 4f5d52b..b651b70 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -44,7 +44,11 @@ def generate(api: AirlockAPIWrapper): print(colorText("Please select a duration in minutes: ", "white")) print(colorText("15 mins, 60 mins, 360 mins(6 Hours), 1440 mins (24 Hours), 10080 mins (7 Days):", "white")) duration_selected = Selector.select_int(possible_durations) - if duration_selected: + + if isinstance(duration_selected, list): + duration_selected = duration_selected[0] if duration_selected else None + + if duration_selected is not None: for agent in agents: otp_code = api.otp_generate(agent.agentid, duration_selected, purpose) logger.info(f"Generated OTP for {agent.hostname}: {otp_code}") diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index d79d242..6caedf8 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -21,7 +21,7 @@ from typing import List import dotenv import pandas as pd -from models.execution import ExecutionHistoryRecord, Hash +from models.execution import ExecutionHistoryRecord from models.policy import Allowlist, Policy from services.API import AirlockAPIWrapper from utils.configmanager import get_protected_value, load_env, load_env_json @@ -29,7 +29,6 @@ from utils.selector import Selector from utils.utils import ( colorText, formatHTML, - import_to_dataframe, regulator, ) @@ -86,46 +85,45 @@ def sortHashes( if history_days is None: logging.warning("No history range selected. Aborting.") return - - executions = [] - hashes = [] - - # Pull execution histories for each policy policy_executions = ExecutionHistoryRecord.from_policies( api, selected_policies, type_=type, history_days=history_days ) - - logger.debug(f"Policy_executions is {policy_executions}") - executions.extend(policy_executions) - logger.debug(f"Executions contains {executions}") - if executions: - hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows() - ] + logger.debug(f"Executions contains {policy_executions}") - if hashes: - unique_hashes = Hash.deduplicate(hashes) + enriched_executions = ExecutionHistoryRecord.enrich_with_hashes(api, policy_executions) + categorized_executions = ExecutionHistoryRecord.categorize_executions_by_hash_decision(enriched_executions) + approved, unapproved, needs_review, unknown = ExecutionHistoryRecord.sort_by_hash_decision(categorized_executions) - needs_review, approved, unapproved = Hash.categorize_hashes( - hashes=unique_hashes - ) - - categories = { + categories = { "needs_review": needs_review, "approved": approved, "unapproved": unapproved, + "unknown" : unknown } - for label, category in categories.items(): - csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv" - html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html" + + for label, records in categories.items(): + csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv" + html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html" + + # Convert ExecutionHistoryRecord objects to dictionaries + df = pd.DataFrame([r.__dict__ for r in records]) + + # Optional: flatten hash_obj if needed + if not df.empty and 'hash_obj' in df.columns: + hash_df = df['hash_obj'].apply(lambda h: h.to_dict() if h else {}) + df = pd.concat([df.drop(columns=['hash_obj']), hash_df], axis=1) + + # Save to CSV + df.to_csv(csv_path, index=False) + logger.info(f"Saved {label} executions to {csv_path}") + + # Generate HTML + formatHTML(df, html_path) + logger.info(f"Generated HTML report at {html_path}") - ExecutionHistoryRecord.enrich_with_hashes_and_export( - executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label - ) - df = import_to_dataframe(csv_path) - formatHTML(df, html_path) def buildPathsandPublishers(split): working_dir = load_env("WORKING_DIR") @@ -255,7 +253,7 @@ def splitFilepathsGrouped(df, col="filename"): def clean_split(path): if not isinstance(path, (str, bytes, os.PathLike)): return [] - parts = os.path.normpath(path).split(os.sep) + parts = str(os.path.normpath(path)).split(os.sep) parts = [p for p in parts if p] # Remove empty strings return parts @@ -268,9 +266,9 @@ def splitFilepathsGrouped(df, col="filename"): df = df.copy() split_paths = df[col].apply(clean_split) - # Filter out paths with fewer than `min_files_for_path` components - df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy() - split_paths = split_paths[df.index] + if min_files_for_path is not None: + df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy() + split_paths = split_paths[df.index] df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant])) grouped = df.groupby("group_key") diff --git a/models/execution.py b/models/execution.py index a7466f6..1ef71d1 100644 --- a/models/execution.py +++ b/models/execution.py @@ -13,18 +13,20 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . +import dataclasses import inspect import json import logging import os import re -from dataclasses import dataclass +from dataclasses import asdict, dataclass from datetime import datetime -from typing import List, Optional +from typing import List, Optional, Tuple import dotenv import pandas as pd +from services.API import AirlockAPIWrapper from services.policyhandler import pullPolicyExechistories from utils.configmanager import get_protected_value, load_env_json from utils.utils import colorText, regulator @@ -33,80 +35,39 @@ logger = logging.getLogger(__name__) dotenv.load_dotenv() - +@dataclass class Hash: """ Hash model representing Hash data """ - - def __init__( - self, - sha256, - applications=None, - baselines=None, - blocklists=None, - createtime=None, - datetime=None, - description=None, - filename=None, - filepath=None, - filesize=None, - md5=None, - modtime=None, - origname=None, - productname=None, - productversion=None, - publisher=None, - reputation=None, - sha128=None, - sha384=None, - sha512=None, - ): - self.sha256 = sha256 - self.applications = applications - self.baselines = baselines - self.blocklists = blocklists - self.createtime = createtime - self.datetime = datetime - self.description = description - self.filename = filename - self.filepath = filepath - self.filesize = filesize - self.md5 = md5 - self.modtime = modtime - self.origname = origname - self.productname = productname - self.productversion = productversion - self.publisher = publisher - self.reputation = reputation - self.sha128 = sha128 - self.sha384 = sha384 - self.sha512 = sha512 - - def __repr__(self): - attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) - return f"" - - def __eq__(self, other): - if isinstance(other, Hash): - return self.sha256 == other.sha256 - return False - - def __hash__(self): - return hash(self.sha256) - + sha256: str + applications: str + baselines: str + blocklists: str + createtime: str + datetime: str + description: str + filename: str + filepath: str + filesize: str + md5: str + modtime: str + origname: str + productname: str + productversion: str + publisher: str + reputation: str + sha128: str + sha384: str + sha512: str + at_decision: Optional[str] = None + def to_dict(self): - """Returns a dictionary representation of the hash.""" - return self.__dict__ - - @staticmethod - def safe_int(value, default=0): - """Safely convert a value to int, returning default on failure.""" - try: - return int(value) - except (TypeError, ValueError): - return default + return asdict(self) + @classmethod + def from_dict(cls, data: dict): + return cls(**data) @classmethod def deduplicate(cls, hash_list): @@ -131,9 +92,9 @@ class Hash: bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]")) pups_pattern = regulator(load_env_json("PUPS", "[]")) - needs_review = [] - approved = [] - unapproved = [] + approved_count = 0 + unapproved_count = 0 + needs_review_count = 0 for hash_obj in hashes: publisher = hash_obj.publisher or "" @@ -147,18 +108,21 @@ class Hash: # 1. Unapproved: bad publisher or PUP if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): logger.debug("Unapproved: Publisher matches bad publisher pattern.") - unapproved.append(hash_obj) + hash_obj.at_decision = "unapproved" + unapproved_count += 1 continue if re.search(pups_pattern, description, re.IGNORECASE): logger.debug("Unapproved: Description matches PUP pattern.") - unapproved.append(hash_obj) + hash_obj.at_decision = "unapproved" + unapproved_count += 1 continue # 2. Approved: signed if publisher != "Not Signed": logger.debug("Approved: File is signed and not flagged.") - approved.append(hash_obj) + hash_obj.at_decision = "approved" + approved_count += 1 continue # 3. Approved or Unapproved based on threat level @@ -167,17 +131,20 @@ class Hash: logger.debug(f"Parsed scannermatch score: {score}") if score > threat_tolerance: # pyright: ignore[reportOperatorIssue] logger.debug("Unapproved: Unsigned file with high threat score.") - unapproved.append(hash_obj) + hash_obj.at_decision = "unapproved" + unapproved_count += 1 else: logger.debug("Approved: Unsigned file with low threat score.") - approved.append(hash_obj) + hash_obj.at_decision = "approved" + approved_count += 1 except (ValueError, TypeError): - logger.debug("Needs Review: Scannermatch score is missing or invalid.") - needs_review.append(hash_obj) + logger.debug("Needs Review: Scannermatch score is missing or invalid. — {e}") + hash_obj.at_decision = "needs_review" + needs_review_count += 1 - logger.debug(f"Final counts — Needs Review: {len(needs_review)}, Approved: {len(approved)}, Unapproved: {len(unapproved)}") - return needs_review, approved, unapproved + logger.debug(f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}") + return hashes @classmethod @@ -209,43 +176,6 @@ class Hash: logger.info(f"CSV file saved to: {file_path}") -""" -#Example - Convert Dataframe returned by hash query into hash objects -hash_objects = [] -for _, row in df.iterrows(): - try: - parsed_data = ast.literal_eval(row['data']) - hash_obj = Hash(sha256=row['sha256'], **parsed_data) - hash_objects.append(hash_obj) - except Exception as e: - print(f"Error parsing row: {e}") - -# Display the created Hash objects -for obj in hash_objects: - print(obj) - - -# Categorize hashes -needs_review, approved, unapproved = Hash.categorize_hashes( - hashes=hash_objects, - threat_tolerance=3, - untrusted_pattern=untrusted_pattern, - pups_pattern=pups_pattern -) - - # Deduplicate - deduped_hashes = Hash.deduplicate(hash_list) - -# Specify the directory where you want to save the CSV -output_directory = "C:/Users/Brandon/Documents/HashExports" - -# Call the export method -Hash.export_to_csv(hashes_for_export, output_directory) - - -""" - - @dataclass class ExecutionHistoryRecord: # Mandatory fields @@ -275,37 +205,7 @@ class ExecutionHistoryRecord: extname: Optional[str] = None exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge - - @staticmethod - def enrich_with_hashes_and_export( - executions: list, hashes: list, directory_path: str, label: str = "enriched" - ): - exec_df = pd.DataFrame([e.__dict__ for e in executions]) - hash_df = pd.DataFrame([h.to_dict() for h in hashes]) - - logger.debug(f"Execution DataFrame columns: {exec_df.columns}") - logger.debug(f"Hash DataFrame columns: {hash_df.columns}") - - if hash_df.empty: - logger.warning(f"hash_df is empty for label: {label}. Skipping merge.") - merged_df = exec_df.copy() - else: - merged_df = pd.merge( - exec_df, - hash_df, - on="sha256", - how="left", # Preserve all executions, enrich where possible - suffixes=("_exec", "_hash") - ) - merged_df.sort_values(by="filename_exec", inplace=True) - logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}") - - filename = f"{label}_executions.csv" - os.makedirs(directory_path, exist_ok=True) - file_path = os.path.join(directory_path, filename) - merged_df.to_csv(file_path, index=False) - - logger.info(f"CSV file saved to: {file_path}") + hash_obj: Optional[Hash] = None @classmethod @@ -354,6 +254,7 @@ class ExecutionHistoryRecord: extname=data.get("extname"), exttype=data.get("exttype"), extbrowser=data.get("extbrowser"), + hash_obj=data.get("hash_obj") ) @classmethod @@ -385,10 +286,181 @@ class ExecutionHistoryRecord: ) return executions + + @staticmethod + def enrich_with_hashes( + api: AirlockAPIWrapper, + executions: List["ExecutionHistoryRecord"] + ) -> List["ExecutionHistoryRecord"]: + """ + Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API. + """ + sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256}) + logger.info(f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records.") + + if not sha256_list: + logger.warning("No sha256 values found in execution records. Skipping enrichment.") + return executions + + logger.debug("Querying hash data from API...") + hash_df = api.hash_query(sha256_list) + logger.info(f"Retrieved {len(hash_df)} hash records from API.") + + hash_objects = [] + required_fields = { + f.name for f in dataclasses.fields(Hash) + if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING + } + + for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()): + row_dict = row.to_dict() + + # Unwrap nested 'data' field if present + if "data" in row_dict and isinstance(row_dict["data"], dict): + row_dict = row_dict["data"] + + # Inject the sha256 back into the row + row_dict["sha256"] = sha256 + + missing = required_fields - row_dict.keys() + if missing: + logger.warning(f"Skipping hash row due to missing fields: {missing}") + logger.debug(f"Row content: {row_dict}") + continue + + try: + hash_obj = Hash.from_dict(row_dict) + hash_objects.append(hash_obj) + except Exception as e: + logger.warning(f"Failed to create Hash from row: {e}") + logger.debug(f"Row content: {row_dict}") + + logger.debug("Converted hash DataFrame to Hash objects.") + + hash_lookup = {h.sha256.strip().lower(): h for h in hash_objects} + logger.debug("Built hash lookup table.") + + enriched_count = 0 + for exec_record in executions: + hash_obj = hash_lookup.get(exec_record.sha256.strip().lower()) + if hash_obj: + exec_record.hash_obj = hash_obj + enriched_count += 1 + + logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.") + return executions + + @staticmethod + def categorize_executions_by_hash_decision(executions: List["ExecutionHistoryRecord"]) -> List["ExecutionHistoryRecord"]: + """ + Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation. + + Modifies the `at_decision` field of each associated Hash object in-place. + + Returns: + List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated. + """ + 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", "[]")) + + approved_count = 0 + unapproved_count = 0 + needs_review_count = 0 + + for record in executions: + hash_obj = record.hash_obj + if not hash_obj: + continue # Skip if no hash object is attached + + publisher = hash_obj.publisher or "" + description = hash_obj.description or "" + reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + scannermatch = reputation.get("scannermatch") + + logger.debug(f"Evaluating hash: {hash_obj}") + logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") + + # 1. Unapproved: bad publisher or PUP + if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): + logger.debug("Unapproved: Publisher matches bad publisher pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + if re.search(pups_pattern, description, re.IGNORECASE): + logger.debug("Unapproved: Description matches PUP pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + # 2. Approved: signed + if publisher != "Not Signed": + logger.debug("Approved: File is signed and not flagged.") + hash_obj.at_decision = "approved" + approved_count += 1 + continue + + # 3. Approved or Unapproved based on threat level + try: + score = int(scannermatch) # pyright: ignore[reportArgumentType] + logger.debug(f"Parsed scannermatch score: {score}") + if threat_tolerance is not None and score >= threat_tolerance: + logger.debug("Unapproved: Unsigned file with high threat score.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + else: + logger.debug("Approved: Unsigned file with low threat score.") + hash_obj.at_decision = "approved" + approved_count += 1 + except (ValueError, TypeError) as e: + logger.debug(f"Needs Review: Scannermatch score is missing or invalid. — {e}") + hash_obj.at_decision = "needs_review" + needs_review_count += 1 + + logger.debug( + f"Final counts — Needs Review: {needs_review_count}, " + f"Approved: {approved_count}, Unapproved: {unapproved_count}" + ) + + return executions + + + @classmethod + def sort_by_hash_decision( + cls, executions: List["ExecutionHistoryRecord"] + ) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]: + """ + Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups + based on the value of hash_obj.at_decision. + + Returns: + Tuple of lists: (approved, unapproved, needs_review, unknown) + """ + approved = [] + unapproved = [] + needs_review = [] + unknown = [] + + for record in executions: + decision = getattr(record.hash_obj, "at_decision", None) + if decision == "approved": + approved.append(record) + elif decision == "unapproved": + unapproved.append(record) + elif decision == "needs_review": + needs_review.append(record) + else: + unknown.append(record) + + logger.info(f"[ExecutionHistoryRecord] Sorted {len(executions)} records by hash_obj.at_decision:") + logger.info(f" Approved: {len(approved)}") + logger.info(f" Unapproved: {len(unapproved)}") + logger.info(f" Needs Review: {len(needs_review)}") + logger.info(f" Unknown/Unset: {len(unknown)}") + + return approved, unapproved, needs_review, unknown - def __repr__(self): - attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) - return f"" """ diff --git a/services/agenthandler.py b/services/agenthandler.py index 290baba..a04fa3b 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -86,12 +86,8 @@ def findAllAgents(api): policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()] agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()] - # Step 2: Create groupid → groupname map - 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(policies) return agents @@ -120,7 +116,7 @@ def findAgents(api, return_dataframe): if user_input == 'y': timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"agentsearch_{timestamp}.csv" - file_path = os.path.join(working_dir, filename) + file_path = os.path.join(str(working_dir), filename) agent_df.to_csv(file_path, index=False) logging.info(f"Exported DataFrame to {file_path}") @@ -138,7 +134,7 @@ def findAgents(api, return_dataframe): def selectAgents(api: AirlockAPIWrapper) -> List[Agent]: print(colorText("🔍 Device Search", "cyan")) print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan")) - print(colorText("When you're done, press Enter twice.\n", "cyan")) + print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan")) print(colorText("Example:", "cyan")) print(colorText("H00000", "cyan")) print(colorText("UTN00000", "cyan")) diff --git a/utils/menus.py b/utils/menus.py index db878c2..e10f8c3 100644 --- a/utils/menus.py +++ b/utils/menus.py @@ -70,17 +70,20 @@ def menu_main(api: AirlockAPIWrapper): elif choice == "2": menu_otp(api) elif choice == "3": - choices = ["audit", "enforcement"] + choices = ["audit", "enforcement", "Cancel"] print(colorText("Move devices to which state?:", "yellow")) direction = Selector.select_string(choices, False, False) - devices = selectAgents(api) - print(colorText("Would you like to continue with these devices?","white")) - for device in devices: - print(device.hostname) - confirm = Selector.confirm() - if direction and devices and confirm: + if direction == "Cancel": + pass + else: + devices = selectAgents(api) + print(colorText("Would you like to continue with these devices?","white")) for device in devices: - moveAgentToRelatedPolicy(api,device, direction) + print(device.hostname) + confirm = Selector.confirm() + if direction and devices and confirm: + for device in devices: + moveAgentToRelatedPolicy(api,device, str(direction)) elif choice == "4": findAgents(api,False) elif choice == "5": @@ -96,8 +99,6 @@ def menu_main(api: AirlockAPIWrapper): else: print(colorText("Invalid choice. Please try again.", "red")) - - def menu_policy_enforce(api: AirlockAPIWrapper): selected_policies = [] destination_policy = [] @@ -243,7 +244,6 @@ def menu_policy_enforce(api: AirlockAPIWrapper): else: print(colorText("Invalid choice. Please try again.", "red")) - def menu_otp(api: AirlockAPIWrapper): working_dir = load_env("WORKING_DIR") while True: diff --git a/utils/setup.py b/utils/setup.py index d2d70ef..a530926 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -39,22 +39,31 @@ def get_base_directory() -> Path: 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)) - # 🔧 Clear existing handlers + # Always allow all messages to propagate to handlers + 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") @@ -65,7 +74,7 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): logger.warning(f"Could not attach Windows Event Log handler: {e}") 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" @@ -111,7 +120,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}") -def setup() -> Path: +def setup(): base_dir = get_base_directory() dirs = { 'config': base_dir / 'config', @@ -172,4 +181,3 @@ def setup() -> Path: write_config_to_env(merged_config, env_path) - return working_dir \ No newline at end of file