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