diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py index c2d4093..47c9562 100644 --- a/AirlockTools_Client.py +++ b/AirlockTools_Client.py @@ -14,76 +14,71 @@ # along with this program. If not, see . -#TODO Continue implementing logger -#TODO Add input sanitation and CSV injection prevention -#TODO Continue OTP and Local approval rewrites -#TODO Explore pywin32 -#TODO Fix Requirements.txt -#TODO Create Generic system_config.json for gitea +# TODO Continue implementing logger +# TODO Add input sanitation and CSV injection prevention +# TODO Continue OTP and Local approval rewrites +# TODO Explore pywin32 +# TODO Fix Requirements.txt +# TODO Create Generic system_config.json for gitea import logging import os import tempfile + import dotenv import urllib3 from services.API import AirlockAPIWrapper from services.security import getAPI from utils.setup import get_base_directory, setup -from utils.TUI import run_AirlockTools +from utils.TUI import run_Loxide from utils.utils import irtang -urllib3.disable_warnings( - urllib3.exceptions.InsecureRequestWarning -) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + def main(): - - + if "NUITKA_ONEFILE_PARENT" in os.environ: splash_filename = os.path.join( tempfile.gettempdir(), - f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp" + f"onefile_{int(os.environ['NUITKA_ONEFILE_PARENT'])}_splash_feedback.tmp", ) if os.path.exists(splash_filename): os.unlink(splash_filename) - - irtang() - #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored + # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored setup() base_dir = get_base_directory() logger = logging.getLogger(__name__) dotenv.load_dotenv(dotenv_path=base_dir / ".env") try: - url = os.getenv("URL") - username = os.getenv("USERNAME") + url = os.getenv("URL") + username = os.getenv("USERNAME") - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") + if not url: + raise ValueError("Missing URL in environment variables.") + if not username: + raise ValueError("Missing USERNAME in environment variables.") - logger.debug(f"Retrieved URL: {url}") - logger.debug(f"Retrieved Username: {username}") + 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 + logger.error(f"Configuration error: {e}", exc_info=True) + raise - - api_key = getAPI(username, "AirlockTools") + api_key = getAPI(username, "Loxide") if api_key is None: - raise ValueError("API key for AirlockTools is missing.") + raise ValueError("API key for Loxide is missing.") api = AirlockAPIWrapper( base_url=str(os.getenv("URL")), api_key=api_key, ) - run_AirlockTools(api) - + run_Loxide(api) if __name__ == "__main__": diff --git a/AirlockTools_Server.py b/AirlockTools_Server.py index 56a58a0..183635c 100644 --- a/AirlockTools_Server.py +++ b/AirlockTools_Server.py @@ -14,12 +14,11 @@ # along with this program. If not, see . - -#TODO Add CSV injection prevention -#TODO Continue OTP and Local approval rewrites -#TODO Explore pywin32 -#TODO Fix Requirements.txt -#TODO Create Generic system_config.json for gitea +# TODO Add CSV injection prevention +# TODO Continue OTP and Local approval rewrites +# TODO Explore pywin32 +# TODO Fix Requirements.txt +# TODO Create Generic system_config.json for gitea import logging @@ -29,62 +28,66 @@ import dotenv import urllib3 import flows.localApproval as la -from Server.scheduler_async import recurring_job, register_function, reload_jobs, start_scheduler +from Server.scheduler_async import ( + recurring_job, + register_function, + reload_jobs, + start_scheduler, +) from services.API import AirlockAPIWrapper from services.policyhandler import updateAuditPoliciesFromEnforcementPolices from services.security import getAPI from utils.setup import setup -urllib3.disable_warnings( - urllib3.exceptions.InsecureRequestWarning -) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + def main(): - #Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - + # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored + working_dir = setup() - + logger = logging.getLogger(__name__) dotenv.load_dotenv(dotenv_path=working_dir / ".env") try: - url = os.getenv("URL") - username = os.getenv("USERNAME") + url = os.getenv("URL") + username = os.getenv("USERNAME") - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") + if not url: + raise ValueError("Missing URL in environment variables.") + if not username: + raise ValueError("Missing USERNAME in environment variables.") + + logger.debug(f"Retrieved URL: {url}") + logger.debug(f"Retrieved Username: {username}") - 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 - + 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") - register_function("monitorLA", la.scheduleAddingLAHashes) register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices) - - if not os.path.exists("scheduling\\jobs.json"): + if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api]) - recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]) + recurring_job( + "updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api] + ) else: reload_jobs() - + start_scheduler() - + + if __name__ == "__main__": main() diff --git a/README.md b/README.md index d647cdc..8257c1c 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ -# šŸ›”ļø Airlock Tools +# šŸ›”ļø Loxide -Python toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction. +Python/Rust/Oxide toolkit for secure, auditable, and automated airlock agent and policy management. Designed for enterprise environments, it supports advanced policy workflows, device tracking, and terminal-based interaction. --- @@ -51,7 +51,7 @@ Python toolkit for secure, auditable, and automated airlock agent and policy man ## šŸ“œ License -**AirlockTools** is licensed under the **GNU Affero General Public License v3.0**. +**Loxide** is licensed under the **GNU Affero General Public License v3.0**. You may copy, distribute, and modify the software under the terms of the AGPL-3.0 license. diff --git a/Server/scheduler_async.py b/Server/scheduler_async.py index 6f90d36..87e8c80 100644 --- a/Server/scheduler_async.py +++ b/Server/scheduler_async.py @@ -30,6 +30,7 @@ scheduled_jobs: Dict[str, asyncio.TimerHandle] = {} # Path to the JSON file for job persistence TODO - pin this to the correct place JOBS_FILE = os.path.join(os.getcwd(), "jobs.json") + def register_function(name: str, func: Callable): """ Register a function so it can be called by name later. @@ -38,6 +39,7 @@ def register_function(name: str, func: Callable): """ FUNCTION_MAP[name] = func + def load_jobs() -> List[Dict[str, Any]]: """ Load jobs from the JSON file, or return [] if none exist. @@ -47,6 +49,7 @@ def load_jobs() -> List[Dict[str, Any]]: with open(JOBS_FILE, "r") as f: return json.load(f) + def save_jobs(jobs: List[Dict[str, Any]]): """ Save jobs to the JSON file (overwrite). @@ -54,6 +57,7 @@ def save_jobs(jobs: List[Dict[str, Any]]): with open(JOBS_FILE, "w") as f: json.dump(jobs, f, indent=4) + def cancel_job(job_id: str): """ Cancel a scheduled job by ID and remove it from the registry and persistence. @@ -66,7 +70,15 @@ def cancel_job(job_id: str): jobs = [j for j in load_jobs() if j.get("id") != job_id] save_jobs(jobs) -def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, kwargs=None, persist=True): + +def run_once_job( + job_id: str, + func_name: str, + delay_seconds: float, + args=None, + kwargs=None, + persist=True, +): """ Schedule a job to run once after a delay (in seconds). """ @@ -87,18 +99,25 @@ def run_once_job(job_id: str, func_name: str, delay_seconds: float, args=None, k if persist: jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append({ - "id": job_id, - "type": "once", - "delay": delay_seconds, - "function": func_name, - "args": args, - "kwargs": kwargs - }) + jobs.append( + { + "id": job_id, + "type": "once", + "delay": delay_seconds, + "function": func_name, + "args": args, + "kwargs": kwargs, + } + ) save_jobs(jobs) - logger.info(f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds.") + logger.info( + f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds." + ) -def recurring_job(job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True): + +def recurring_job( + job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True +): """ Schedule a recurring job. """ @@ -121,17 +140,20 @@ def recurring_job(job_id: str, func_name: str, interval: float, args=None, kwarg if persist: jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append({ - "id": job_id, - "type": "recurring", - "interval": interval, - "function": func_name, - "args": args, - "kwargs": kwargs - }) + jobs.append( + { + "id": job_id, + "type": "recurring", + "interval": interval, + "function": func_name, + "args": args, + "kwargs": kwargs, + } + ) save_jobs(jobs) logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") + def reload_jobs(): """ Reload jobs from JSON and reschedule them. @@ -145,7 +167,7 @@ def reload_jobs(): job["delay"], job.get("args"), job.get("kwargs"), - persist=False + persist=False, ) elif job["type"] == "recurring": recurring_job( @@ -154,9 +176,10 @@ def reload_jobs(): job["interval"], job.get("args"), job.get("kwargs"), - persist=False + persist=False, ) + async def start_scheduler(): """ Start the asynchronous scheduler loop. @@ -189,4 +212,4 @@ async def start_scheduler(): while True: await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later except asyncio.CancelledError: - logger.critical("Scheduler stopped.") \ No newline at end of file + logger.critical("Scheduler stopped.") diff --git a/airlock_libs/airlock_libs.pyi b/airlock_libs/airlock_libs.pyi index 2e061f1..137dc92 100644 --- a/airlock_libs/airlock_libs.pyi +++ b/airlock_libs/airlock_libs.pyi @@ -1,6 +1,9 @@ -from typing import Dict, List, Optional -def pull_policy_exec_histories(self, type: List[str], checkpoint: str, policy: List[str]) -> str: - """Retrieve execution history logs.""" +from typing import Dict, List + +def pull_policy_exec_histories( + self, type: List[str], checkpoint: str, policy: List[str] +) -> str: + """Retrieve execution history logs.""" def api(AirlockAPIWrapper): """ @@ -22,66 +25,66 @@ def api(AirlockAPIWrapper): """ def history_logging( - api, - exec_types: str, - checkpoint_number: str, - policy_names: str, - ) -> List[Dict[str, Any]]: - """ - Query execution history logs from the Airlock API. + api, + exec_types: str, + checkpoint_number: str, + policy_names: str, +) -> List[Dict[str, Any]]: + """ + Query execution history logs from the Airlock API. - Parameters - ---------- - exec_types : str - A JSON-style string list of execution types to retrieve. - Example: "[3,5,8]" - - 0 = Trusted Execution - - 1 = Blocked Execution - - 2 = Untrusted Execution [Audit] - - 3 = Untrusted Execution [OTP] - - 5 = Trusted Publisher Execution - - 8 = Trusted Process Execution - (etc.) + Parameters + ---------- + exec_types : str + A JSON-style string list of execution types to retrieve. + Example: "[3,5,8]" + - 0 = Trusted Execution + - 1 = Blocked Execution + - 2 = Untrusted Execution [Audit] + - 3 = Untrusted Execution [OTP] + - 5 = Trusted Publisher Execution + - 8 = Trusted Process Execution + (etc.) - checkpoint_number : str - The checkpoint ID. Used to fetch results after a certain event. - Example: "601d275487bacb01e3470713" + checkpoint_number : str + The checkpoint ID. Used to fetch results after a certain event. + Example: "601d275487bacb01e3470713" - policy_names : str - A comma-separated or JSON-style list of policy group names. - Example: "Apple Mac" or "["Apple Mac", "Servers London"]" + policy_names : str + A comma-separated or JSON-style list of policy group names. + Example: "Apple Mac" or "["Apple Mac", "Servers London"]" - Returns - ------- - List[Dict[str, Any]] - A list of dictionaries, where each dictionary represents an - execution history record. Each record can include fields like: + Returns + ------- + List[Dict[str, Any]] + A list of dictionaries, where each dictionary represents an + execution history record. Each record can include fields like: - - checkpoint: str - - type: int - - username: str - - hostname: str - - filename: str - - ppolicy: str - - policyname: str - - policyver: str - - commandline: str - - publisher: str - - pprocess: str - - gprocess: str - - sha256: str - - datetime: str - - ip: str - - localip: str - Raises - ------ - RuntimeError - If the request fails or the response cannot be parsed. + - checkpoint: str + - type: int + - username: str + - hostname: str + - filename: str + - ppolicy: str + - policyname: str + - policyver: str + - commandline: str + - publisher: str + - pprocess: str + - gprocess: str + - sha256: str + - datetime: str + - ip: str + - localip: str + Raises + ------ + RuntimeError + If the request fails or the response cannot be parsed. - Example - ------- - >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac") - >>> print(histories[0]["filename"]) - 'chrome.exe' - """ - ... \ No newline at end of file + Example + ------- + >>> histories = await airlock_libs.history_logging("[3,5,8]", "601d275487bacb01e3470713", "Apple Mac") + >>> print(histories[0]["filename"]) + 'chrome.exe' + """ + ... diff --git a/flows/localApproval.py b/flows/localApproval.py index 789b917..5cc69a2 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -41,7 +41,9 @@ def getLocalApprovals(api: AirlockAPIWrapper): result = api.otp_find_awaiting() local_approval = pd.DataFrame(result["response"]["otpusage"]) if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"): - previous_run = pd.read_parquet(f"{base_dir}\\cache\\newest_local_approval.parquet") + previous_run = pd.read_parquet( + f"{base_dir}\\cache\\newest_local_approval.parquet" + ) previous_run.to_parquet( f"{base_dir}\\cache\\last_local_approval.parquet", index=False ) @@ -66,10 +68,10 @@ def getLocalApprovals(api: AirlockAPIWrapper): def scheduleAddingLAHashes(api: AirlockAPIWrapper): - policy_relationship_map = get_protected_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) + threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type=int) try: register_function("add_hash", returnFromLocalApproval) @@ -93,7 +95,9 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper): duration_minutes = int(batch_df["duration"].iloc[0]) start_time = datetime.datetime.now() run_time = start_time + datetime.timedelta(minutes=duration_minutes) - early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95)) + early_time = start_time + datetime.timedelta( + minutes=np.floor(duration_minutes * 0.95) + ) early_timestamp = early_time.timestamp() run_timestamp = run_time.timestamp() @@ -148,7 +152,13 @@ def scheduleAddingLAHashes(api: AirlockAPIWrapper): logger.warning(f"Failed to process batch {batchid}: {e}") -def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant +def returnFromLocalApproval( + api, + device_df, + policy_relationship_map, + bad_publisher_list, + pups, + threat_tolerance_constant, ): """ # Get unique policy names from device list @@ -166,11 +176,14 @@ 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 = get_protected_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") - print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}") + print( + f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}" + ) + def moveToLocalApproval(api: AirlockAPIWrapper): possible_durations = [15, 60, 360, 1440, 10080] @@ -213,10 +226,9 @@ def moveToLocalApproval(api: AirlockAPIWrapper): def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid): - + purpose = f"šŸŽ« Local Approval šŸŽ« - {duration_selected} mins - batch:{batchid} Client:{agentid}" api.otp_generate(agentid, duration_selected, purpose) - def monitorAuditStatus(api: AirlockAPIWrapper): @@ -224,11 +236,13 @@ def monitorAuditStatus(api: AirlockAPIWrapper): last_agents = [] if not last_agents: last_agents = current_agents - policy_relationship_map = get_protected_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()} - known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items()) + known_transitions = set(policy_relationship_map.items()) | set( + reverse_policy_map.items() + ) # Index last_agents by hostname for quick lookup last_agent_map = {agent.hostname: agent for agent in last_agents} @@ -261,8 +275,8 @@ def monitorAuditStatus(api: AirlockAPIWrapper): def getNewLocalApprovals(api: AirlockAPIWrapper): - - working_dir = load_env("WORKING_DIR") + + working_dir = load_env("WORKING_DIR") current_la = getLocalApprovals(api) # Load old approval list @@ -273,15 +287,21 @@ def getNewLocalApprovals(api: AirlockAPIWrapper): old_la = pd.DataFrame(columns=current_la.columns) # Create composite keys - current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str) + current_la["key"] = ( + current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str) + ) old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str) # Find new entries new_entries = current_la[~current_la["key"].isin(old_la["key"])] # Convert 'granted' to datetime and filter by last 10 minutes - new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce") - ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10) + new_entries["granted"] = pd.to_datetime( + new_entries["granted"], utc=True, errors="coerce" + ) + ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + minutes=10 + ) recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago] # Save current approvals for next run diff --git a/flows/otp.py b/flows/otp.py index 6c09db2..8b15fab 100644 --- a/flows/otp.py +++ b/flows/otp.py @@ -14,7 +14,6 @@ # along with this program. If not, see . - from datetime import datetime import logging import os @@ -33,19 +32,24 @@ logger = logging.getLogger(__name__) def otp_generate(api: AirlockAPIWrapper): otp_dict = {} agents = selectAgents(api) - print(colorText("Would you like to continue with these devices?","white")) + print(colorText("Would you like to continue with these devices?", "white")) for agent in agents: print(agent.hostname) confirm = Selector.confirm() if agents and confirm: requester = get_sanitized_input("Who is requesting the OTP: ") because = get_sanitized_input("Why/What work are they doing?: ") - + purpose = f"Requester: {requester} - for : {because}" possible_durations = [15, 60, 360, 1440, 10080] 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")) + 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 isinstance(duration_selected, list): @@ -60,56 +64,68 @@ def otp_generate(api: AirlockAPIWrapper): print(colorText("Requested Codes:", "green")) for key, value in otp_dict.items(): - print(colorText(f"{key} | {value}","green")) + print(colorText(f"{key} | {value}", "green")) + def otp_activities_by_agent(api: AirlockAPIWrapper): activeagents = api.otp_find_active() awaitingagents = api.otp_find_awaiting() enforcedagents = api.otp_find_enforced() revokedagents = api.otp_find_revoked() - - + # Add a 'status' column to each DataFrame - activeagents['status'] = 'active' - awaitingagents['status'] = 'awaiting' - enforcedagents['status'] = 'enforced' - revokedagents['status'] = 'revoked' + activeagents["status"] = "active" + awaitingagents["status"] = "awaiting" + enforcedagents["status"] = "enforced" + revokedagents["status"] = "revoked" # Combine all into one DataFrame - combined_agents = pd.concat([activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + combined_agents = pd.concat( + [activeagents, awaitingagents, enforcedagents, revokedagents], ignore_index=True + ) + combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - #Optionally, select specific hosts - user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() - if user_input == 'y': + # Optionally, select specific hosts + user_input = ( + get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") + .strip() + .lower() + ) + if user_input == "y": agentnames = [] agents = selectAgents(api) for agent in agents: agentnames.append(agent.hostname) - combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] + combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - #Present and select rows + # Present and select rows selected_rows = Selector.select_dataframe_with_mode( combined_agents, - columns=['otpid', 'hostname', 'status','purpose','granted'], - header="OTP Sessions" + columns=["otpid", "hostname", "status", "purpose", "granted"], + header="OTP Sessions", ) combined_df = pd.DataFrame() for row in selected_rows: - otpid = row['otpid'] - hostname = row['hostname'] + otpid = row["otpid"] + hostname = row["hostname"] result = api.otp_get_activities(otpid) - result['hostname'] = hostname + result["hostname"] = hostname if not result.empty: logger.info(f"Activities for {hostname} (otpid: {otpid}):\n{result}") combined_df = pd.concat([combined_df, result], ignore_index=True) else: logger.info(f"No activities found for {hostname} (otpid: {otpid})") - user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower() - if user_input == 'y': + user_input = ( + get_sanitized_input( + "\nWould you like to export the results to a CSV file? (y/n): " + ) + .strip() + .lower() + ) + if user_input == "y": working_dir = load_env("WORKING_DIR") timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"otp_activities_{timestamp}.csv" @@ -128,43 +144,44 @@ def otp_activities_by_agent(api: AirlockAPIWrapper): logging.debug("User declined to export the DataFrame.") - def otp_revoke(api: AirlockAPIWrapper): activeagents = api.otp_find_active() awaitingagents = api.otp_find_awaiting() - - activeagents['status'] = 'active' - awaitingagents['status'] = 'awaiting' + + activeagents["status"] = "active" + awaitingagents["status"] = "awaiting" combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - # Combine all into one DataFrame + # Combine all into one DataFrame combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) - combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + combined_agents = combined_agents.sort_values(by="otpid", ascending=False) - #Optionally, select specific hosts - user_input = get_sanitized_input("\nWould you like to search for a specific device? (y/n): ").strip().lower() - if user_input == 'y': + # Optionally, select specific hosts + user_input = ( + get_sanitized_input("\nWould you like to search for a specific device? (y/n): ") + .strip() + .lower() + ) + if user_input == "y": agentnames = [] agents = selectAgents(api) for agent in agents: agentnames.append(agent.hostname) - combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] + combined_agents = combined_agents[combined_agents["hostname"].isin(agentnames)] - #Present and select rows + # Present and select rows selected_rows = Selector.select_dataframe_with_mode( combined_agents, - columns=['otpid', 'hostname', 'status','purpose','granted'], - header="OTP Sessions" + columns=["otpid", "hostname", "status", "purpose", "granted"], + header="OTP Sessions", ) for row in selected_rows: - otpid = row['otpid'] - hostname = row['hostname'] + otpid = row["otpid"] + hostname = row["hostname"] result = api.otp_revoke(otpid) logger.info(f"{hostname} (otpid: {otpid}):\n{result}") - - diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index f66eee7..6f4243c 100644 --- a/flows/prepPolicy.py +++ b/flows/prepPolicy.py @@ -44,7 +44,6 @@ logger = logging.getLogger(__name__) dotenv.load_dotenv() - def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] @@ -60,9 +59,18 @@ def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]: return selected if isinstance(selected, list) else [selected] -def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]: - if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()] - else: allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()] +def selectAllowlists( + api: AirlockAPIWrapper, policy=all, allow_multiple=True +) -> List[Allowlist]: + if policy == "all": + allowlists = [ + Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows() + ] + else: + allowlists = [ + Allowlist(**row.to_dict()) + for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows() + ] logger.debug("Prompting for Allowlist(s)") print(colorText("Please select allowlist(s)", "white")) selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) @@ -76,9 +84,7 @@ def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) def sortHashes( - api: AirlockAPIWrapper, - selected_policies: List[Policy], - type=[1, 2, 6, 7] + api: AirlockAPIWrapper, selected_policies: List[Policy], type=[1, 2, 6, 7] ): working_dir = load_env("WORKING_DIR") history_days = Selector.select_value( @@ -86,34 +92,41 @@ def sortHashes( value_type=int, valid_range=(1, 150), ) - + logger.debug(f"{history_days} day selected for history") - + if history_days is None: logging.warning("No history range selected. Aborting.") return - + policy_executions = ExecutionHistoryRecord.from_policies( api, selected_policies, type_=type, history_days=history_days ) logger.debug(f"Executions contains {policy_executions}") - 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) + 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) + ) categories = { - "needs_review": needs_review, - "approved": approved, - "unapproved": unapproved, - "leftover" : unknown - } + "needs_review": needs_review, + "approved": approved, + "unapproved": unapproved, + "leftover": unknown, + } - for label, records in categories.items(): if not records: - continue # Skip empty or falsy categories + continue # Skip empty or falsy categories csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_{label}_executions.csv" html_path = f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{label}.html" @@ -122,9 +135,9 @@ def sortHashes( 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) + 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) @@ -140,9 +153,11 @@ def buildPathsandPublishers(selected_policies: List[Policy], split): df1 = pd.DataFrame() df2 = pd.DataFrame() all_approved_hashes = pd.DataFrame() - path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" + path1 = ( + f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" + ) path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_needs_review_executions.csv" - path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int) + path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type=int) if os.path.exists(path1): df1 = pd.read_csv(path1) @@ -163,37 +178,48 @@ def buildPathsandPublishers(selected_policies: List[Policy], split): if "filename" in all_approved_hashes.columns: all_approved_hashes = all_approved_hashes.sort_values(by="filename") else: - logger.warning("Warning: 'filename' column not found in concatenated DataFrame.") + logger.warning( + "Warning: 'filename' column not found in concatenated DataFrame." + ) if not all_approved_hashes.empty and path_exclusion_constant: primary_path_exclusions = calculatePath( - all_approved_hashes, path_exclusion_constant, + all_approved_hashes, + path_exclusion_constant, split, ) remaining_hashes = all_approved_hashes[ ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) ] secondary_path_exclusions = calculatePath( - remaining_hashes,(path_exclusion_constant - 1), split + remaining_hashes, (path_exclusion_constant - 1), split ) remaining_hashes = remaining_hashes[ ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) ] dataframes = { - "all_approved_hashes" : all_approved_hashes, + "all_approved_hashes": all_approved_hashes, "primary_Paths": primary_path_exclusions, "secondary_Paths": secondary_path_exclusions, - "hashes_not_approvable_by_path": remaining_hashes + "hashes_not_approvable_by_path": remaining_hashes, } logger.debug("Preparing to sort dataframes") for name, df in dataframes.items(): logger.debug(f" DataFrame headers: {list(df.columns)}") - if "hashes" in name : df.sort_values(by="filename", inplace=True) - else: df.sort_values(by="longestcfp", inplace=True) - - df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", index=False) - formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html") + if "hashes" in name: + df.sort_values(by="filename", inplace=True) + else: + df.sort_values(by="longestcfp", inplace=True) + + df.to_csv( + f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_{name}.csv", + index=False, + ) + formatHTML( + df, + f"{working_dir}\\Needs_Review\\HTML\\{selected_policies[0].name}_{name}.html", + ) if not all_approved_hashes.empty: # Drop all not signed, only keep unique values @@ -201,14 +227,18 @@ def buildPathsandPublishers(selected_policies: List[Policy], split): all_approved_hashes["publisher"] != "Not Signed" ].drop_duplicates(subset=["publisher"]) # Remove Bad publisher if somehow they made it this far - pattern = regulator(load_env_json("BAD_PUBLISHERS","[]")) + pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]")) publist = publist[~publist["publisher"].str.contains(pattern, na=False)] publist = publist[["publisher"]] publist.sort_values(by="publisher", inplace=True) - publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", index=False) - else: + publist.to_csv( + f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_publishers.csv", + index=False, + ) + else: logger.debug("Approved Hashes list appears empty") + def buildPreflights(selected_policies: List[Policy]): working_dir = load_env("WORKING_DIR") @@ -222,8 +252,7 @@ def buildPreflights(selected_policies: List[Policy]): path2 = f"{working_dir}\\Approved\\{selected_policies[0].name}_secondary_Paths.csv" publishers = f"{working_dir}\\Approved\\{selected_policies[0].name}_publishers.csv" - - #Read in and combine the two path generations + # Read in and combine the two path generations if os.path.exists(path1): df1 = pd.read_csv(path1) else: @@ -239,19 +268,18 @@ def buildPreflights(selected_policies: List[Policy]): approved_paths = pd.DataFrame() else: approved_paths = pd.concat([df1, df2], ignore_index=True) - - approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep ="first") - #We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions. + approved_paths = approved_paths.drop_duplicates(subset="longestcfp", keep="first") + + # We create a list of hashes that are left over if we exclude the ones that are covered by the path exclusions. if os.path.exists(hash): hashes = pd.read_csv(hash) - approved_hashes = hashes[~hashes['filename'].isin(approved_paths['longestcfp'])] + approved_hashes = hashes[~hashes["filename"].isin(approved_paths["longestcfp"])] - approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep ="first") + approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep="first") else: - logger.warning(f"File not found: {hash}") - + logger.warning(f"File not found: {hash}") if os.path.exists(publishers): approved_publishers = pd.read_csv(publishers) @@ -259,19 +287,33 @@ def buildPreflights(selected_policies: List[Policy]): else: logger.warning(f"File not found: {publishers}") - dataframes = {"approved_paths": approved_paths, "approved_hashes": approved_hashes, "approved_publishers": approved_publishers} + dataframes = { + "approved_paths": approved_paths, + "approved_hashes": approved_hashes, + "approved_publishers": approved_publishers, + } for name, df in dataframes.items(): logger.debug(f" DataFrame headers: {list(df.columns)}") - if name == "approved_paths":df.sort_values(by="longestcfp", inplace=True) - elif name == "approved_hashes":df.sort_values(by="filename", inplace=True) - elif name == "approved_publishers" : df.sort_values(by="publisher", inplace=True) - - df.to_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", index=False) - formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html") + if name == "approved_paths": + df.sort_values(by="longestcfp", inplace=True) + elif name == "approved_hashes": + df.sort_values(by="filename", inplace=True) + elif name == "approved_publishers": + df.sort_values(by="publisher", inplace=True) + + df.to_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_{name}.csv", + index=False, + ) + formatHTML( + df, + f"{working_dir}\\Preflight\\HTML\\{selected_policies[0].name}_{name}.html", + ) + def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): - min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", 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)): @@ -281,7 +323,9 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): return parts # Diagnostic: log any non-string entries - non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))] + non_string_entries = df[ + ~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike))) + ] if not non_string_entries.empty: print(f"[WARNING] Non-string entries found in column '{col}':") print(non_string_entries) @@ -290,10 +334,14 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): split_paths = df[col].apply(clean_split) if min_files_for_path is not None: - df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy() + 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])) + df["group_key"] = split_paths.apply( + lambda parts: os.sep.join(parts[:path_exclusion_constant]) + ) grouped = df.groupby("group_key") new_rows = [] @@ -317,7 +365,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): for i, parts in enumerate(split_parts): filename = parts[-1] middle = ( - os.sep.join(parts[len(common_prefix):-1]) + os.sep.join(parts[len(common_prefix) : -1]) if len(parts) > len(common_prefix) + 1 else "" ) @@ -330,6 +378,7 @@ def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): return pd.DataFrame(new_rows).drop(columns=["group_key"]) + def calculatePath(approved_hashes, path_exclusion_constant, split): if split: dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] @@ -337,7 +386,7 @@ def calculatePath(approved_hashes, path_exclusion_constant, split): dfs_by_policy = [approved_hashes] badpathparts = load_env_json("BAD_PATH_PARTS", "[]") - min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int) + min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type=int) processed_dfs = [] @@ -364,7 +413,9 @@ def calculatePath(approved_hashes, path_exclusion_constant, split): ] unique_sha_counts = ( - lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index() + lcp_not_forbidden_review.groupby("longestcfp")["sha256"] + .nunique() + .reset_index() ) unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] @@ -380,51 +431,64 @@ def calculatePath(approved_hashes, path_exclusion_constant, split): return pathExclusions + def testChange(selected_policies, destination_policy, destination_allowlist): - working_dir = load_env("WORKING_DIR") + working_dir = load_env("WORKING_DIR") - logger.info("These path exclusions would be added to:") - logger.info(destination_policy) + logger.info("These path exclusions would be added to:") + logger.info(destination_policy) - pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv") - hashes = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv") + pathexclusions = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ) + hashes = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ) - unique_combinations = pathexclusions[["longestcfp", "file_extension"]].drop_duplicates() + unique_combinations = pathexclusions[ + ["longestcfp", "file_extension"] + ].drop_duplicates() - drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") - processed_paths = [ - (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" - for path, ext in unique_combinations.itertuples(index=False, name=None) - ] + drive_letter_pattern = re.compile(r"^[a-zA-Z]:\\") + processed_paths = [ + (path if drive_letter_pattern.match(path) else f"\\\\{path}") + f"\\**{ext}" + for path, ext in unique_combinations.itertuples(index=False, name=None) + ] - for path in processed_paths: - logger.info(path) + for path in processed_paths: + logger.info(path) - print(colorText("These publishers would added", "yellow")) - processed_publishers = [] - if os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv"): - publishers = pd.read_csv(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv") - if publishers.empty: - print(colorText("The publishers list is empty.", "red")) - else: - processed_publishers = ( - publishers[publishers["publisher"] != "Not Signed"] - ["publisher"] - .drop_duplicates() - .tolist() - ) - for publisher in processed_publishers: - print(publisher) + print(colorText("These publishers would added", "yellow")) + processed_publishers = [] + if os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ): + publishers = pd.read_csv( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_publishers.csv" + ) + if publishers.empty: + print(colorText("The publishers list is empty.", "red")) + else: + processed_publishers = ( + publishers[publishers["publisher"] != "Not Signed"]["publisher"] + .drop_duplicates() + .tolist() + ) + for publisher in processed_publishers: + print(publisher) - print(colorText("These hashes would be added to:", "yellow")) - print(destination_allowlist) + print(colorText("These hashes would be added to:", "yellow")) + print(destination_allowlist) - processed_hashes = hashes["sha256"].unique().tolist() - print_x_wide(processed_hashes, 3) + processed_hashes = hashes["sha256"].unique().tolist() + print_x_wide(processed_hashes, 3) - return processed_paths, processed_hashes, processed_publishers - -def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 into functions + return processed_paths, processed_hashes, processed_publishers + + +def menu_policy_enforce( + api: AirlockAPIWrapper, +): # TODO Need to clean up 6 and 7 into functions selected_policies = [] destination_policy = [] destination_allowlist = [] @@ -434,16 +498,22 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 working_dir = load_env("WORKING_DIR") while True: - printEnforceChecklist(selected_policies, destination_policy, destination_allowlist) + printEnforceChecklist( + selected_policies, destination_policy, destination_allowlist + ) choice = get_sanitized_input("\nEnter your choice: ") if choice == "1": clear_screen() - selected_policies = selectPolicies(api,True) + selected_policies = selectPolicies(api, True) elif choice == "2": clear_screen() - print(colorText("Please choose destination_name Policy for Path Exclusions", "white")) + print( + colorText( + "Please choose destination_name Policy for Path Exclusions", "white" + ) + ) destination_policy = selectPolicies(api, False) @@ -461,35 +531,53 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 elif choice == "4": clear_screen() - if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"): + if os.path.exists( + f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv" + ): buildPathsandPublishers(selected_policies, False) else: - print("File not found. Please make sure it's saved correctly and try again.") + print( + "File not found. Please make sure it's saved correctly and try again." + ) elif choice == "5": clear_screen() - if os.path.exists(f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv") and os.path.exists( + if os.path.exists( + f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv" + ) and os.path.exists( f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" ): buildPreflights(selected_policies) else: - print("File not found. Please make sure it's saved correctly and try again.") + print( + "File not found. Please make sure it's saved correctly and try again." + ) elif choice == "6": clear_screen() if ( - os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv") - and os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv") + os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ) + and os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ) and destination_policy and destination_allowlist ): - processed_paths, processed_hashes, processed_publishers = testChange(selected_policies, destination_policy, destination_allowlist) + processed_paths, processed_hashes, processed_publishers = testChange( + selected_policies, destination_policy, destination_allowlist + ) else: # Log which condition(s) failed missing_items = [] - if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"): + if not os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv" + ): missing_items.append("approved_paths.csv not found") - if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"): + if not os.path.exists( + f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv" + ): missing_items.append("approved_hashes.csv not found") if not destination_policy: missing_items.append("destination_policy is empty or None") @@ -513,13 +601,19 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 and confirmation.strip() == "I AGREE" ): print(colorText("Proceeding with the code...", "yellow")) - api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes) - api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths) + api.hash_add_to_allowlist( + destination_allowlist[0].applicationid, processed_hashes + ) + api.policy_add_path_exclusions( + destination_policy[0].groupid, processed_paths + ) if processed_publishers: - api.policy_add_publishers(destination_policy[0].groupid, processed_publishers) + api.policy_add_publishers( + destination_policy[0].groupid, processed_publishers + ) locked() - + else: logger.error("Confirmation block failed. Reasons:") if not processed_publishers or processed_hashes or processed_paths: @@ -529,30 +623,52 @@ def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 if not destination_allowlist: logger.error(" - `destination_allowlist` is missing or invalid.") if confirmation.strip() != "I AGREE": - logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip()) + logger.error( + " - User did not confirm with 'I AGREE'. Received: '%s'", + confirmation.strip(), + ) elif choice.upper() == "F": open_directory(working_dir) elif choice.upper() == "B": break - else: print(colorText("Invalid choice. Please try again.", "red")) + def section_header(title): - print(colorText("\n --------------------------------------------------------------------", "cyan")) + print( + colorText( + "\n --------------------------------------------------------------------", + "cyan", + ) + ) print(colorText(f" ------------- {title} -------------", "cyan")) - print(colorText(" --------------------------------------------------------------------", "cyan")) + print( + colorText( + " --------------------------------------------------------------------", + "cyan", + ) + ) def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): working_dir = load_env("WORKING_DIR") section_header("šŸ› ļø šŸ”’ Prepare to Enforce Policy šŸ› ļø šŸ”’") - print(colorText("\nSequentially follow these steps to prepare a policy for enforcement:", "white")) + print( + colorText( + "\nSequentially follow these steps to prepare a policy for enforcement:", + "white", + ) + ) # Step 1: Originating Policies - print(colorText("\n1. Choose which policy or policies to gather execution info from", "cyan")) + print( + colorText( + "\n1. Choose which policy or policies to gather execution info from", "cyan" + ) + ) if not selected_policies: print(colorText(" [āœ—] No policies have been chosen", "red")) else: @@ -561,67 +677,194 @@ def printEnforceChecklist(selected_policies, destination_policy, destination_all print(colorText(f" [āœ“] {policy.name}", "green")) # Step 2: Destination Policy and Allowlist - print(colorText("2. Choose the destination policy and associated allowlist", "cyan")) + print( + colorText("2. Choose the destination policy and associated allowlist", "cyan") + ) if destination_policy: - print(colorText(f" [āœ“] {destination_policy[0].name} has been selected as the destination policy", "green")) + print( + colorText( + f" [āœ“] {destination_policy[0].name} has been selected as the destination policy", + "green", + ) + ) else: print(colorText(" [āœ—] No destination policy has been chosen", "red")) if destination_allowlist: - print(colorText(f" [āœ“] {destination_allowlist[0].name} has been selected as allowlist", "green")) + print( + colorText( + f" [āœ“] {destination_allowlist[0].name} has been selected as allowlist", + "green", + ) + ) else: print(colorText(" [āœ—] No allowlist has been chosen", "red")) # Step 3: Data Preparation - print(colorText(f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review", "cyan")) + print( + colorText( + f"3. Select to begin pulling execution history. The executions will be sorted and placed in {working_dir}\\data\\Needs_Review", + "cyan", + ) + ) if selected_policies: policy_id = selected_policies[0].name review_path = f"{working_dir}\\Needs_Review\\Review_First\\{policy_id}_approved_executions.csv" - print(colorText(" [āœ“] Data has been fetched" if os.path.exists(review_path) else " [āœ—] Data has not been fetched", "green" if os.path.exists(review_path) else "red")) + print( + colorText( + ( + " [āœ“] Data has been fetched" + if os.path.exists(review_path) + else " [āœ—] Data has not been fetched" + ), + "green" if os.path.exists(review_path) else "red", + ) + ) else: - print(colorText(" [āœ—] No policies selected, cannot check data fetch status", "red")) + print( + colorText( + " [āœ—] No policies selected, cannot check data fetch status", "red" + ) + ) # Step 4: Manual Review print(colorText("4. Manually review the files:", "cyan")) - print(colorText(" Remove the rows containing hashes you do not approve of", "cyan")) - print(colorText(f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.", "cyan")) - print(colorText(" This will start the process to generate possible filepath approvals", "cyan")) + print( + colorText( + " Remove the rows containing hashes you do not approve of", "cyan" + ) + ) + print( + colorText( + f" When complete, save both csv files to {working_dir}\\data\\Approved and choose this option.", + "cyan", + ) + ) + print( + colorText( + " This will start the process to generate possible filepath approvals", + "cyan", + ) + ) if selected_policies: policy_id = selected_policies[0].name approved_path = f"{working_dir}\\Approved\\{policy_id}_approved_executions.csv" - second_review_path = f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" - print(colorText(" [āœ“] Reviewed hashes have been loaded" if os.path.exists(approved_path) else " [āœ—] Reviewed hashes have not been loaded", "green" if os.path.exists(approved_path) else "red")) - print(colorText(" [āœ“] Path review list created" if os.path.exists(second_review_path) else " [āœ—] Path review list has not been created", "green" if os.path.exists(second_review_path) else "red")) + second_review_path = ( + f"{working_dir}\\Needs_Review\\Review_Second\\{policy_id}_primary_Paths.csv" + ) + print( + colorText( + ( + " [āœ“] Reviewed hashes have been loaded" + if os.path.exists(approved_path) + else " [āœ—] Reviewed hashes have not been loaded" + ), + "green" if os.path.exists(approved_path) else "red", + ) + ) + print( + colorText( + ( + " [āœ“] Path review list created" + if os.path.exists(second_review_path) + else " [āœ—] Path review list has not been created" + ), + "green" if os.path.exists(second_review_path) else "red", + ) + ) else: - print(colorText(" [āœ—] No policies selected, cannot check reviewed hashes or path list", "red")) + print( + colorText( + " [āœ—] No policies selected, cannot check reviewed hashes or path list", + "red", + ) + ) # Step 5: Path Review - print(colorText(f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", "cyan")) - print(colorText(" Remove the rows containing path exclusions or publishers you do not approve of.", "cyan")) - print(colorText(f" When complete, save the files to {working_dir}\\data\\Approved", "cyan")) - print(colorText(" Choose this option when done to build your preflights", "cyan")) + print( + colorText( + f"5. Manually review the files in {working_dir}\\Needs_Review\\Review_Second\\", + "cyan", + ) + ) + print( + colorText( + " Remove the rows containing path exclusions or publishers you do not approve of.", + "cyan", + ) + ) + print( + colorText( + f" When complete, save the files to {working_dir}\\data\\Approved", + "cyan", + ) + ) + print( + colorText(" Choose this option when done to build your preflights", "cyan") + ) if selected_policies: policy_id = selected_policies[0].name reviewed_path = f"{working_dir}\\Approved\\{policy_id}_primary_Paths.csv" preflight_paths = f"{working_dir}\\Preflight\\{policy_id}_approved_paths.csv" preflight_hashes = f"{working_dir}\\Preflight\\{policy_id}_approved_hashes.csv" - print(colorText(" [āœ“] Reviewed path list detected" if os.path.exists(reviewed_path) else " [āœ—] Path review list has not been detected", "green" if os.path.exists(reviewed_path) else "red")) - preflight_ready = os.path.exists(preflight_paths) and os.path.exists(preflight_hashes) - print(colorText(" [āœ“] Preflight Path Exclusion List has been generated" if preflight_ready else " [āœ—] Preflight Path Exclusion List has not been generated", "green" if preflight_ready else "red")) + print( + colorText( + ( + " [āœ“] Reviewed path list detected" + if os.path.exists(reviewed_path) + else " [āœ—] Path review list has not been detected" + ), + "green" if os.path.exists(reviewed_path) else "red", + ) + ) + preflight_ready = os.path.exists(preflight_paths) and os.path.exists( + preflight_hashes + ) + print( + colorText( + ( + " [āœ“] Preflight Path Exclusion List has been generated" + if preflight_ready + else " [āœ—] Preflight Path Exclusion List has not been generated" + ), + "green" if preflight_ready else "red", + ) + ) else: - print(colorText(" [āœ—] No policies selected, cannot check preflight status", "red")) + print( + colorText( + " [āœ—] No policies selected, cannot check preflight status", "red" + ) + ) # Final Steps - print(colorText("6. Test ------------------------------------------------------", "cyan")) - print(colorText(" Prints to console the changes that would be made, must be done to proceed. ", "cyan")) + print( + colorText( + "6. Test ------------------------------------------------------", "cyan" + ) + ) + print( + colorText( + " Prints to console the changes that would be made, must be done to proceed. ", + "cyan", + ) + ) - print(colorText("7. Liftoff ------------------------------------------------------", "cyan")) - print(colorText(" Apply path exclusions and approved publishers to selected policy", "cyan")) + print( + colorText( + "7. Liftoff ------------------------------------------------------", "cyan" + ) + ) + print( + colorText( + " Apply path exclusions and approved publishers to selected policy", + "cyan", + ) + ) print(colorText(" Apply approved hashes to allowlist", "cyan")) - # Utility Options print(colorText("F. šŸ“‚ - Open Working Directory", "cyan")) print(colorText("B. šŸ”š - Back", "cyan")) diff --git a/flows/quietAgent.py b/flows/quietAgent.py index 6bc05ea..d723da7 100644 --- a/flows/quietAgent.py +++ b/flows/quietAgent.py @@ -51,20 +51,22 @@ def findQuietAgents(api: AirlockAPIWrapper): valid_range=(1, 150), ) - confirm = Selector.confirm(f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : ") + confirm = Selector.confirm( + f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N : " + ) # Get execution history as a DataFrame if confirm: policy_exec_history = getPolicyInfo( - api, selected_policy[0], [1, 2, 6, 7], history_days + api, selected_policy[0], [1, 2, 6, 7], history_days ) - - + if policy_exec_history.empty: - logging.info("No execution history found for the selected policy and time range.") + logging.info( + "No execution history found for the selected policy and time range." + ) get_sanitized_input("Press enter to continue") return - # Convert 'datetime' column to timezone-aware datetime objects policy_exec_history["datetime"] = pd.to_datetime( policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True @@ -82,12 +84,14 @@ def findQuietAgents(api: AirlockAPIWrapper): hostname_counts = policy_exec_history["hostname"].value_counts() # Map execution counts to agents - agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int) + agents["execution_count"] = ( + agents["hostname"].map(hostname_counts).fillna(0).astype(int) + ) # Find most recent execution per hostname - most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates( - subset="hostname", keep="first" - ) + most_recent_exec = policy_exec_history.sort_values( + by="days_ago" + ).drop_duplicates(subset="hostname", keep="first") # Map most recent execution age to agents agents["days_since"] = agents["hostname"].map( @@ -101,7 +105,9 @@ def findQuietAgents(api: AirlockAPIWrapper): ) # Sort agents by execution count and hostname - agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True]) + agents = agents.sort_values( + by=["execution_count", "hostname"], ascending=[True, True] + ) # Save to CSV filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv" @@ -116,8 +122,7 @@ def findQuietAgents(api: AirlockAPIWrapper): ready_percentage = (ready_agents / total_agents) * 100 # Print results - - + message = ( f"Total agents: {total_agents}\n" f"Agents marked as 'enforce_ready': {ready_agents}\n" @@ -125,5 +130,5 @@ def findQuietAgents(api: AirlockAPIWrapper): f"Percentage ready for enforcement: {ready_percentage:.2f}%" ) logger.debug(message) - colorText(message,"green") + colorText(message, "green") get_sanitized_input("Press enter to continue") diff --git a/loading.png b/loading.png index 1dc81bd..a10998c 100644 Binary files a/loading.png and b/loading.png differ diff --git a/models/agent.py b/models/agent.py index 9a4c580..a41a284 100644 --- a/models/agent.py +++ b/models/agent.py @@ -38,23 +38,19 @@ class Agent: status_text: Optional[str] = field(default=None) # Class-level status map - status_map: ClassVar[dict] = { - 0: "Offline", - 1: "Online", - 2: "Hidden", - 3: "Safemode" - } - + status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"} def enrich_with_policies(self, policies: List[Policy]): - """Enrich the agent with groupname and human-readable status.""" - self.status_text = self.status_map.get(self.status, "Unknown") - for policy in policies: - if policy.groupid == self.groupid: - self.groupname = policy.name - break - if not self.groupname: - self.groupname = "Unknown" + """Enrich the agent with groupname and human-readable status.""" + self.status_text = self.status_map.get(self.status, "Unknown") + for policy in policies: + if policy.groupid == self.groupid: + self.groupname = policy.name + break + if not self.groupname: + self.groupname = "Unknown" + + """ from models.agent import Agent diff --git a/models/execution.py b/models/execution.py index 13281d1..e1cf514 100644 --- a/models/execution.py +++ b/models/execution.py @@ -28,7 +28,6 @@ import pandas as pd import airlock_libs 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 @@ -36,11 +35,13 @@ logger = logging.getLogger(__name__) dotenv.load_dotenv() + @dataclass class Hash: """ Hash model representing Hash data """ + sha256: str applications: str baselines: str @@ -62,7 +63,7 @@ class Hash: sha384: str sha512: str at_decision: Optional[str] = None - + def to_dict(self): return asdict(self) @@ -100,11 +101,15 @@ class Hash: for hash_obj in hashes: publisher = hash_obj.publisher or "" description = hash_obj.description or "" - reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + 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}") + logger.debug( + f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}" + ) # 1. Unapproved: bad publisher or PUP if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): @@ -128,9 +133,9 @@ class Hash: # 3. Approved or Unapproved based on threat level try: - score = int(scannermatch) # pyright: ignore[reportArgumentType] + score = int(scannermatch) # pyright: ignore[reportArgumentType] logger.debug(f"Parsed scannermatch score: {score}") - if score > threat_tolerance: # pyright: ignore[reportOperatorIssue] + if score > threat_tolerance: # pyright: ignore[reportOperatorIssue] logger.debug("Unapproved: Unsigned file with high threat score.") hash_obj.at_decision = "unapproved" unapproved_count += 1 @@ -139,15 +144,17 @@ class Hash: hash_obj.at_decision = "approved" approved_count += 1 except (ValueError, TypeError): - logger.debug("Needs Review: Scannermatch score is missing or invalid. ā€” {e}") + 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: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}") + logger.debug( + f"Final counts ā€” Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}" + ) return hashes - @classmethod def export_to_csv(cls, hash_list, directory_path): """ @@ -204,11 +211,12 @@ class ExecutionHistoryRecord: localip: Optional[str] = None extid: Optional[str] = None extname: Optional[str] = None - exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension + exttype: Optional[int] = ( + None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension + ) extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge hash_obj: Optional[Hash] = None - @classmethod def from_dict(cls, data: dict): mandatory_fields = [ @@ -225,7 +233,9 @@ class ExecutionHistoryRecord: "datetime", ] missing_fields = [ - field for field in mandatory_fields if field not in data or data[field] is None + field + for field in mandatory_fields + if field not in data or data[field] is None ] if missing_fields: raise ValueError(f"Missing mandatory fields: {missing_fields}") @@ -255,7 +265,7 @@ class ExecutionHistoryRecord: extname=data.get("extname"), exttype=data.get("exttype"), extbrowser=data.get("extbrowser"), - hash_obj=data.get("hash_obj") + hash_obj=data.get("hash_obj"), ) @classmethod @@ -264,7 +274,9 @@ class ExecutionHistoryRecord: ) -> List["ExecutionHistoryRecord"]: executions = [] for policy in selected_policies: - execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days) + execs = airlock_libs.pull_policy_exec_histories( + api, policy.name, str([1, 2, 6, 7]), history_days + ) if execs: data = json.loads(execs) exechistories = data.get("response", {}).get("exechistories", []) @@ -275,8 +287,12 @@ class ExecutionHistoryRecord: df = df.drop_duplicates(subset=["sha256", "filename", "hostname"]) df = df.sort_values(by=["sha256", "filename"]) - executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()]) - logger.debug(f"Staging of Execution history for policy: {policy.name} is complete") + executions.extend( + [cls.from_dict(row.to_dict()) for _, row in df.iterrows()] + ) + logger.debug( + f"Staging of Execution history for policy: {policy.name} is complete" + ) print( colorText( f"Staging of Execution history for policy: {policy.name} is complete", @@ -285,20 +301,23 @@ class ExecutionHistoryRecord: ) return executions - + @staticmethod def enrich_with_hashes( - api: AirlockAPIWrapper, - executions: List["ExecutionHistoryRecord"] + 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.") + 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.") + logger.warning( + "No sha256 values found in execution records. Skipping enrichment." + ) return executions logger.debug("Querying hash data from API...") @@ -307,8 +326,10 @@ class ExecutionHistoryRecord: hash_objects = [] required_fields = { - f.name for f in dataclasses.fields(Hash) - if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING + 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()): @@ -346,11 +367,15 @@ class ExecutionHistoryRecord: exec_record.hash_obj = hash_obj enriched_count += 1 - logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.") + 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"]: + def categorize_executions_by_hash_decision( + executions: List["ExecutionHistoryRecord"], + ) -> List["ExecutionHistoryRecord"]: """ Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation. @@ -374,11 +399,15 @@ class ExecutionHistoryRecord: publisher = hash_obj.publisher or "" description = hash_obj.description or "" - reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + 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}") + logger.debug( + f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}" + ) # 1. Unapproved: bad publisher or PUP if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): @@ -402,7 +431,7 @@ class ExecutionHistoryRecord: # 3. Approved or Unapproved based on threat level try: - score = int(scannermatch) # pyright: ignore[reportArgumentType] + 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.") @@ -413,7 +442,9 @@ class ExecutionHistoryRecord: 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}") + logger.debug( + f"Needs Review: Scannermatch score is missing or invalid. ā€” {e}" + ) hash_obj.at_decision = "needs_review" needs_review_count += 1 @@ -423,12 +454,14 @@ class ExecutionHistoryRecord: ) return executions - - + @classmethod - def sort_by_hash_decision( - cls, executions: List["ExecutionHistoryRecord"] - ) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]: + 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. @@ -454,7 +487,9 @@ class ExecutionHistoryRecord: else: unknown.append(record) - logger.info(f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:") + logger.info( + f"[ExecutionHistoryRecord] Sorted {len(sorted_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)}") @@ -463,7 +498,6 @@ class ExecutionHistoryRecord: return approved, unapproved, needs_review, unknown - """ executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30) diff --git a/models/policy.py b/models/policy.py index 7040f32..dbfb357 100644 --- a/models/policy.py +++ b/models/policy.py @@ -29,7 +29,9 @@ class Policy: def __repr__(self): # Show all current attributes, including dynamically added ones - attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + attrs = ", ".join( + f"{key}={repr(value)}" for key, value in self.__dict__.items() + ) return f"" def to_dict(self): @@ -53,7 +55,9 @@ class Allowlist: def __repr__(self): # Show all current attributes, including dynamically added ones - attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + attrs = ", ".join( + f"{key}={repr(value)}" for key, value in self.__dict__.items() + ) return f"" def to_dict(self): diff --git a/services/API.py b/services/API.py index e5a3c76..34ce788 100644 --- a/services/API.py +++ b/services/API.py @@ -23,7 +23,6 @@ import requests logger = logging.getLogger(__name__) - class AirlockAPIWrapper: """ A wrapper class for interacting with the Airlock API. @@ -141,25 +140,25 @@ class AirlockAPIWrapper: payload = {"status": "0"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_enforced(self) -> pd.DataFrame: """Find OTPs that are awaiting activation.""" payload = {"status": "2"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_revoked(self) -> pd.DataFrame: """Find OTPs that are awaiting activation.""" payload = {"status": "3"} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_find_by_agent(self, agentid) -> pd.DataFrame: """Find OTP by agent.""" payload = {"agentid": agentid} result = self._post("/v1/otp/usage", payload) return pd.DataFrame(result["response"]["otpusage"]) - + def otp_generate(self, agentid: str, duration: int, purpose: str) -> str: """Generate a new OTP for an agent.""" payload = { @@ -175,7 +174,7 @@ class AirlockAPIWrapper: payload = {"otpid": otpid} result = self._post("/v1/otp/activities", payload) return pd.DataFrame(result["response"]["otpactivities"]) - + def otp_revoke(self, otpid: str) -> dict: """ Revoke an active OTP. @@ -186,18 +185,17 @@ class AirlockAPIWrapper: """ payload = {"otpid": otpid} return self._post("/v1/otp/revoke", payload) - - def otp_validate(self, otpcode: str) -> dict: - """ - Validate an OTP code. - Parameters: - - otpcode (str): The OTP code to validate. - Returns: - - dict: JSON response indicating validity. - """ - payload = {"otpcode": otpcode} - return self._post("/v1/otp/validate", payload) + def otp_validate(self, otpcode: str) -> dict: + """ + Validate an OTP code. + Parameters: + - otpcode (str): The OTP code to validate. + Returns: + - dict: JSON response indicating validity. + """ + payload = {"otpcode": otpcode} + return self._post("/v1/otp/validate", payload) # Policy Management def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict: @@ -225,7 +223,7 @@ class AirlockAPIWrapper: payload = {"groupid": groupid} result = self._post("/v1/group/agents", payload) return pd.DataFrame(result["response"]["agents"]) - + def policy_list_allowlists(self, groupid: str) -> pd.DataFrame: """List allowlists assigned to a specific policy group.""" payload = {"groupid": groupid} @@ -236,31 +234,37 @@ class AirlockAPIWrapper: """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" payload = {"groupid": groupid, "auditmode": auditmode} return self._post("/v1/group/settings/auditmode", payload) - - def policy_set_script_custom(self, - groupid: str, - script_custom: int, - scripts_audit: List[str], - scripts_disabled: List[str], - scripts_respect: List[str], - ) -> dict: + + def policy_set_script_custom( + self, + groupid: str, + script_custom: int, + scripts_audit: List[str], + scripts_disabled: List[str], + scripts_respect: List[str], + ) -> dict: """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" - payload = {"groupid": groupid, - "script_custom": script_custom, - "scripts_audit": scripts_audit, - "scripts_disabled": scripts_disabled, - "scripts_respect": scripts_respect - } + payload = { + "groupid": groupid, + "script_custom": script_custom, + "scripts_audit": scripts_audit, + "scripts_disabled": scripts_disabled, + "scripts_respect": scripts_respect, + } return self._post("/v1/group/settings/script_custom", payload) # Execution History - def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> str: + def history_logging( + self, type: List[str], checkpoint: str, policy: List[str] + ) -> str: """Retrieve execution history logs.""" payload = {"type": type, "checkpoint": checkpoint, "policy": policy} result = self._post("/v1/logging/exechistories", payload) return result["response"]["exechistories"] - def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]: + def history_execution( + self, today: str, date_selected: str, agent_name: str + ) -> List[Dict]: """ Retrieve execution history logs. diff --git a/services/agenthandler.py b/services/agenthandler.py index 531fc41..0ec3d3b 100644 --- a/services/agenthandler.py +++ b/services/agenthandler.py @@ -47,7 +47,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): print(colorText("No agents selected or invalid history range.", "red")) return - historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d") + historical_date = (datetime.now() - timedelta(days=history_days)).strftime( + "%Y-%m-%d" + ) today = datetime.now().strftime("%Y-%m-%d") all_history = [] @@ -56,7 +58,11 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): try: exechistory = api.history_execution(today, historical_date, agent.hostname) except Exception as e: - print(colorText(f"āŒ Error retrieving history for {agent.hostname}: {e}", "red")) + print( + colorText( + f"āŒ Error retrieving history for {agent.hostname}: {e}", "red" + ) + ) continue if isinstance(exechistory, list): @@ -76,7 +82,9 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool): print(colorText(f"{key}: {value}", "green")) print("\n") else: - print(colorText(f"No execution history found for {agent.hostname}.", "yellow")) + print( + colorText(f"No execution history found for {agent.hostname}.", "yellow") + ) if outputjson: print(json.dumps(all_history, indent=2)) @@ -92,6 +100,7 @@ def findAllAgents(api): return agents + def findAgents(api, return_dataframe): agents = selectAgents(api) working_dir = load_env("WORKING_DIR") @@ -113,8 +122,14 @@ def findAgents(api, return_dataframe): print(agent_df) logging.debug("Displayed DataFrame to console.") - user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower() - if user_input == 'y': + user_input = ( + get_sanitized_input( + "\nWould you like to export the results to a CSV file? (y/n): " + ) + .strip() + .lower() + ) + if user_input == "y": timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") filename = f"agentsearch_{timestamp}.csv" file_path = os.path.join(str(working_dir), filename) @@ -131,17 +146,27 @@ def findAgents(api, return_dataframe): else: logging.debug("User declined to export the DataFrame.") + def collect_device_names() -> List[str]: 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 (Three times if you have a single device).\n", "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 (Three times if you have a single device).\n", + "cyan", + ) + ) print(colorText("Example:", "cyan")) print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan")) print(colorText("Paste or type your device names below:", "white")) device_input_lines = [] empty_line_count = 0 - valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$') + valid_line_pattern = re.compile(r"^[a-zA-Z0-9_\- ]+$") while True: line = get_sanitized_input("") @@ -158,7 +183,12 @@ def collect_device_names() -> List[str]: if valid_line_pattern.match(stripped_line): device_input_lines.append(stripped_line) else: - print(colorText(f"āš ļø Invalid input: '{stripped_line}' ā€” only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow")) + print( + colorText( + f"āš ļø Invalid input: '{stripped_line}' ā€” only letters, numbers, underscores, spaces, and hyphens are allowed.", + "yellow", + ) + ) return [name for name in device_input_lines if name] @@ -168,10 +198,13 @@ def choose_match_type() -> bool: return get_sanitized_input("").strip().lower() in ["y", "yes"] -def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']: +def match_agents( + device_names: List[str], agents: List["Agent"], use_exact: bool +) -> List["Agent"]: if use_exact: return [ - agent for agent in agents + agent + for agent in agents if agent.hostname.lower() in [name.lower() for name in device_names] ] else: @@ -180,23 +213,38 @@ def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool return [agent for agent in agents if regex.search(agent.hostname)] -def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool): +def show_unmatched( + device_names: List[str], matched_agents: List["Agent"], use_exact: bool +): if use_exact: - unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)] + unmatched = [ + name + for name in device_names + if not any( + agent.hostname.lower() == name.lower() for agent in matched_agents + ) + ] else: - unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)] + unmatched = [ + name + for name in device_names + if not any( + re.search(re.escape(name), agent.hostname, re.IGNORECASE) + for agent in matched_agents + ) + ] if unmatched: logger.debug(f"āš ļø No matches for: {', '.join(unmatched)}") print(colorText(f"āš ļø No matches for: {', '.join(unmatched)}", "yellow")) -def enrich_agents(agents: List['Agent'], policies: List['Policy']): +def enrich_agents(agents: List["Agent"], policies: List["Policy"]): for agent in agents: agent.enrich_with_policies(policies) -def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: +def selectAgents(api: "AirlockAPIWrapper") -> List["Agent"]: device_names = collect_device_names() if not device_names: logger.debug("No device names entered") @@ -227,11 +275,11 @@ def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: if idx < len(matched_agents): line += f"{matched_agents[idx].hostname:<30}" logger.info(line) - + matched_agents = Selector.select_with_mode( matched_agents, label_func=lambda agent: agent.hostname, - header="Matched Devices:" + header="Matched Devices:", ) if not matched_agents: @@ -263,11 +311,17 @@ def moveAgentToRelatedPolicy( if agent.groupid in policy_relationship_map: target_policy = policy_relationship_map[agent.groupid] elif agent.groupid in policy_relationship_map.values(): - logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.") - print(f"Agent {agent.hostname} is already in an audit group. No action needed.") + logger.debug( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) + print( + f"Agent {agent.hostname} is already in an audit group. No action needed." + ) return else: - logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.") + logger.warning( + f"Error: No corresponding audit policy found for groupid: {agent.groupid}." + ) return elif mode == "enforcement": @@ -275,10 +329,14 @@ def moveAgentToRelatedPolicy( if agent.groupid in inverse_map: target_policy = inverse_map[agent.groupid] elif agent.groupid in inverse_map.values(): - logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.") + logger.info( + f"Agent {agent.hostname} is already in an enforcement group. No action needed." + ) return else: - logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.") + logger.warning( + f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}." + ) return else: @@ -299,25 +357,32 @@ def toggleEnforcement(api: AirlockAPIWrapper): devices = selectAgents(api) for device in devices: print(device.hostname) - confirm = Selector.confirm("Would you like to continue with these devices? Y/N: ") + confirm = Selector.confirm( + "Would you like to continue with these devices? Y/N: " + ) if direction and devices and confirm: for device in devices: - result = moveAgentToRelatedPolicy(api,device, str(direction).lower()) + result = moveAgentToRelatedPolicy(api, device, str(direction).lower()) logger.info(f"{device.hostname}: result: {result}") get_sanitized_input("Press enter to continue") + def moveAgents(api: AirlockAPIWrapper): devices = selectAgents(api) for device in devices: print(device.hostname) - confirm_devices = Selector.confirm("Would you like to continue with these devices? Y/N: ") + confirm_devices = Selector.confirm( + "Would you like to continue with these devices? Y/N: " + ) if devices and confirm_devices: policies = selectPolicies(api, False) - confirm_move = Selector.confirm(f"Would you like to move these devices to {policies[0].name}?") + confirm_move = Selector.confirm( + f"Would you like to move these devices to {policies[0].name}?" + ) if confirm_move: for device in devices: result = api.agent_move(device.agentid, policies[0].groupid) logger.info(f"{device.hostname}: result: {result}") else: logger.info("Exiting without change") - get_sanitized_input("Press enter to continue") \ No newline at end of file + get_sanitized_input("Press enter to continue") diff --git a/services/policyhandler.py b/services/policyhandler.py index 5acc422..c32cdfc 100644 --- a/services/policyhandler.py +++ b/services/policyhandler.py @@ -34,7 +34,6 @@ from utils.utils import areYouSure, colorText, get_sanitized_input logger = logging.getLogger(__name__) - def pullPolicyExechistories( api: AirlockAPIWrapper, policy: Policy, @@ -72,7 +71,7 @@ def pullPolicyExechistories( ) as pbar: while True: histories = api.history_logging( - type=type, checkpoint=checkpoint, policy= [policy.name] + type=type, checkpoint=checkpoint, policy=[policy.name] ) # Ensure histories is a list of dictionaries @@ -98,13 +97,17 @@ def pullPolicyExechistories( # Update checkpoint on last item if index == len(histories) - 1: - checkpoint = history_item["checkpoint"] # pyright: ignore[reportArgumentType] + checkpoint = history_item[ + "checkpoint" + ] # pyright: ignore[reportArgumentType] filebar.desc = f"Checkpoint Progress: {checkpoint}" break try: history_date = datetime.datetime.strptime( - history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportArgumentType] + history_item["datetime"].replace( + " +0000 UTC", "" + ), # pyright: ignore[reportArgumentType] "%Y-%m-%dT%H:%M:%SZ", ).date() except ValueError: @@ -177,8 +180,10 @@ def pullPolicyExechistories( def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): + import airlock_libs + executionhist_policy = pd.DataFrame() - exehist = pullPolicyExechistories(api, policy, type, days, True) + exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days) if exehist is not None: data = json.loads(exehist) executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) @@ -203,7 +208,7 @@ def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): executionhist_policy = executionhist_policy.sort_values( by=["sha256", "filename"] ) - logger.debug( f"Staging of Execution history for policy: {policy} is complete") + logger.debug(f"Staging of Execution history for policy: {policy} is complete") print( colorText( f"Staging of Execution history for policy: {policy} is complete", @@ -235,10 +240,10 @@ def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper): for enforcement_policy, audit_policy in policy_relationship_map.items(): api.policy_clone(enforcement_policy, audit_policy) api.policy_set_auditmode(audit_policy, "1") - + def confirmUpdateAfromE(api: AirlockAPIWrapper): areYouSure() confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") if confirmation.strip() == "I AGREE": - updateAuditPoliciesFromEnforcementPolices(api) \ No newline at end of file + updateAuditPoliciesFromEnforcementPolices(api) diff --git a/services/security.py b/services/security.py index aa95929..70229e2 100644 --- a/services/security.py +++ b/services/security.py @@ -28,8 +28,8 @@ import keyring # Constants KDF_ITERATIONS = 200_000 -SALT_SIZE = 16 # 128-bit Salt -NONCE_SIZE = 12 # AES-GCM +SALT_SIZE = 16 # 128-bit Salt +NONCE_SIZE = 12 # AES-GCM KEY_SIZE = 32 # AES-256 logger = logging.getLogger(__name__) @@ -49,9 +49,11 @@ def configure_keyring_backend(): system = platform.system() if system == "Windows": import keyring.backends.Windows + keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring()) elif system == "Linux": import keyring.backends.kwallet + keyring.set_keyring(keyring.backends.kwallet.DBusKeyring()) else: raise EnvironmentError(f"Unsupported OS: {system}") @@ -68,8 +70,9 @@ def store_api_key(service: str, username: str, api_key: str, password: str): b64 = base64.b64encode(blob).decode() keyring.set_password(service, username, b64) - - logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.") + logger.debug( + f"API key for service '{service}' and user '{username}' stored successfully." + ) print("\nāœ… API key stored securely.") print("The program will now exit. Press Enter to continue...") @@ -90,8 +93,8 @@ def retrieve_api_key(service: str, username: str, password: str) -> str: raise ValueError("No stored secret for this service/username.") blob = base64.b64decode(b64) salt = blob[:SALT_SIZE] - nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE] - ct = blob[SALT_SIZE + NONCE_SIZE:] + nonce = blob[SALT_SIZE : SALT_SIZE + NONCE_SIZE] + ct = blob[SALT_SIZE + NONCE_SIZE :] key = _derive_key(password.encode(), salt) aesgcm = AESGCM(key) pt = aesgcm.decrypt(nonce, ct, associated_data=None) @@ -124,7 +127,9 @@ def getAPI(USERNAME, SERVICE_NAME): if api_key_exists(SERVICE_NAME, USERNAME): for attempt in range(1, 4): - password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ") + password = getpass( + f"Attempt {attempt}/3 - Enter password to unlock your API key: " + ) try: apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password) logging.debug("API key successfully retrieved.") @@ -134,9 +139,15 @@ def getAPI(USERNAME, SERVICE_NAME): logging.error("Failed to retrieve API key after 3 incorrect attempts.") raise ValueError("Failed to retrieve API key after 3 incorrect attempts.") else: - logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.") - api_key = getpass(f"No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip() - print("Please exit and relaunch program after saving your credential to avoid errors") + logging.warning( + f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'." + ) + api_key = getpass( + f"No API key found. Please enter your API key for '{SERVICE_NAME}': " + ).strip() + print( + "Please exit and relaunch program after saving your credential to avoid errors" + ) while True: password = getpass("Create a password to encrypt your API key: ") @@ -155,7 +166,9 @@ def getAPI(USERNAME, SERVICE_NAME): logging.error(f"Failed to store API key: {e}") break else: - logging.warning("Password does not meet complexity requirements. Try again.") + logging.warning( + "Password does not meet complexity requirements. Try again." + ) class APIKeyManager: @@ -169,4 +182,4 @@ class APIKeyManager: def get(cls) -> str: if cls._api_key is None: raise ValueError("API key not loaded. Call APIKeyManager.load() first.") - return cls._api_key \ No newline at end of file + return cls._api_key diff --git a/utils/configmanager.py b/utils/configmanager.py index 62c2bdf..b099fa3 100644 --- a/utils/configmanager.py +++ b/utils/configmanager.py @@ -29,14 +29,15 @@ PROTECTED_KEYS = [ "PATH_EXCLUSION_CONST", "MIN_FILES_FOR_PATH", "VT_THREAT_TOLERANCE", - "POLICY_MAP_ENF_AUD" + "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_dir = Path(getattr(sys, "_MEIPASS", "")) bundled_path = bundled_dir / "system_config.json" if bundled_path.exists(): return bundled_path @@ -44,6 +45,7 @@ def get_system_config_path() -> Path: # Fallback to external location return Path(__file__).parent.parent / "system_config.json" + def load_protected_config() -> dict: global _protected_config try: @@ -52,19 +54,20 @@ def load_protected_config() -> dict: except FileNotFoundError: logging.warning("āš ļø system_config.json not found. Using built-in defaults.") system_config = { - "APPNAME": "AirlockTools", + "APPNAME": "Loxide", "PATH_EXCLUSION_CONST": 4, "MIN_FILES_FOR_PATH": 4, "VT_THREAT_TOLERANCE": 4, - "POLICY_MAP_ENF_AUD": { - "enforced_id": "audit_id" - } + "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]: + +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.") @@ -74,9 +77,12 @@ def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: 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__}.") + 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): @@ -85,13 +91,11 @@ def get_protected_json(key: str, default: str = "{}") -> dict: return json.loads(raw) except json.JSONDecodeError: try: - escaped = raw.encode('unicode_escape').decode('utf-8') + 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): @@ -100,13 +104,16 @@ def load_env_json(key: str, default: str): return json.loads(raw) except json.JSONDecodeError: try: - escaped = raw.encode('unicode_escape').decode('utf-8') + 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]: + +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. @@ -126,5 +133,7 @@ def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] 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 + logger.warning( + f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}." + ) + return default diff --git a/utils/selector.py b/utils/selector.py index 714e197..0a0f498 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -21,9 +21,12 @@ from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) + class Selector: @staticmethod - def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]: + def _get_sorted_items( + items: List[Any], label_func: Callable[[Any], str] + ) -> List[Any]: return sorted(items, key=lambda item: label_func(item).lower()) @staticmethod @@ -31,10 +34,10 @@ class Selector: items: List[Any], label_func: Callable[[Any], str], num_columns: int = 4, - header: str = "Available Choices:" + header: str = "Available Choices:", ) -> None: # Force single column if items are DataFrame rows - + if items and isinstance(items[0], (pd.Series, dict)): num_columns = 1 @@ -51,9 +54,7 @@ class Selector: @staticmethod def _display_selected_items( - selected: List[Any], - label_func: Callable[[Any], str], - num_columns: int = 4 + selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4 ) -> None: print(colorText("\nCurrent selections:", "cyan")) if not selected: @@ -92,7 +93,7 @@ class Selector: allow_multiple: bool = False, prompt_each: bool = False, header: str = "Available Choices:", - num_columns: int = 4 + num_columns: int = 4, ) -> Union[Optional[Any], List[Any]]: if not items: logger.warning("No items available for selection.") @@ -104,9 +105,19 @@ class Selector: if allow_multiple: while True: - Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header) - Selector._display_selected_items(selected, label_func, num_columns=num_columns) - choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower() + Selector._display_choices( + remaining_items, label_func, num_columns=num_columns, header=header + ) + Selector._display_selected_items( + selected, label_func, num_columns=num_columns + ) + choice = ( + get_sanitized_input( + "Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: " + ) + .strip() + .lower() + ) if choice == "q": break elif choice == "r": @@ -125,10 +136,14 @@ class Selector: logger.info(f"Selected: {label_func(item)}") else: logger.warning("Item already selected.") - remaining_items = [item for item in remaining_items if item not in newly_selected] + remaining_items = [ + item for item in remaining_items if item not in newly_selected + ] return selected if selected else None else: - Selector._display_choices(full_sorted_items, label_func, num_columns=num_columns, header=header) + Selector._display_choices( + full_sorted_items, label_func, num_columns=num_columns, header=header + ) try: choice = int(get_sanitized_input("Select one item by number: ")) if 1 <= choice <= len(full_sorted_items): @@ -145,9 +160,14 @@ class Selector: def select_with_mode( items: List[Any], label_func: Callable[[Any], str], - header: str = "Available Choices:" + header: str = "Available Choices:", ) -> List[Any]: - print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) + print( + colorText( + "Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", + "white", + ) + ) mode = get_sanitized_input("").strip().lower() if mode == "a": return items @@ -156,7 +176,7 @@ class Selector: label_func=label_func, allow_multiple=True, prompt_each=False, - header=header + header=header, ) if not selected: return items @@ -172,44 +192,38 @@ class Selector: @staticmethod def select_objects( - objects: List[Any], - allow_multiple: bool = False, - prompt_each: bool = False + objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[Any], List[Any]]: return Selector._select_from_list( objects, label_func=lambda obj: getattr(obj, "name", str(obj)), allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Objects:" + header="Available Objects:", ) @staticmethod def select_string( - options: List[str], - allow_multiple: bool = False, - prompt_each: bool = False + options: List[str], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[str], List[str]]: return Selector._select_from_list( options, label_func=str, allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Options:" + header="Available Options:", ) @staticmethod def select_int( - options: List[int], - allow_multiple: bool = False, - prompt_each: bool = False + options: List[int], allow_multiple: bool = False, prompt_each: bool = False ) -> Union[Optional[int], List[int]]: return Selector._select_from_list( options, label_func=lambda x: str(x), allow_multiple=allow_multiple, prompt_each=prompt_each, - header="Available Integers:" + header="Available Integers:", ) @staticmethod @@ -217,7 +231,7 @@ class Selector: prompt: str, value_type: type = int, valid_range: Optional[tuple] = None, - allow_quit: bool = False + allow_quit: bool = False, ) -> Optional[Any]: while True: user_input = get_sanitized_input(prompt).strip().lower() @@ -255,7 +269,7 @@ class Selector: columns: Optional[List[str]] = None, allow_multiple: bool = False, prompt_each: bool = False, - header: str = "Available Rows:" + header: str = "Available Rows:", ) -> List[pd.Series]: if df.empty: print("DataFrame is empty.") @@ -272,7 +286,7 @@ class Selector: label_func=label_func, allow_multiple=allow_multiple, prompt_each=prompt_each, - header=header + header=header, ) if isinstance(result, pd.Series): @@ -286,7 +300,7 @@ class Selector: def select_dataframe_with_mode( df: pd.DataFrame, columns: Optional[List[str]] = None, - header: str = "Available Rows:" + header: str = "Available Rows:", ) -> List[pd.Series]: if df.empty: print("āš ļø DataFrame is empty.") @@ -305,7 +319,12 @@ class Selector: print(f"{i}: {label_func(row)}") # Prompt for mode once - print(colorText("\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white")) + print( + colorText( + "\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", + "white", + ) + ) mode = get_sanitized_input("").strip().lower() if mode == "a": @@ -317,7 +336,7 @@ class Selector: label_func=label_func, allow_multiple=True, prompt_each=False, - header=header + header=header, ) if not selected: @@ -331,4 +350,4 @@ class Selector: return [pd.Series(row) for row in items if row not in selected] else: print(colorText("āš ļø Invalid mode. Returning no rows.", "yellow")) - return [] \ No newline at end of file + return [] diff --git a/utils/setup.py b/utils/setup.py index f7ba7f2..4277389 100644 --- a/utils/setup.py +++ b/utils/setup.py @@ -30,16 +30,16 @@ from utils.configmanager import PROTECTED_KEYS, load_protected_config def get_base_directory() -> Path: system = platform.system() home = Path.home() - if system == 'Windows': - return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools" - elif system == 'Darwin': - return home / 'Library' / 'Application Support' / "AirlockTools" + if system == "Windows": + return Path(os.getenv("APPDATA", home / "AppData" / "Roaming")) / "Loxide" + elif system == "Darwin": + return home / "Library" / "Application Support" / "Loxide" else: - return home / '.local' / 'share' / "AirlockTools" + return home / ".local" / "share" / "Loxide" def configure_logging(log_dir: Path, log_level: str = "DEBUG"): - log_file = log_dir / "airlocktools.log" + log_file = log_dir / "Loxide.log" config = { "version": 1, # Required key for dictConfig format version @@ -58,17 +58,17 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): "file": { "class": "logging.handlers.TimedRotatingFileHandler", "filename": str(log_file), - "when": "midnight", # Rotate logs at midnight - "interval": 1, # Every 1 day - "backupCount": 7, # Keep 7 days of logs - "encoding": "utf-8", # Ensure UTF-8 encoding - "level": "DEBUG", # Always log DEBUG and above - "formatter": "detailed", # Use detailed format + "when": "midnight", # Rotate logs at midnight + "interval": 1, # Every 1 day + "backupCount": 7, # Keep 7 days of logs + "encoding": "utf-8", # Ensure UTF-8 encoding + "level": "DEBUG", # Always log DEBUG and above + "formatter": "detailed", # Use detailed format }, "console": { "class": "logging.StreamHandler", "level": log_level.upper(), # Configurable log level - "formatter": "simple", # Use simple format + "formatter": "simple", # Use simple format }, }, "root": { @@ -82,9 +82,9 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): try: config["handlers"]["eventlog"] = { "class": "logging.handlers.NTEventLogHandler", - "appname": "AirlockTools", # Event log source name - "level": "CRITICAL", # Only log critical errors - "formatter": "simple", # Use simple format + "appname": "Loxide", # Event log source name + "level": "CRITICAL", # Only log critical errors + "formatter": "simple", # Use simple format } config["root"]["handlers"].append("eventlog") except Exception as e: @@ -94,9 +94,11 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"): logging.config.dictConfig(config) logging.getLogger().debug("āœ… Logging configured.") - + def get_system_config_path() -> Path: - base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))) + base_path = Path( + getattr(sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__))) + ) return base_path.parent / "system_config.json" @@ -108,45 +110,45 @@ def load_system_config() -> dict: except FileNotFoundError: logging.warning("āš ļø system_config.json not found. Using built-in defaults.") return { - "APPNAME": "AirlockTools", + "APPNAME": "Loxide", "LOG_LEVEL": "DEBUG", "PATH_EXCLUSION_CONST": 4, "MIN_FILES_FOR_PATH": 4, "VT_THREAT_TOLERANCE": 4, - "POLICY_MAP_ENF_AUD": { - "enforced_id": "audit_id" - } + "POLICY_MAP_ENF_AUD": {"enforced_id": "audit_id"}, } + 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" - } + default_user_config = {"URL": "", "LOG_LEVEL": "INFO"} 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) + def write_config_to_env(config: dict, env_path: Path): 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) + serialized = ( + json.dumps(value) if isinstance(value, (list, dict)) else str(value) + ) set_key(env_path, key, serialized) except Exception as e: logging.warning(f"Failed to write {key} to .env: {e}") + def setup(): base_dir = get_base_directory() dirs = { - 'config': base_dir / 'config', - 'cache': base_dir / 'cache', - 'logs': base_dir / 'logs', + "config": base_dir / "config", + "cache": base_dir / "cache", + "logs": base_dir / "logs", } for name, path in dirs.items(): @@ -154,7 +156,7 @@ def setup(): logging.debug(f"{name.capitalize()} directory ensured at: {path}") system_config = load_system_config() - configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG")) + configure_logging(dirs["logs"], system_config.get("LOG_LEVEL", "DEBUG")) env_path = base_dir / ".env" if not env_path.exists(): @@ -171,7 +173,7 @@ def setup(): "Approved": [], "Needs_Review": ["Review_First", "Review_Second", "HTML"], "Preflight": ["HTML"], - "Archived": [] + "Archived": [], } for folder_name, subfolders in folders_structure.items(): @@ -183,7 +185,7 @@ def setup(): subfolder_path.mkdir(parents=True, exist_ok=True) logging.debug(f" ā””ā”€ '{subfolder}' subfolder created at: {subfolder_path}") - user_config = load_user_config(dirs['config']) + user_config = load_user_config(dirs["config"]) merged_config = {**system_config, **user_config} protected_config = load_protected_config() @@ -194,10 +196,12 @@ def setup(): if not url: url = os.getenv("URL") if not url: - url = input("šŸŒ Enter the service URL (e.g., https://example.com/api): ").strip() + 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 logging.debug(f"Service URL set to: {url}") - write_config_to_env(merged_config, env_path) \ No newline at end of file + write_config_to_env(merged_config, env_path) diff --git a/utils/test.py b/utils/test.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/tui.py b/utils/tui.py index 36deeda..5d026b8 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -58,7 +58,9 @@ def _persist_user_theme(theme_name: str) -> None: config_dir.mkdir(parents=True, exist_ok=True) if not user_config_path.exists(): # minimal default like your load_user_config does - user_config_path.write_text('{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8") + user_config_path.write_text( + '{"URL": "", "LOG_LEVEL": "INFO"}\n', encoding="utf-8" + ) # load existing user config user_conf = load_user_config(config_dir) @@ -86,8 +88,6 @@ def _persist_user_theme(theme_name: str) -> None: logger.debug("Reloaded .env from %s", env_path) - - # --------------------------------------------------------------------------- # 1) SCREEN # --------------------------------------------------------------------------- @@ -138,7 +138,6 @@ class MainMenuScreen(Screen): wd = os.getcwd() self.working_dir = wd - def _make_buttons_for(self, tab_id: str) -> Vertical: defs = self.BUTTON_DEFS.get(tab_id, []) buttons = [] @@ -148,9 +147,6 @@ class MainMenuScreen(Screen): buttons.append(btn) return Vertical(*buttons) - - - def compose(self) -> ComposeResult: yield Header(show_clock=True, icon="āš™") @@ -208,8 +204,7 @@ class MainMenuScreen(Screen): new_index = current + direction if 0 <= new_index < len(buttons): buttons[new_index].focus() - - + def switch_tab(self, tab_id: str) -> None: self.current_tab = tab_id content = self.query_one("#content", Vertical) @@ -230,7 +225,9 @@ class MainMenuScreen(Screen): layout.mount(policy_tree) # Right: Details pane - details_pane = Static("Select a policy or device to view details", id="details-pane") + details_pane = Static( + "Select a policy or device to view details", id="details-pane" + ) details_pane.styles.width = "3fr" layout.mount(details_pane) @@ -240,7 +237,9 @@ class MainMenuScreen(Screen): # Top-level policies for _, policy in self.app.policies.iterrows(): if policy["parent"] == "global-policy-settings": - node = policy_tree.root.add(label=policy["name"], data=policy.to_dict()) + node = policy_tree.root.add( + label=policy["name"], data=policy.to_dict() + ) node_map[policy["groupid"]] = node # Child policies @@ -259,7 +258,6 @@ class MainMenuScreen(Screen): label = device["hostname"] # Keep tree clean parent_node.add(label=label, data=device.to_dict()) - elif tab_id == "settings": # Create and mount the horizontal container horizontal_container = Horizontal(id="settings_grid") @@ -279,7 +277,7 @@ class MainMenuScreen(Screen): if j < len(self.THEME_BUTTONS): label, btn_id = self.THEME_BUTTONS[j] button = Button(label, id=f"set_theme_{btn_id}", compact=True) - #button.styles.width = "100%" + # button.styles.width = "100%" column.mount(button) # Mount each button else: @@ -300,9 +298,10 @@ class MainMenuScreen(Screen): details = f"Selected: {node.label}" details_pane.update(details) - - def on_directory_tree_file_selected(self, event: DirectoryTree.FileSelected) -> None: + def on_directory_tree_file_selected( + self, event: DirectoryTree.FileSelected + ) -> None: path = event.path logger.debug("Directory file selected: %s", path) try: @@ -359,14 +358,10 @@ class MainMenuScreen(Screen): self.app.exit() - - - - # --------------------------------------------------------------------------- # 2) APP # --------------------------------------------------------------------------- -class AirlockTools(App): +class Loxide(App): CSS = """ #logo { width: 100%; @@ -407,7 +402,6 @@ class AirlockTools(App): screen.switch_tab("dir") - # --------------------------------------------------------------------------- # 3) TERMINAL + LEGACY # --------------------------------------------------------------------------- @@ -422,6 +416,7 @@ def _restore_terminal_for_legacy() -> None: if os.name == "nt": try: import ctypes + kernel32 = ctypes.windll.kernel32 handle = kernel32.GetStdHandle(-11) mode = ctypes.c_ulong() @@ -447,7 +442,7 @@ def _run_legacy_job(func, args, kwargs) -> None: # --------------------------------------------------------------------------- # 4) PUBLIC ENTRYPOINT # --------------------------------------------------------------------------- -def run_AirlockTools(api: AirlockAPIWrapper) -> None: +def run_Loxide(api: AirlockAPIWrapper) -> None: global _PENDING_JOB while True: @@ -456,7 +451,7 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None: dotenv.load_dotenv(dotenv_path=env_path, override=True) _PENDING_JOB = None - app = AirlockTools(api) + app = Loxide(api) try: app.run() @@ -486,4 +481,4 @@ def run_AirlockTools(api: AirlockAPIWrapper) -> None: # --------------------------------------------------------------------------- if __name__ == "__main__": api = AirlockAPIWrapper() - run_AirlockTools(api) + run_Loxide(api) diff --git a/utils/utils.py b/utils/utils.py index 539a588..26b2023 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -28,8 +28,6 @@ import pandas as pd logger = logging.getLogger(__name__) - - def import_to_dataframe(file_path: str) -> pd.DataFrame: df = pd.DataFrame() @@ -96,17 +94,17 @@ def choose_file(initial_directory=None, required_substring=None): return file_path - - def get_sanitized_input(prompt: str) -> str: while True: user_input = input(prompt) if user_input.strip() == "": return user_input # Allow blank lines - if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()): + if re.match(r"^[a-zA-Z0-9_\- .]+$", user_input.strip()): return user_input else: - print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.") + print( + "Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed." + ) def regulator(paths, case_insensitive=True): @@ -119,8 +117,10 @@ def regulator(paths, case_insensitive=True): pattern = "(?i)" + pattern # Add inline case-insensitive flag print(f"Regulator is providing: {pattern}") return pattern + + def irtang(): - print( + print( colorText( r""" ā–ˆā–ˆā–ˆ @@ -149,6 +149,8 @@ def irtang(): "yellow", ) ) + + def displayIntro(): print( @@ -164,6 +166,8 @@ def displayIntro(): "cyan", ) ) + + def welcome(): print( colorText( @@ -184,11 +188,21 @@ def welcome(): ) ) -def section_header(title): - print(colorText("\n --------------------------------------------------------------------", "cyan")) - print(colorText(f" ------------- {title} -------------", "cyan")) - print(colorText(" --------------------------------------------------------------------", "cyan")) +def section_header(title): + print( + colorText( + "\n --------------------------------------------------------------------", + "cyan", + ) + ) + print(colorText(f" ------------- {title} -------------", "cyan")) + print( + colorText( + " --------------------------------------------------------------------", + "cyan", + ) + ) def areYouSure(): @@ -348,7 +362,11 @@ def printDeviceEnforceChecklist(): "cyan", ) ) - print(colorText(" When complete, save the csv file to the directory 'approved'", "cyan")) + print( + colorText( + " When complete, save the csv file to the directory 'approved'", "cyan" + ) + ) print( colorText( " Do the same process with the list of publishers forthe same directories", @@ -357,20 +375,38 @@ def printDeviceEnforceChecklist(): ) print(colorText(" Preflight Lists will be generated", "cyan")) - print(colorText("5. Choose the destination policy and parent and child allow list", "cyan")) + print( + colorText( + "5. Choose the destination policy and parent and child allow list", "cyan" + ) + ) - print(colorText("6. Test ------------------------------------------------------", "cyan")) + print( + colorText( + "6. Test ------------------------------------------------------", "cyan" + ) + ) print(colorText(" Print rather than apply selected data.", "cyan")) - print(colorText("7. Liftoff ------------------------------------------------------", "cyan")) + print( + colorText( + "7. Liftoff ------------------------------------------------------", "cyan" + ) + ) print( colorText( " Apply path exclusions according to allowed and approved paths", "cyan", ) ) - print(colorText(" Apply signed or attested hashes to Parent Allow List", "cyan")) - print(colorText(" Apply approved, but unsigned hashes to the Child Allow List", "cyan")) + print( + colorText(" Apply signed or attested hashes to Parent Allow List", "cyan") + ) + print( + colorText( + " Apply approved, but unsigned hashes to the Child Allow List", "cyan" + ) + ) print( colorText( @@ -523,10 +559,9 @@ def formatHTML(df, output_html_path=None, overwrite=True): return styled_html - def open_directory(path): system = platform.system() - + if system == "Windows": os.startfile(path) elif system == "Linux": @@ -537,9 +572,9 @@ def open_directory(path): def print_x_wide(items: list, width: int): for i in range(0, len(items), width): - row = items[i:i+width] + row = items[i : i + width] print(" | ".join(row)) - + def clear_screen(): - os.system('cls' if os.name == 'nt' else 'clear') + os.system("cls" if os.name == "nt" else "clear")