diff --git a/.gitignore b/.gitignore index 5145354..7d2a7e0 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,12 @@ *.csv *__pycache__* *.parquet -chunkinator.json \ No newline at end of file +chunkinator.json +jobs.json +*.xl* +*.exe +securitytest.py +*.toml +system_config.json +Development/ +AirlockTools_client*/ \ No newline at end of file diff --git a/AirlockTools.py b/AirlockTools.py deleted file mode 100644 index d86dcb2..0000000 --- a/AirlockTools.py +++ /dev/null @@ -1,257 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import dotenv -import os -import pandas as pd -import urllib3 -import utils.allowlist -import utils.getdeviceevents -import utils.hashfunctions -import utils.pathfunctions -import utils.policyfunctions -import utils.pretty as ct - - -urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) - -dotenv.load_dotenv() - -#Constants -url = os.getenv('url') -bad_publisher_list = ["Brave","Zoom", "GlavSoft", "VNC"] -pups = ["logmein", "invalid", "nmap"] -badpathparts = ["users", "wwwroot", "windows\\temp", "windows\\task", "windows\\system32", "startup", "windows\\fonts", "Recycle.Bin", "AppData", "programdata", "Solarwinds", "kaseya"] -path_exclusion_constant = 4 -min_files_for_path = 4 -threat_tolerance_constant = 4 - -def apivalidation(): - match os.getenv('APIKEY'): - case '': - print(ct.colorText("Please add your API Key to the .env file", "red")) - case _: - menu_main() - -def tryToReadCSV(csv): - try: - df =pd.read_csv(csv) - if df.empty: - print(ct.colorText("Error: CSV file has headers but no data rows.", "red")) - else: - print(ct.colorText(f"Data loaded successfully from {csv}", "green")) - except pd.errors.EmptyDataError: - print(ct.colorText("Notice : CSV file is completely empty (no headers, no data), falling back to empty frame", "white")) - df = pd.DataFrame() # Create an empty DataFrame as fallback - return df - -def tryToReadParquet(parquet): - try: - df = pd.read_parquet(parquet) - if df.empty: - print(ct.colorText("Error: Parquet file has headers but no data rows.", "red")) - else: - print(ct.colorText(f"Data loaded successfully from {parquet}", "green")) - except pd.errors.EmptyDataError: - print(ct.colorText("Notice : Parquet file is completely empty (no headers, no data), falling back to empty frame", "white")) - df = pd.DataFrame() # Create an empty DataFrame as fallback - return df - -def deduplicate_list(lst): - seen = set() - return [x for x in lst if not (x in seen or seen.add(x))] - -def menu_main(): - while True: - ct.displayIntro(); - print(ct.colorText("1. Get All Events for Single Device", "yellow")) - print(ct.colorText("2. Placeholder for Local Approval", "yellow")) - print(ct.colorText("3. Placeholder for Another Tool", "yellow")) - print(ct.colorText("4. Prepare Policy For Enforcement", "yellow")) - print(ct.colorText("Q. Quit", "yellow")) - - choice = input(ct.colorText("\nEnter Menu Item: ", "white")) - if choice == '1': - utils.getdeviceevents.devicehistory(url,False) - elif choice == "2": - menu_local_approve() - elif choice == "3": - menu_feature2() - elif choice == "4": - menu_prepare_to_enforce() - elif choice == "Q": - break - else: - print(ct.colorText("Invalid choice. Please try again.","red")) - -def menu_local_approve(): - while True: - print("\n--- Submenu ---") - print("1. Sub-option A") - print("2. Sub-option B") - print("3. Return to Main Menu") - choice = input("Enter your choice: ") - - if choice == "1": - print("You selected Sub-option A") - elif choice == "2": - print("You selected Sub-option B") - elif choice == "3": - print("Returning to Main Menu...") - break - else: - print("Invalid choice. Please try again.") - -def menu_feature2(): - while True: - print("\n--- Submenu ---") - print("1. Sub-option A") - print("2. Sub-option B") - print("3. Return to Main Menu") - choice = input("Enter your choice: ") - - if choice == "1": - print("You selected Sub-option A") - elif choice == "2": - print("You selected Sub-option B") - elif choice == "3": - print("Returning to Main Menu...") - break - else: - print("Invalid choice. Please try again.") - -def menu_prepare_to_enforce(): - - first_policy = " " - second_policy = " " - destination_name = " " - destination_id = " " - allowlist_parent_name = " " - allowlist_parent_id = " " - allowlist_child_name = " " - allowlist_child_id = " " - - #If the directorys where we're going to store our output dont exist, make them. - if not os.path.exists("parquet"): os.makedirs("parquet") - if not os.path.exists("needs_approved"): os.makedirs("needs_approved") - if not os.path.exists("approved"): os.makedirs("approved") - if not os.path.exists("preflight"): os.makedirs("preflight") - - while True: - - ct.printEnforceChecklist(first_policy, second_policy, allowlist_child_name, allowlist_parent_name, destination_name) - - choice = input(ct.colorText("\nEnter your choice: ", "white")) - - if choice == "1": - - choice, policynames, policyid = utils.allowlist.listPolicies(url) - first_policy = policynames[choice] - while True: - answer = input(ct.colorText(f"{"Do you want to load a second policy?"} (yes/no): ", "white").strip().lower()) - if answer in ("yes", "y"): - choice, policynames, policyid = utils.allowlist.listPolicies(url) - second_policy = policynames[choice] - - break - elif answer in ("no", "n"): - second_policy = first_policy - break - else: - print(ct.colorText("Please answer with 'yes' or 'no'.", "red")) - - elif choice == "2": - - if not os.path.exists(f"parquet\\execution_history_{first_policy}.parquet"): - utils.policyfunctions.getPolicyInfo(url, first_policy, 60) - - if not os.path.exists(f"parquet\\execution_history_{second_policy}.parquet"): - utils.policyfunctions.getPolicyInfo(url, second_policy, 60) - - if not os.path.exists(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"): - utils.hashfunctions.combineHashes(url, first_policy, second_policy) - - if not os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"): - utils.hashfunctions.categorizeHashes( - first_policy, - second_policy, - pd.read_parquet(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"), - threat_tolerance_constant, - bad_publisher_list, - pups - ) - - if not os.path.exists(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet"): - utils.hashfunctions.condenseExecutions(first_policy,second_policy) - - if os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") & os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") & os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"): - utils.hashfunctions.divideSortedHashExecutions(first_policy,second_policy,pups) - - elif choice == "3": - - if os.path.exists(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv") and os.path.exists(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv"): - utils.pathfunctions.generatePathReview(first_policy, second_policy, badpathparts, min_files_for_path) - else: - print(ct.colorText(f"Please manually approve hashes prior to this step","red")) - - elif choice == "4": - - if os.path.exists(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv"): - if not os.path.exists(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet") and not os.path.exists(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet"): - utils.hashfunctions.generatePreflights(first_policy, second_policy) - - elif choice == "5": - - print(ct.colorText(f"Please choose destination_name Policy for Path Exclusions","white")) - choice, policynames, policyid = utils.allowlist.listPolicies(url) - #print(allowlist_parent_tuple) - destination_name = policynames[choice] - destination_id = policyid[choice] - - print(ct.colorText(f"Please choose Parent Allowlist for Known Hashes","white")) - choice, allowlists,allowid = utils.allowlist.listAllowlists(url) - #print(allowlist_parent_tuple) - allowlist_parent_name = allowlists[choice] - allowlist_parent_id = allowid[choice] - - print(ct.colorText(f"Please choose Child Allowlist for Less-Known Hashes","white")) - choice, allowlists, allowid = utils.allowlist.listAllowlists(url) - #print(allowlist_child_tuple) - allowlist_child_name = allowlists[choice] - allowlist_child_id = allowid[choice] - - elif choice == "6": - if os.path.exists(f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html") and os.path.exists(f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html") and allowlist_parent_name != " " and allowlist_child_name != " " and destination_name != " ": - utils.policyfunctions.sendToPolicy( - url, - first_policy, - second_policy, - destination_name, - destination_id, - allowlist_parent_name, - allowlist_parent_id, - allowlist_child_name, - allowlist_child_id - ) - - elif choice == "Q": - break - else: - print(ct.colorText("Invalid choice. Please try again.", "red")) - - -if __name__ == "__main__": - apivalidation() - diff --git a/AirlockTools_Client.py b/AirlockTools_Client.py new file mode 100644 index 0000000..0554d43 --- /dev/null +++ b/AirlockTools_Client.py @@ -0,0 +1,78 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + +#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 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.utils import irtang + +urllib3.disable_warnings( + urllib3.exceptions.InsecureRequestWarning +) + +def main(): + irtang() + #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") + + 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}") + + except ValueError as e: + logger.error(f"Configuration error: {e}", exc_info=True) + raise + + + api_key = getAPI(username, "AirlockTools") + if api_key is None: + raise ValueError("API key for AirlockTools is missing.") + + api = AirlockAPIWrapper( + base_url=str(os.getenv("URL")), + api_key=api_key, + ) + run_AirlockTools(api) + + + +if __name__ == "__main__": + main() diff --git a/AirlockTools_Server.py b/AirlockTools_Server.py new file mode 100644 index 0000000..56a58a0 --- /dev/null +++ b/AirlockTools_Server.py @@ -0,0 +1,90 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + + +#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 +import os + +import dotenv +import urllib3 + +import flows.localApproval as la +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 +) + +def main(): + + #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") + + 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}") + + except ValueError as e: + logger.error(f"Configuration error: {e}", exc_info=True) + raise + + api = AirlockAPIWrapper( + base_url=str(os.getenv("URL")), + api_key = getAPI(username, "AirlockTools"), + ) + + + 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"): + recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api]) + recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api]) + else: + reload_jobs() + + start_scheduler() + +if __name__ == "__main__": + main() diff --git a/IRT_icon_32-512.ico b/IRT_icon_32-512.ico new file mode 100644 index 0000000..103b72e Binary files /dev/null and b/IRT_icon_32-512.ico differ diff --git a/README.md b/README.md index 86edfdb..626650e 100644 --- a/README.md +++ b/README.md @@ -1,17 +1,59 @@ -[![AGPL License](https://img.shields.io/badge/license-AGPL-blue.svg)](http://www.gnu.org/licenses/agpl-3.0) -# Airlock Digital Local Approval +# šŸ›”ļø Airlock Tools -Python based Carbon Black App Control feature implementation for Airlock -## Features +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. -- "Local Approval Initialization" -This programmatically scans devices in audit mode within Airlock and subsequently adds the identified blocks to a user-specified whitelist. +--- + + +## šŸš€ Features +- šŸ” **Fuzzy Device Search** + Quickly locate devices using partial or approximate matches. + +- šŸ“¦ **Batch Move Devices** + Move multiple devices between groups or policies easily. + +- šŸ”„ **Toggle Enforcement/Audit Policies** + Seamlessly switch devices between enforcement and audit modes. + +- šŸ•µļøā€ā™‚ļø **Device History Search** + Track agent executions. + +- šŸ§° **Prepare Policies for Enforcement** + Validate and stage policies before pushing them to enforcement. + +- šŸ’¤ **Find Quiet Hosts** + Identify devices ready for enforcement. + +- šŸŽ›ļø **TUI** + Navigate with arrow keys and F-key shortcuts using a custom ANSI-colored terminal UI. + +--- + +## šŸ§­ Roadmap + +- āš™ļø **Rust-based Async API Calls** + Improve performance and concurrency with a Rust-powered backend. + +- āœ… **Carbon Black-style Local Approval** + Enable local user approvals for policy exceptions and enforcement actions. + +- šŸ“Š **Audit Logging & Export** + Add detailed logging and export capabilities for compliance and analysis. + +--- + +## šŸ§‘ā€šŸ’» Requirements + +TBD + +--- + +## šŸ“œ License -## License **AirlockTools** 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. -See the [LICENSE](LICENSE.md) file for full details, or visit -[https://www.gnu.org/license/agpl-3.0.html](https://www.gnu.org/license/agpl-3.0.html) \ No newline at end of file +See the [LICENSE](LICENSE.md) file for full details, or visit +[https://www.gnu.org/license/agpl-3.0.html](https://www.gnu.org/license/agpl-3.0.html) diff --git a/Server/scheduler_async.py b/Server/scheduler_async.py new file mode 100644 index 0000000..6f90d36 --- /dev/null +++ b/Server/scheduler_async.py @@ -0,0 +1,192 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import asyncio +import json +import logging +import os +from typing import Any, Callable, Dict, List + +logger = logging.getLogger(__name__) + +# Registry of functions that can be scheduled +FUNCTION_MAP: Dict[str, Callable] = {} + +# Dictionary to manually track scheduled jobs by ID +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. + Example: + register_function("say_hello", say_hello) + """ + FUNCTION_MAP[name] = func + +def load_jobs() -> List[Dict[str, Any]]: + """ + Load jobs from the JSON file, or return [] if none exist. + """ + if not os.path.exists(JOBS_FILE): + return [] + with open(JOBS_FILE, "r") as f: + return json.load(f) + +def save_jobs(jobs: List[Dict[str, Any]]): + """ + Save jobs to the JSON file (overwrite). + """ + 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. + """ + handle = scheduled_jobs.pop(job_id, None) + if handle: + handle.cancel() + logger.info(f"Cancelled job '{job_id}'") + + 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): + """ + Schedule a job to run once after a delay (in seconds). + """ + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + func = FUNCTION_MAP.get(func_name) + if func is None: + logger.error(f"Function '{func_name}' is not registered.") + return + func(*args, **kwargs) + cancel_job(job_id) + + loop = asyncio.get_event_loop() + handle = loop.call_later(delay_seconds, job_wrapper) + scheduled_jobs[job_id] = handle + + 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 + }) + save_jobs(jobs) + 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): + """ + Schedule a recurring job. + """ + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + func = FUNCTION_MAP.get(func_name) + if func is None: + logger.error(f"Function '{func_name}' is not registered.") + return + func(*args, **kwargs) + # Reschedule the job + handle = asyncio.get_event_loop().call_later(interval, job_wrapper) + scheduled_jobs[job_id] = handle + + cancel_job(job_id) + handle = asyncio.get_event_loop().call_later(interval, job_wrapper) + scheduled_jobs[job_id] = handle + + 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 + }) + save_jobs(jobs) + logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") + +def reload_jobs(): + """ + Reload jobs from JSON and reschedule them. + """ + jobs = load_jobs() + for job in jobs: + if job["type"] == "once": + run_once_job( + job["id"], + job["function"], + job["delay"], + job.get("args"), + job.get("kwargs"), + persist=False + ) + elif job["type"] == "recurring": + recurring_job( + job["id"], + job["function"], + job["interval"], + job.get("args"), + job.get("kwargs"), + persist=False + ) + +async def start_scheduler(): + """ + Start the asynchronous scheduler loop. + + This function is a placeholder to keep the event loop alive. + Jobs are scheduled using asyncio.call_later and do not require polling. + """ + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + logger.critical("Scheduler stopped.") + + """ + Start the asynchronous scheduler loop. + + This function is a placeholder for compatibility. Since we use asyncio.call_later, + jobs are scheduled directly on the event loop and no polling is required. + + Usage: + # In an async app (e.g., Textual) + asyncio.create_task(start_scheduler()) + + # Or in a standalone script + async def main(): + await start_scheduler() + + asyncio.run(main()) + """ + try: + 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 diff --git a/allowlistandgroup.py b/allowlistandgroup.py deleted file mode 100644 index 5fcf86f..0000000 --- a/allowlistandgroup.py +++ /dev/null @@ -1,22 +0,0 @@ -import requests -import dotenv -import json -import os -import utils.pretty as ct - -url = 'https://172.17.22.240:3129' -policiesnames = [] -policyids=[] -dotenv.load_dotenv() -endpoint = url + '/v1/application' -print(ct.colorText("[+] Grabbing All Allowlists", "cyan")) -payload = {} -headers = { - "X-APIKey": os.getenv('APIKEY') -} -response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) -parse_text = json.loads(response.text) -for index, list in enumerate(parse_text['response']['applications'], start=1): - if index >= 38: - print(list) -#Need else and catch for upper bound diff --git a/default_system_config.json b/default_system_config.json new file mode 100644 index 0000000..b09f901 --- /dev/null +++ b/default_system_config.json @@ -0,0 +1,14 @@ +{ + "APPNAME": "AirlockTools", + "URL": "https://server:3129", + "LOG_LEVEL": "INFO", + "BAD_PATH_PARTS": ["users","wwwroot","windows\\temp","windows\\task","windows\\system32","startup", "windows\\fonts","Recycle.Bin","AppData","programdata", "Solarwinds","kaseya"], + "BAD_PUBLISHERS": ["Brave", "Zoom", "GlavSoft", "VNC"], + "PUPS":["logmein","invalid","nmap","LTSvc","VNC","Kaseya","Solarwinds","mRemoteNG"], + "PATH_EXCLUSION_CONST": 4, + "MIN_FILES_FOR_PATH": 4, + "VT_THREAT_TOLERANCE": 4, + "POLICY_MAP_ENF_AUD": { + + } +} \ No newline at end of file diff --git a/flows/localApproval.py b/flows/localApproval.py new file mode 100644 index 0000000..789b917 --- /dev/null +++ b/flows/localApproval.py @@ -0,0 +1,291 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + +import datetime +import logging +import os +import re +import time + +import dotenv +import numpy as np +import pandas as pd + +from models.agent import Agent +from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents +from services.API import AirlockAPIWrapper +from utils.configmanager import get_protected_json, load_env, load_env_json +from utils.setup import get_base_directory +from utils.utils import colorText, get_sanitized_input + +logger = logging.getLogger(__name__) + +dotenv.load_dotenv() + + +def getLocalApprovals(api: AirlockAPIWrapper): + base_dir = get_base_directory + 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.to_parquet( + f"{base_dir}\\cache\\last_local_approval.parquet", index=False + ) + os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet") + + # Only keep rows presumably created by the generate local approval function + local_approval = local_approval[ + local_approval["purpose"].str.startswith("šŸŽ« Local Approval šŸŽ«") + ] + + local_approval["batchid"] = local_approval["purpose"].apply( + lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1) + ) + + if not local_approval.empty: + local_approval.to_parquet( + f"{base_dir}\\cache\\newest_local_approval.parquet", index=False + ) + + return local_approval + + +def scheduleAddingLAHashes(api: AirlockAPIWrapper): + + 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) + + try: + register_function("add_hash", returnFromLocalApproval) + register_function("move_device", moveAgentToRelatedPolicy) + except Exception as e: + logger.warning(f"Failed to register functions: {e}") + return + + try: + approvals_df = getNewLocalApprovals(api) + if approvals_df.empty: + logger.debug("No new local approvals found. Nothing to schedule.") + return + batches = approvals_df.groupby("batchid") + except Exception as e: + logger.warning(f"Failed to retrieve or group local approvals: {e}") + return + + for batchid, batch_df in batches: + try: + 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_timestamp = early_time.timestamp() + run_timestamp = run_time.timestamp() + + # Schedule add_hash job + try: + run_once_job( + f"add_hash_{batchid}", + "add_hash", + early_timestamp, + [ + api, + batch_df, + policy_relationship_map, + bad_publisher_list, + pups, + threat_tolerance_constant, + ], + None, + ) + logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}") + except Exception: + logger.debug("Failed to schedule add_hash for batch {batchid}: {e}") + + # Schedule move_device jobs + devices = batch_df["agentid"].drop_duplicates().tolist() + agents = [] + + for device in devices: + rows = api.agent_find_by_hostname(device).iterrows() + agents += [Agent(**row["data"]) for _, row in rows] + + for agent in agents: + try: + run_once_job( + f"move_device_{agent.hostame}_{batchid}", + "move_device", + run_timestamp, + [api, agent, policy_relationship_map], + "enforcement", + ) + + print( + f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}" + ) + except Exception as e: + print( + f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}" + ) + + except Exception as e: + logger.warning(f"Failed to process batch {batchid}: {e}") + + +def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant +): + """ + # Get unique policy names from device list + policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist()) + + # Create inverse map to go from Audit to Enforcement + inverse_map = {v: k for k, v in policy_relationship_map.items()} + + # Fetch all policies + all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()] + + # Define policy types + policy_types = [1, 2, 6, 7] + + #TODO finish logic for adding hashes + """ + working_dir = load_env("WORKING_DIR") + 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}") + +def moveToLocalApproval(api: AirlockAPIWrapper): + possible_durations = [15, 60, 360, 1440, 10080] + duration_selected = None + + print(colorText("Please select a duration:", "white")) + for i, option in enumerate(possible_durations, start=1): + print(f"{i}. {option}") + + try: + + choice = int(get_sanitized_input("Enter the number of your choice:")) + if 1 <= choice <= len(possible_durations): + duration_selected = possible_durations[choice - 1] + print(colorText(f"You selected: {duration_selected}", "yellow")) + logger.debug(f"You selected: {duration_selected}") + else: + print(colorText("āŒ Invalid choice.", "red")) + logger.debug("Invalid Input") + return + except ValueError: + print(colorText("āŒ Invalid input. Please enter a number.", "red")) + logger.debug("Invalid Input") + return + + agents = selectAgents(api) + batch = int(time.time()) + + if not agents: + print(colorText("āŒ No agents found or error retrieving agents.", "red")) + logger.debug("No agents found or error retrieving agents") + return + + for agent in agents: + try: + addLocalApproval(api, batch, duration_selected, agent.agentid) + moveAgentToRelatedPolicy(api, agent, "audit") + except Exception as e: + print(colorText(f"āŒ Error processing agent {agent.hostname}: {e}", "red")) + + +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): + current_agents = findAllAgents(api) + last_agents = [] + if not last_agents: + last_agents = current_agents + 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()) + + # Index last_agents by hostname for quick lookup + last_agent_map = {agent.hostname: agent for agent in last_agents} + + # Result buckets + newly_added = [] + same_policy = [] + moved_to_audit = [] + moved_to_enforcement = [] + unusual_move = [] + + for current in current_agents: + previous = last_agent_map.get(current.hostname) + + if not previous: + newly_added.append(current) + continue + + if current.groupid == previous.groupid: + same_policy.append(current) + elif (previous.groupid, current.groupid) in known_transitions: + moved_to_audit.append(current) + elif (current.groupid, previous.groupid) in known_transitions: + moved_to_enforcement.append(current) + else: + unusual_move.append(current) + + # Return all five DataFrames + return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move + + +def getNewLocalApprovals(api: AirlockAPIWrapper): + + working_dir = load_env("WORKING_DIR") + current_la = getLocalApprovals(api) + + # Load old approval list + old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet" + if os.path.exists(old_la_path): + old_la = pd.read_parquet(old_la_path) + else: + old_la = pd.DataFrame(columns=current_la.columns) + + # Create composite keys + 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) + recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago] + + # Save current approvals for next run + current_la.drop(columns=["key"], inplace=True) + current_la.to_parquet(old_la_path, index=False) + + return recent_entries diff --git a/flows/otp.py b/flows/otp.py new file mode 100644 index 0000000..6c09db2 --- /dev/null +++ b/flows/otp.py @@ -0,0 +1,170 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + + +from datetime import datetime +import logging +import os + +import pandas as pd + +from services.agenthandler import selectAgents +from services.API import AirlockAPIWrapper +from utils.configmanager import load_env +from utils.selector import Selector +from utils.utils import colorText, get_sanitized_input + +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")) + 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")) + duration_selected = Selector.select_int(possible_durations) + + if isinstance(duration_selected, list): + duration_selected = duration_selected[0] if duration_selected else None + + if duration_selected is not None: + for agent in agents: + logging.info(f"Querying API for {agent.hostname}") + otp_code = api.otp_generate(agent.agentid, duration_selected, purpose) + logger.debug(f"Generated OTP for {agent.hostname}: {otp_code}") + otp_dict[agent.hostname] = otp_code + + print(colorText("Requested Codes:", "green")) + for key, value in otp_dict.items(): + 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' + + # 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) + + #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)] + + #Present and select rows + selected_rows = Selector.select_dataframe_with_mode( + combined_agents, + columns=['otpid', 'hostname', 'status','purpose','granted'], + header="OTP Sessions" + ) + combined_df = pd.DataFrame() + + for row in selected_rows: + otpid = row['otpid'] + hostname = row['hostname'] + result = api.otp_get_activities(otpid) + 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': + working_dir = load_env("WORKING_DIR") + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + filename = f"otp_activities_{timestamp}.csv" + file_path = os.path.join(str(working_dir), filename) + + combined_df.to_csv(file_path, index=False) + logging.info(f"Exported Data to {file_path}") + + print( + colorText( + f"\nāœ… OTP Activity exported to: {working_dir}\\{filename}", + "green", + ) + ) + else: + 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' + + combined_agents = pd.concat([activeagents, awaitingagents], ignore_index=True) + combined_agents = combined_agents.sort_values(by='otpid', ascending=False) + + # Combine all into one DataFrame + combined_agents = pd.concat([activeagents, awaitingagents], 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': + agentnames = [] + agents = selectAgents(api) + for agent in agents: + agentnames.append(agent.hostname) + + combined_agents = combined_agents[combined_agents['hostname'].isin(agentnames)] + + #Present and select rows + selected_rows = Selector.select_dataframe_with_mode( + combined_agents, + columns=['otpid', 'hostname', 'status','purpose','granted'], + header="OTP Sessions" + ) + + for row in selected_rows: + 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 new file mode 100644 index 0000000..f66eee7 --- /dev/null +++ b/flows/prepPolicy.py @@ -0,0 +1,627 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import logging +import os +import os.path +import re +from typing import List + +import dotenv +import pandas as pd + +from models.execution import ExecutionHistoryRecord +from models.policy import Allowlist, Policy +from services.API import AirlockAPIWrapper +from utils.configmanager import get_protected_value, load_env, load_env_json +from utils.selector import Selector +from utils.utils import ( + areYouSure, + clear_screen, + colorText, + formatHTML, + get_sanitized_input, + locked, + open_directory, + print_x_wide, + regulator, +) + +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()] + logger.debug("Prompting for Policies") + print(colorText("Please select policy/policies", "white")) + selected = Selector.select_objects(policies, allow_multiple, prompt_each=True) + + if selected is None: + return [] + + # Normalize to always return a list + logger.debug("Returning {selected.dict}") + 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()] + logger.debug("Prompting for Allowlist(s)") + print(colorText("Please select allowlist(s)", "white")) + selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True) + + if selected is None: + return [] + + # Normalize to always return a list + logger.debug(f"Returning {selected}") + return selected if isinstance(selected, list) else [selected] + + +def sortHashes( + api: AirlockAPIWrapper, + selected_policies: List[Policy], + type=[1, 2, 6, 7] +): + working_dir = load_env("WORKING_DIR") + history_days = Selector.select_value( + prompt="Enter how many days of history to pull (1ā€“150): ", + 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) + + categories = { + "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 + + 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" + + # Convert ExecutionHistoryRecord objects to dictionaries + df = pd.DataFrame([r.__dict__ for r in records]) + + # Optional: flatten hash_obj if needed + if not df.empty and 'hash_obj' in df.columns: + hash_df = df['hash_obj'].apply(lambda h: h.to_dict() if h else {}) + df = pd.concat([df.drop(columns=['hash_obj']), hash_df], axis=1) + + # Save to CSV + df.to_csv(csv_path, index=False) + logger.info(f"Saved {label} executions to {csv_path}") + + # Generate HTML + formatHTML(df, html_path) + logger.info(f"Generated HTML report at {html_path}") + + +def buildPathsandPublishers(selected_policies: List[Policy], split): + working_dir = load_env("WORKING_DIR") + df1 = pd.DataFrame() + df2 = pd.DataFrame() + all_approved_hashes = pd.DataFrame() + 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) + + if os.path.exists(path1): + df1 = pd.read_csv(path1) + else: + logger.warning(f"File not found: {path1}") + + if os.path.exists(path2): + df2 = pd.read_csv(path2) + else: + logger.warning(f"File not found: {path2}") + + if df1.empty and df2.empty: + logger.warning("Both DataFrames are empty. Skipping sort.") + all_approved_hashes = pd.DataFrame() + logger.debug(all_approved_hashes.head) + else: + all_approved_hashes = pd.concat([df1, df2], ignore_index=True) + 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.") + + if not all_approved_hashes.empty and path_exclusion_constant: + + primary_path_exclusions = calculatePath( + 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 = remaining_hashes[ + ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) + ] + dataframes = { + "all_approved_hashes" : all_approved_hashes, + "primary_Paths": primary_path_exclusions, + "secondary_Paths": secondary_path_exclusions, + "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 not all_approved_hashes.empty: + # Drop all not signed, only keep unique values + publist = all_approved_hashes[ + 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","[]")) + 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: + logger.debug("Approved Hashes list appears empty") + +def buildPreflights(selected_policies: List[Policy]): + working_dir = load_env("WORKING_DIR") + + df1 = pd.DataFrame() + df2 = pd.DataFrame() + approved_hashes = pd.DataFrame() + approved_publishers = pd.DataFrame() + + hash = f"{working_dir}\\Needs_Review\\Review_Second\\{selected_policies[0].name}_all_approved_hashes.csv" + path1 = f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv" + 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 + if os.path.exists(path1): + df1 = pd.read_csv(path1) + else: + logger.warning(f"File not found: {path1}") + + if os.path.exists(path2): + df2 = pd.read_csv(path2) + else: + logger.warning(f"File not found: {path2}") + + if df1.empty and df2.empty: + logger.warning("Both DataFrames are empty. Skipping sort.") + 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. + if os.path.exists(hash): + hashes = pd.read_csv(hash) + approved_hashes = hashes[~hashes['filename'].isin(approved_paths['longestcfp'])] + + approved_hashes = approved_hashes.drop_duplicates(subset="sha256", keep ="first") + + else: + logger.warning(f"File not found: {hash}") + + + if os.path.exists(publishers): + approved_publishers = pd.read_csv(publishers) + + else: + logger.warning(f"File not found: {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") + +def splitFilepathsGrouped(df, path_exclusion_constant, col="filename"): + 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)): + return [] + parts = str(os.path.normpath(path)).split(os.sep) + parts = [p for p in parts if p] # Remove empty strings + return parts + + # Diagnostic: log any non-string entries + 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) + + df = df.copy() + 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() + split_paths = split_paths[df.index] + + df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant])) + grouped = df.groupby("group_key") + new_rows = [] + + for _, group_df in grouped: + paths = group_df[col].tolist() + split_parts = [clean_split(p) for p in paths] + + def longest_common_prefix(paths): + if not paths: + return [] + prefix = paths[0] + for path in paths[1:]: + prefix = [a for a, b in zip(prefix, path) if a == b] + if not prefix: + break + return prefix + + common_prefix = longest_common_prefix(split_parts) + prefix_str = os.sep.join(common_prefix) + + for i, parts in enumerate(split_parts): + filename = parts[-1] + middle = ( + os.sep.join(parts[len(common_prefix):-1]) + if len(parts) > len(common_prefix) + 1 + else "" + ) + row = group_df.iloc[i].copy() + row["longestcfp"] = prefix_str + row["middle"] = middle + row["filename_only"] = filename + row["file_extension"] = os.path.splitext(filename)[1].lower() + new_rows.append(row) + + 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")] + else: + 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) + + processed_dfs = [] + + for df in dfs_by_policy: + haslcp = splitFilepathsGrouped(df, path_exclusion_constant, "filename") + haslcp = haslcp.drop_duplicates() + + forbidden = regulator(badpathparts, True) + forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False) + + logger.debug("Removing forbidden filepaths for path exceptions") + print(colorText("Removing forbidden filepaths for path exceptions", "green")) + lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() + + lcp_not_forbidden_review = lcp_not_forbidden[ + [ + "policyname", + "longestcfp", + "middle", + "filename_only", + "file_extension", + "sha256", + ] + ] + + unique_sha_counts = ( + lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index() + ) + unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"] + + lcp_not_forbidden_review = lcp_not_forbidden_review.merge( + unique_sha_counts, on="longestcfp", how="left" + ) + lcp_not_forbidden_review = lcp_not_forbidden_review[ + lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path + ] + processed_dfs.append(lcp_not_forbidden_review) + + pathExclusions = pd.concat(processed_dfs, ignore_index=True) + + return pathExclusions + +def testChange(selected_policies, destination_policy, destination_allowlist): + working_dir = load_env("WORKING_DIR") + + 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") + + 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) + ] + + 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 hashes would be added to:", "yellow")) + print(destination_allowlist) + + 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 + selected_policies = [] + destination_policy = [] + destination_allowlist = [] + processed_paths = [] + processed_hashes = [] + processed_publishers = [] + working_dir = load_env("WORKING_DIR") + + while True: + printEnforceChecklist(selected_policies, destination_policy, destination_allowlist) + choice = get_sanitized_input("\nEnter your choice: ") + + if choice == "1": + clear_screen() + selected_policies = selectPolicies(api,True) + + elif choice == "2": + clear_screen() + print(colorText("Please choose destination_name Policy for Path Exclusions", "white")) + + destination_policy = selectPolicies(api, False) + + print(colorText("Please choose Allowlist for Hashes", "white")) + + destination_allowlist = selectAllowlists(api, destination_policy, False) + + elif choice == "3": + clear_screen() + sortHashes( + api, + selected_policies, + type=[1, 2, 6, 7], + ) + + elif choice == "4": + clear_screen() + 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.") + + elif choice == "5": + clear_screen() + 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.") + + 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") + and destination_policy + and 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"): + missing_items.append("approved_paths.csv not found") + 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") + if not destination_allowlist: + missing_items.append("destination_allowlist is empty or None") + + logger.error("Preflight check failed due to the following:") + for item in missing_items: + logger.error(f" - {item}") + + elif choice == "7": + clear_screen() + areYouSure() + confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") + if ( + processed_paths + and processed_hashes + and processed_publishers + and destination_policy + and destination_allowlist + 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) + if 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: + logger.error(" - Test not performed.") + if not destination_policy: + logger.error(" - `destination_policy` is missing or invalid.") + 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()) + + 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(f" ------------- {title} -------------", "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")) + + # Step 1: Originating Policies + 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: + print(colorText("The following policies have been chosen:", "green")) + for policy in selected_policies: + print(colorText(f" [āœ“] {policy.name}", "green")) + + # Step 2: Destination Policy and Allowlist + 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")) + 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")) + 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")) + 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")) + else: + 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")) + + 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")) + else: + 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")) + + 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")) + else: + 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("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 new file mode 100644 index 0000000..6bc05ea --- /dev/null +++ b/flows/quietAgent.py @@ -0,0 +1,129 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import datetime +import logging + +import dotenv +import pandas as pd + +from flows.prepPolicy import selectPolicies +from services.API import AirlockAPIWrapper +from services.policyhandler import getPolicyInfo +from utils.configmanager import load_env +from utils.selector import Selector +from utils.utils import colorText, get_sanitized_input + +logger = logging.getLogger(__name__) + + +dotenv.load_dotenv() + + +def findQuietAgents(api: AirlockAPIWrapper): + working_dir = load_env("WORKING_DIR") + # Get policy selection and agent list + selected_policy = selectPolicies(api, False) + if selected_policy: + agents = api.agents_find_by_group(selected_policy[0].groupid) + + # Prompt user for history range + history_days = Selector.select_value( + prompt="Enter how many days of history to pull (1ā€“150): ", + value_type=int, + valid_range=(1, 150), + ) + required_quiet = Selector.select_value( + prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1ā€“365): ", + value_type=int, + valid_range=(1, 150), + ) + + 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 + ) + + + if policy_exec_history.empty: + 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 + ) + + # Get current UTC time + now = datetime.datetime.now(datetime.timezone.utc) + + # Calculate days ago + policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply( + lambda dt: (now - dt).days + ) + + # Count total executions per hostname + 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) + + # Find most recent execution per hostname + 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( + most_recent_exec.set_index("hostname")["days_ago"] + ) + + # Check for enforcement readiness + agents["required_quiet"] = required_quiet + agents["enforce_ready"] = agents["days_since"].apply( + lambda x: True if pd.isna(x) or x > required_quiet else False + ) + + # Sort agents by execution count and hostname + 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" + logging.debug(f"Saving CSV to {filename}") + print(colorText(f"Saving CSV to {filename}", "green")) + agents.to_csv(filename, index=False) + + # Summary statistics + total_agents = len(agents) + ready_agents = agents["enforce_ready"].sum() + not_ready_agents = total_agents - ready_agents + 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" + f"Agents not ready: {not_ready_agents}\n" + f"Percentage ready for enforcement: {ready_percentage:.2f}%" + ) + logger.debug(message) + colorText(message,"green") + get_sanitized_input("Press enter to continue") diff --git a/models/agent.py b/models/agent.py new file mode 100644 index 0000000..9a4c580 --- /dev/null +++ b/models/agent.py @@ -0,0 +1,75 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +from dataclasses import dataclass, field +from typing import ClassVar, List, Optional + +from models.policy import Policy + + +@dataclass +class Agent: + agentid: str + clientversion: str + domain: str + freespace: int + groupid: str # Changed to str to match UUID-style IDs + hostname: str + ip: str + localip: str + lastcheckin: str + os: str + policyversion: str + status: int # raw status code + username: str + groupname: Optional[str] = field(default=None) + status_text: Optional[str] = field(default=None) + + # Class-level status map + 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" +""" + +from models.agent import Agent +from modesls.policy + +# Step 1: Load data from API +policies = [Policy(**row['data']) for _, row in api.policy_find_all().iterrows()] +agents = [Agent(**row['data']) for _, row in api.agent_find_all().iterrows()] + +# Step 2: Create groupid ā†’ groupname map +groupid_to_name = {policy.groupid: policy.name for policy in policies} + +# Step 3: Enrich agents +for agent in agents: + agent.enrich_with_policies(groupid_to_name) + + +""" diff --git a/models/execution.py b/models/execution.py new file mode 100644 index 0000000..c5375df --- /dev/null +++ b/models/execution.py @@ -0,0 +1,472 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import dataclasses +from dataclasses import asdict, dataclass +from datetime import datetime +import inspect +import json +import logging +import os +import re +from typing import List, Optional, Tuple + +import dotenv +import pandas as pd + +from services.API import AirlockAPIWrapper +from services.policyhandler import pullPolicyExechistories +from utils.configmanager import get_protected_value, load_env_json +from utils.utils import colorText, regulator + +logger = logging.getLogger(__name__) + +dotenv.load_dotenv() + +@dataclass +class Hash: + """ + Hash model representing Hash data + """ + sha256: str + applications: str + baselines: str + blocklists: str + createtime: str + datetime: str + description: str + filename: str + filepath: str + filesize: str + md5: str + modtime: str + origname: str + productname: str + productversion: str + publisher: str + reputation: str + sha128: str + sha384: str + sha512: str + at_decision: Optional[str] = None + + def to_dict(self): + return asdict(self) + + @classmethod + def from_dict(cls, data: dict): + return cls(**data) + + @classmethod + def deduplicate(cls, hash_list): + """ + Deduplicates a list of Hash objects based on sha256. + Args: + hash_list (list): List of Hash instances. + Returns: + list: Deduplicated list of Hash instances. + """ + seen = set() + deduped = [] + for h in hash_list: + if h.sha256 not in seen: + seen.add(h.sha256) + deduped.append(h) + return deduped + + @classmethod + def categorize_hashes(cls, hashes): + threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int) + bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]")) + pups_pattern = regulator(load_env_json("PUPS", "[]")) + + approved_count = 0 + unapproved_count = 0 + needs_review_count = 0 + + for 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 {} + scannermatch = reputation.get("scannermatch") + + logger.debug(f"Evaluating hash: {hash_obj}") + logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") + + # 1. Unapproved: bad publisher or PUP + if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): + logger.debug("Unapproved: Publisher matches bad publisher pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + if re.search(pups_pattern, description, re.IGNORECASE): + logger.debug("Unapproved: Description matches PUP pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + # 2. Approved: signed + if publisher != "Not Signed": + logger.debug("Approved: File is signed and not flagged.") + hash_obj.at_decision = "approved" + approved_count += 1 + continue + + # 3. Approved or Unapproved based on threat level + try: + score = int(scannermatch) # pyright: ignore[reportArgumentType] + logger.debug(f"Parsed scannermatch score: {score}") + if score > threat_tolerance: # pyright: ignore[reportOperatorIssue] + logger.debug("Unapproved: Unsigned file with high threat score.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + else: + logger.debug("Approved: Unsigned file with low threat score.") + hash_obj.at_decision = "approved" + approved_count += 1 + except (ValueError, TypeError): + 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}") + return hashes + + + @classmethod + def export_to_csv(cls, hash_list, directory_path): + """ + Exports a list of Hash objects to a CSV file in the specified directory. + The filename is derived from the variable name of the list if possible, + and includes a timestamp to ensure uniqueness. + """ + filename = "hashes_export.csv" + frame = inspect.currentframe() + if frame is not None and frame.f_back is not None: + callers_local_vars = frame.f_back.f_locals.items() + for var_name, var_val in callers_local_vars: + if var_val is hash_list: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"{var_name}_{timestamp}.csv" + break + else: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"hashes_export_{timestamp}.csv" + + os.makedirs(directory_path, exist_ok=True) + file_path = os.path.join(directory_path, filename) + + df = pd.DataFrame([h.to_dict() for h in hash_list]) + df.to_csv(file_path, index=False) + + logger.info(f"CSV file saved to: {file_path}") + + +@dataclass +class ExecutionHistoryRecord: + # Mandatory fields + username: str + hostname: str + netdomain: str + filename: str + ppolicy: str + policyname: str + policyver: str + commandline: str + publisher: str + sha256: str + datetime: str + + # Optional fields + type: Optional[int] = None + pprocess: Optional[str] = None + gprocess: Optional[str] = None + md5: Optional[str] = None + sha128: Optional[str] = None + sha384: Optional[str] = None + sha512: Optional[str] = None + ip: Optional[str] = None + localip: Optional[str] = None + extid: Optional[str] = None + extname: Optional[str] = None + exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension + extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge + hash_obj: Optional[Hash] = None + + + @classmethod + def from_dict(cls, data: dict): + mandatory_fields = [ + "username", + "hostname", + "netdomain", + "filename", + "ppolicy", + "policyname", + "policyver", + "commandline", + "publisher", + "sha256", + "datetime", + ] + missing_fields = [ + 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}") + + return cls( + username=data["username"], + hostname=data["hostname"], + netdomain=data["netdomain"], + filename=data["filename"], + ppolicy=data["ppolicy"], + policyname=data["policyname"], + policyver=data["policyver"], + commandline=data["commandline"], + publisher=data["publisher"], + sha256=data["sha256"], + datetime=data["datetime"], + type=data.get("type"), + pprocess=data.get("pprocess"), + gprocess=data.get("gprocess"), + md5=data.get("md5"), + sha128=data.get("sha128"), + sha384=data.get("sha384"), + sha512=data.get("sha512"), + ip=data.get("ip"), + localip=data.get("localip"), + extid=data.get("extid"), + extname=data.get("extname"), + exttype=data.get("exttype"), + extbrowser=data.get("extbrowser"), + hash_obj=data.get("hash_obj") + ) + + @classmethod + def from_policies( + cls, api, selected_policies, type_: list, history_days: int + ) -> List["ExecutionHistoryRecord"]: + executions = [] + for policy in selected_policies: + execs = pullPolicyExechistories( + api, policy, type_, history_days, True + ) + if execs: + data = json.loads(execs) + exechistories = data.get("response", {}).get("exechistories", []) + if not exechistories: + continue + + df = pd.DataFrame(exechistories) + 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") + print( + colorText( + f"Staging of Execution history for policy: {policy.name} is complete", + "green", + ) + ) + + return executions + + @staticmethod + def enrich_with_hashes( + api: AirlockAPIWrapper, + executions: List["ExecutionHistoryRecord"] + ) -> List["ExecutionHistoryRecord"]: + """ + Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API. + """ + sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256}) + logger.info(f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records.") + + if not sha256_list: + logger.warning("No sha256 values found in execution records. Skipping enrichment.") + return executions + + logger.debug("Querying hash data from API...") + hash_df = api.hash_query(sha256_list) + logger.info(f"Retrieved {len(hash_df)} hash records from API.") + + hash_objects = [] + required_fields = { + f.name for f in dataclasses.fields(Hash) + if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING + } + + for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()): + row_dict = row.to_dict() + + # Unwrap nested 'data' field if present + if "data" in row_dict and isinstance(row_dict["data"], dict): + row_dict = row_dict["data"] + + # Inject the sha256 back into the row + row_dict["sha256"] = sha256 + + missing = required_fields - row_dict.keys() + if missing: + logger.warning(f"Skipping hash row due to missing fields: {missing}") + logger.debug(f"Row content: {row_dict}") + continue + + try: + hash_obj = Hash.from_dict(row_dict) + hash_objects.append(hash_obj) + except Exception as e: + logger.warning(f"Failed to create Hash from row: {e}") + logger.debug(f"Row content: {row_dict}") + + logger.debug("Converted hash DataFrame to Hash objects.") + + hash_lookup = {h.sha256.strip().lower(): h for h in hash_objects} + logger.debug("Built hash lookup table.") + + enriched_count = 0 + for exec_record in executions: + hash_obj = hash_lookup.get(exec_record.sha256.strip().lower()) + if hash_obj: + exec_record.hash_obj = hash_obj + enriched_count += 1 + + logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.") + return executions + + @staticmethod + def categorize_executions_by_hash_decision(executions: List["ExecutionHistoryRecord"]) -> List["ExecutionHistoryRecord"]: + """ + Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation. + + Modifies the `at_decision` field of each associated Hash object in-place. + + Returns: + List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated. + """ + threat_tolerance = get_protected_value("VT_THREAT_TOLERANCE", cast_type=int) + bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]")) + pups_pattern = regulator(load_env_json("PUPS", "[]")) + + approved_count = 0 + unapproved_count = 0 + needs_review_count = 0 + + for record in executions: + hash_obj = record.hash_obj + if not hash_obj: + continue # Skip if no hash object is attached + + publisher = hash_obj.publisher or "" + description = hash_obj.description or "" + reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {} + scannermatch = reputation.get("scannermatch") + + logger.debug(f"Evaluating hash: {hash_obj}") + logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}") + + # 1. Unapproved: bad publisher or PUP + if re.search(bad_publishers_pattern, publisher, re.IGNORECASE): + logger.debug("Unapproved: Publisher matches bad publisher pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + if re.search(pups_pattern, description, re.IGNORECASE): + logger.debug("Unapproved: Description matches PUP pattern.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + continue + + # 2. Approved: signed + if publisher != "Not Signed": + logger.debug("Approved: File is signed and not flagged.") + hash_obj.at_decision = "approved" + approved_count += 1 + continue + + # 3. Approved or Unapproved based on threat level + try: + score = int(scannermatch) # pyright: ignore[reportArgumentType] + logger.debug(f"Parsed scannermatch score: {score}") + if threat_tolerance is not None and score >= threat_tolerance: + logger.debug("Unapproved: Unsigned file with high threat score.") + hash_obj.at_decision = "unapproved" + unapproved_count += 1 + else: + logger.debug("Approved: Unsigned file with low threat score.") + hash_obj.at_decision = "approved" + approved_count += 1 + except (ValueError, TypeError) as e: + logger.debug(f"Needs Review: Scannermatch score is missing or invalid. ā€” {e}") + hash_obj.at_decision = "needs_review" + needs_review_count += 1 + + logger.debug( + f"Final counts ā€” Needs Review: {needs_review_count}, " + f"Approved: {approved_count}, Unapproved: {unapproved_count}" + ) + + return executions + + + @classmethod + def sort_by_hash_decision( + cls, executions: List["ExecutionHistoryRecord"] + ) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]: + """ + Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups + based on the value of hash_obj.at_decision. + + Returns: + Tuple of lists: (approved, unapproved, needs_review, unknown) + """ + approved = [] + unapproved = [] + needs_review = [] + unknown = [] + + sorted_executions = sorted(executions, key=lambda x: x.filename) + + for record in sorted_executions: + decision = getattr(record.hash_obj, "at_decision", None) + if decision == "approved": + approved.append(record) + elif decision == "unapproved": + unapproved.append(record) + elif decision == "needs_review": + needs_review.append(record) + else: + unknown.append(record) + + logger.info(f"[ExecutionHistoryRecord] Sorted {len(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)}") + logger.info(f" Unknown/Unset: {len(unknown)}") + + return approved, unapproved, needs_review, unknown + + + +""" +executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30) + +ExecutionHistoryRecord.enrich_with_hashes_and_export(executions, hash_objects, "C:/Users/Brandon/Documents/EnrichedExports") +""" diff --git a/models/policy.py b/models/policy.py new file mode 100644 index 0000000..7040f32 --- /dev/null +++ b/models/policy.py @@ -0,0 +1,65 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import json + +""" +Policy model representing policy data and relationships. +""" + + +class Policy: + def __init__(self, groupid, hidden, name, parent): + self.groupid = groupid + self.hidden = hidden + self.name = name + self.parent = parent + + def __repr__(self): + # Show all current attributes, including dynamically added ones + attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + return f"" + + def to_dict(self): + # Return all attributes as a dictionary + return self.__dict__ + + def to_json(self): + # Convert to JSON string, handling non-serializable types gracefully + return json.dumps(self.to_dict(), default=str) + + +class Allowlist: + """ + Represents Allowlist + """ + + def __init__(self, applicationid, name, version): + self.applicationid = applicationid + self.name = name + self.version = version + + def __repr__(self): + # Show all current attributes, including dynamically added ones + attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + return f"" + + def to_dict(self): + # Return all attributes as a dictionary + return self.__dict__ + + def to_json(self): + # Convert to JSON string, handling non-serializable types gracefully + return json.dumps(self.to_dict(), default=str) diff --git a/requirements.txt b/requirements.txt index 8a3c64f..bb00042 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,29 +1,11 @@ -bson==0.5.10 -certifi==2025.8.3 -charset-normalizer==3.4.3 -colorama==0.4.6 -cramjam==2.11.0 -docopt==0.6.2 -dotenv==0.9.9 -fastparquet==2024.11.0 -fsspec==2025.9.0 -idna==3.10 -ijson==3.4.0 -lxml==6.0.0 -markdown-it-py==4.0.0 -mdurl==0.1.2 +cryptography==46.0.1 +keyring==25.6.0 numpy==2.3.2 -packaging==25.0 pandas==2.3.1 -pretty-tables==3.1.0 -pyarrow==21.0.0 -Pygments==2.19.2 -python-dateutil==2.9.0.post0 python-dotenv==1.1.1 -pytz==2025.2 -requests==2.32.4 -six==1.17.0 +pymongo +requests==2.32.5 +schedule==1.2.2 tqdm==4.67.1 -tzdata==2025.2 urllib3==2.5.0 -yarg==0.1.10 \ No newline at end of file +bson==0.5.10 \ No newline at end of file diff --git a/services/API.py b/services/API.py new file mode 100644 index 0000000..c618a79 --- /dev/null +++ b/services/API.py @@ -0,0 +1,286 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import json +import logging +from typing import Dict, List, Optional + +import pandas as pd +import requests + +logger = logging.getLogger(__name__) + + + +class AirlockAPIWrapper: + """ + A wrapper class for interacting with the Airlock API. + Provides methods for managing agents, policies, hashes, OTPs, and execution history. + """ + + def __init__(self, base_url: str, api_key: str): + """ + Initialize the API wrapper. + + Parameters: + - base_url (str): Base URL of the Airlock API. + - api_key (str): API key for authentication. + """ + self.base_url = base_url.rstrip("/") + self.api_key = api_key + self.headers = {"X-APIKey": self.api_key} + + def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict: + """ + Internal method to send POST requests to the API. + + Parameters: + - endpoint (str): API endpoint. + - payload (dict, optional): Request payload. + + Returns: + - dict: JSON response from the API. + """ + url = f"{self.base_url}{endpoint}" + data = json.dumps(payload or {}) + try: + logger.debug(f"POST Request to {url} with payload: {payload}") + response = requests.post(url, headers=self.headers, data=data, verify=False) + response.raise_for_status() + logger.debug(f"Response received from {url}") + return response.json() + except requests.exceptions.RequestException as e: + logger.error(f"API request failed: {e}") + raise + + # Allowlist Management + def allowlist_find_all(self) -> pd.DataFrame: + """ + Retrieve all applications in the allowlist. + + Returns: + - pd.DataFrame: DataFrame containing allowlisted applications. + """ + result = self._post("/v1/application", {}) + return pd.DataFrame(result["response"]["applications"]) + + # Agent Management + def agent_find_all(self) -> pd.DataFrame: + """Retrieve all agents.""" + result = self._post("/v1/agent/find", {}) + return pd.DataFrame(result["response"]["agents"]) + + def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame: + """Find agents by hostname.""" + payload = {"hostname": hostname} + result = self._post("/v1/agent/find", payload) + return pd.DataFrame(result["response"]["agents"]) + + def agent_find_by_id(self, agentid: str) -> pd.DataFrame: + """Find agents by agent ID.""" + payload = {"agentid": agentid} + result = self._post("/v1/agent/find", payload) + return pd.DataFrame(result["response"]["agents"]) + + def agent_find_by_status(self, status: int) -> pd.DataFrame: + """Find agents by status (0 = Offline, 1 = Online, 3 = Safemode).""" + payload = {"status": status} + result = self._post("/v1/agent/find", payload) + return pd.DataFrame(result["response"]["agents"]) + + def agent_find_by_username(self, username: str) -> pd.DataFrame: + """Find agents by username.""" + payload = {"username": username} + result = self._post("/v1/agent/find", payload) + return pd.DataFrame(result["response"]["agents"]) + + def agent_move(self, agentid: str, groupid: str) -> dict: + """Move an agent to a different group.""" + payload = {"agentid": agentid, "groupid": groupid} + return self._post("/v1/agent/move", payload) + + def agents_find_by_group(self, groupid: str) -> pd.DataFrame: + """Find agents by group ID.""" + payload = {"groupid": groupid} + result = self._post("/v1/agent/find", payload) + return pd.DataFrame(result["response"]["agents"]) + + # Hash Management + def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict: + """Add hashes to the allowlist for a specific application.""" + payload = {"applicationid": applicationid, "hashes": hashes} + return self._post("/v1/hash/application/add", payload) + + def hash_query(self, hashes: List[str]) -> pd.DataFrame: + """Query information about specific hashes.""" + payload = {"hashes": hashes} + result = self._post("/v1/hash/query", payload) + return pd.DataFrame(result["response"]["results"]) + + # OTP Management + def otp_find_active(self) -> pd.DataFrame: + """Find active OTPs.""" + payload = {"status": "1"} + result = self._post("/v1/otp/usage", payload) + return pd.DataFrame(result["response"]["otpusage"]) + + def otp_find_awaiting(self) -> pd.DataFrame: + """Find OTPs that are awaiting activation.""" + 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 = { + "duration": str(duration), + "agentid": str(agentid), + "purpose": purpose, + } + result = self._post("/v1/otp/retrieve", payload) + return result["response"]["otpcode"] + + def otp_get_activities(self, otpid: str) -> pd.DataFrame: + """Retrieve activities associated with a specific OTP.""" + 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. + Parameters: + - otpid (str): The ID of the OTP to revoke. + Returns: + - dict: JSON response from the API. + """ + 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) + + + # Policy Management + def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict: + """Add path exclusions to a policy group.""" + payload = {"groupid": groupid, "path": paths} + return self._post("/v1/group/path/add", payload) + + def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict: + """Add publishers to a policy group.""" + payload = {"groupid": groupid, "publisher": publishers} + return self._post("/v1/group/publisher/add", payload) + + def policy_clone(self, source_groupid: str, target_groupid: str) -> dict: + """Clone a policy from one group to another.""" + payload = {"groupid": source_groupid, "targetgroupid": target_groupid} + return self._post("/v1/group/assign", payload) + + def policy_find_all(self) -> pd.DataFrame: + """Retrieve all policy groups.""" + result = self._post("/v1/group") + return pd.DataFrame(result["response"]["groups"]) + + def policy_list_agents(self, groupid: str) -> pd.DataFrame: + """List agents assigned to a specific policy group.""" + 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} + result = self._post("/v1/group/policies", payload) + return pd.DataFrame(result["response"]["applications"]) + + def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict: + """Set audit mode for a policy group. 1=Audit, 0=Enforcement""" + payload = {"groupid": groupid, "auditmode": auditmode} + return self._post("/v1/group/settings/auditmode", payload) + + # Execution History + 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]: + """ + Retrieve execution history logs. + + "datefrom":"", //(Optional) Datefrom is for date range search, formatted as "YYYY-MM-DD" + "dateto":"", //(Optional) Dateto is for date range search, formatted as "YYYY-MM-DD" + "category":"", //(Optional) Category for filtering type + "hostname":"", //(Optional) Hostname to filter + "username":"admin", //(Optional) Username to filter + "netdomain":"", //(Optional) Domain (or group) to filter + "filename":"", //(Optional) Filename to filter + "ppolicy":"", //(Optional) Parent Policy name to filter + "policyname":"", //(Optional) Policy name to filter + "policyver":"", //(Optional) Policy version to filter (e.g. "v95") + "commandline":"", //(Optional) Commandline to filter + "publisher":"", //(Optional) Publisher to filter + "pprocess":"", //(Optional) Parent Process to filter + "sha256":"", //(Optional) SHA256 hash to filter + "contains":["hostname"], //(Optional) Contains is an array for wildcard searches on a filter + "limit":"5" //(Optional) Limit the amount of results returned, default set to 50 + + """ + + payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name} + result = self._post("/v1/getexechistory", payload) + return result["response"]["exechistory"] + + +""" +from services.API import AirlockAPIWrapper + + +api = AirlockAPIWrapper(base_url="https://airlock.example.com/api", api_key="your_api_key_here") + +#Example: Get all agents + +agents_df = api.agent_find_all() +print("All Agents:") +print(agents_df) +""" diff --git a/services/agenthandler.py b/services/agenthandler.py new file mode 100644 index 0000000..531fc41 --- /dev/null +++ b/services/agenthandler.py @@ -0,0 +1,323 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + +from dataclasses import asdict +from datetime import datetime, timedelta +import json +import logging +import os +import re +from typing import List + +import pandas as pd + +from flows.prepPolicy import selectPolicies +from models.agent import Agent +from models.policy import Policy +from services.API import AirlockAPIWrapper +from utils.configmanager import get_protected_json, load_env +from utils.selector import Selector +from utils.utils import colorText, get_sanitized_input + +logger = logging.getLogger(__name__) + + +def devicehistory(api: AirlockAPIWrapper, outputjson: bool): + agents = selectAgents(api) + history_days = Selector.select_value( + prompt="Enter how many days of history to pull (1ā€“150): ", + value_type=int, + valid_range=(1, 150), + ) + + if not agents or not history_days: + print(colorText("No agents selected or invalid history range.", "red")) + return + + historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d") + today = datetime.now().strftime("%Y-%m-%d") + + all_history = [] + + for agent in agents: + 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")) + continue + + if isinstance(exechistory, list): + for block in exechistory: + record = { + "Command": block.get("commandline", "N/A"), + "Date": block.get("datetime", "N/A"), + "Filename": block.get("filename", "N/A"), + "Policy Name": block.get("policyname", "N/A"), + "Hostname": block.get("hostname", "N/A"), + "Hash": block.get("sha256", "N/A"), + } + all_history.append(record) + + if not outputjson: + for key, value in record.items(): + print(colorText(f"{key}: {value}", "green")) + print("\n") + else: + print(colorText(f"No execution history found for {agent.hostname}.", "yellow")) + + if outputjson: + print(json.dumps(all_history, indent=2)) + + +def findAllAgents(api): + # Step 1: Load data from API + policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()] + agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()] + + for agent in agents: + agent.enrich_with_policies(policies) + + return agents + +def findAgents(api, return_dataframe): + agents = selectAgents(api) + working_dir = load_env("WORKING_DIR") + + if not agents: + logging.warning("No agents or policies found.") + print("No agents matched the criteria.") + return + + # Convert enriched agents to DataFrame + agent_dicts = [asdict(agent) for agent in agents] + agent_df = pd.DataFrame(agent_dicts) + + if return_dataframe: + logging.debug("Returning DataFrame to caller.") + return agent_df + + # Otherwise, print and optionally export + 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': + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + filename = f"agentsearch_{timestamp}.csv" + file_path = os.path.join(str(working_dir), filename) + + agent_df.to_csv(file_path, index=False) + logging.info(f"Exported DataFrame to {file_path}") + + print( + colorText( + f"\nāœ… Matched devices exported to: {working_dir}\\{filename}", + "green", + ) + ) + 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("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_\- ]+$') + + while True: + line = get_sanitized_input("") + stripped_line = line.strip() + + if stripped_line == "": + empty_line_count += 1 + if empty_line_count == 2: + break + continue + else: + empty_line_count = 0 + + 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")) + + return [name for name in device_input_lines if name] + + +def choose_match_type() -> bool: + print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white")) + return get_sanitized_input("").strip().lower() in ["y", "yes"] + + +def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']: + if use_exact: + return [ + agent for agent in agents + if agent.hostname.lower() in [name.lower() for name in device_names] + ] + else: + pattern = "|".join(map(re.escape, device_names)) + regex = re.compile(pattern, re.IGNORECASE) + 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): + if use_exact: + 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)] + + 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']): + for agent in agents: + agent.enrich_with_policies(policies) + + +def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']: + device_names = collect_device_names() + if not device_names: + logger.debug("No device names entered") + print(colorText("āš ļø No device names entered.", "red")) + return [] + + use_exact = choose_match_type() + + policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] + agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()] + matched_agents = match_agents(device_names, agents, use_exact) + matched_agents.sort(key=lambda agent: agent.hostname.lower()) + + show_unmatched(device_names, matched_agents, use_exact) + + if not matched_agents: + logger.debug("āŒ No matching devices found.") + print(colorText("āŒ No matching devices found.", "red")) + return [] + + print(colorText(f"āœ… Found {len(matched_agents)} matching device(s).", "green")) + logger.info("Matched agent hostnames:") + rows = (len(matched_agents) + 2) // 3 # 3 columns + for row in range(rows): + line = "" + for col in range(3): + idx = row + col * rows + 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:" + ) + + if not matched_agents: + logger.debug("āŒ No matching devices remain after refinement.") + print(colorText("āŒ No matching devices remain after refinement.", "red")) + return [] + + enrich_agents(matched_agents, policies) + return matched_agents + + +def moveAgentToRelatedPolicy( + api: AirlockAPIWrapper, + agent: Agent, + mode: str = "audit", +): + """ + Moves an agent between audit and enforcement policies based on the mode. + + Args: + api: AirlockAPIWrapper instance. + agent: Agent object. + policy_relationship_map: Dict mapping enforcement ā†’ audit. + mode: 'audit' to move to audit, 'enforcement' to move to enforcement. + """ + policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + + if mode == "audit": + 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.") + return + else: + logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.") + return + + elif mode == "enforcement": + inverse_map = {v: k for k, v in policy_relationship_map.items()} + 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.") + return + else: + logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.") + return + + else: + logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.") + return + + result = api.agent_move(agent.agentid, target_policy) + return result + + +def toggleEnforcement(api: AirlockAPIWrapper): + choices = ["Audit", "Enforcement", "Exit"] + print(colorText("Move devices to which state?:", "yellow")) + direction = Selector.select_string(choices, False, False) + if direction == "Exit": + pass + else: + devices = selectAgents(api) + for device in devices: + print(device.hostname) + 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()) + 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: ") + 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}?") + 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 diff --git a/services/policyhandler.py b/services/policyhandler.py new file mode 100644 index 0000000..5acc422 --- /dev/null +++ b/services/policyhandler.py @@ -0,0 +1,244 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + + +import datetime +import gc +import json +import logging +import os +import sys + +from bson import ObjectId +import pandas as pd +import tqdm + +from models.policy import Policy +from services.API import AirlockAPIWrapper +from utils.configmanager import get_protected_json +from utils.setup import get_base_directory +from utils.utils import areYouSure, colorText, get_sanitized_input + +logger = logging.getLogger(__name__) + + + +def pullPolicyExechistories( + api: AirlockAPIWrapper, + policy: Policy, + type: list, + days, + outputjson: bool, +): + + file_path = f"{get_base_directory()}\\cache\\chunkinator.json" + + # Ensure the file exists + if not os.path.exists(file_path): + with open(file_path, "w") as file: + json.dump({"error": "Success", "response": {"exechistories": []}}, file) + logger.debug(f"File '{file_path}' has been created.") + else: + logger.debug(f"File '{file_path}' already exists.") + + checkpoint = str(skipback(days)) + json_output = {"error": "Success", "response": {"exechistories": []}} + + with tqdm.tqdm( + file=sys.stdout, + leave=True, + total=10000, + desc=f"Checkpoint Progress: {checkpoint}", + colour="blue", + initial=1, + ) as filebar: + with tqdm.tqdm( + file=sys.stdout, + leave=True, + total=100, + desc=f"Total of {policy} Complete: ", + ) as pbar: + while True: + histories = api.history_logging( + type=type, checkpoint=checkpoint, policy= [policy.name] + ) + + # Ensure histories is a list of dictionaries + if not isinstance(histories, list) or not all( + isinstance(h, dict) for h in histories + ): + logger.error( + "Unexpected response format from API. Expected list of dictionaries." + ) + break + + filebar.total = len(histories) + + if not histories: + break + + for index, history_item in enumerate(histories): + if ( + "checkpoint" not in history_item + or "datetime" not in history_item + ): + continue # Skip malformed entries + + # Update checkpoint on last item + if index == len(histories) - 1: + 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] + "%Y-%m-%dT%H:%M:%SZ", + ).date() + except ValueError: + continue # Skip if date format is invalid + + if ( + datetime.date.today() - datetime.timedelta(days=days) + ) <= history_date: + json_output["response"]["exechistories"].append(history_item) + + filebar.update(1) + filebar.refresh() + + # Deduplicate entries + seen = {} + if os.path.exists(file_path): + with open(file_path, "r") as file: + existing_data = json.load(file) + combined = ( + existing_data["response"]["exechistories"] + + json_output["response"]["exechistories"] + ) + else: + combined = json_output["response"]["exechistories"] + + for entry in combined: + key = ( + entry.get("sha256"), + entry.get("filename"), + entry.get("hostname"), + ) + seen[key] = entry + + deduplicated = list(seen.values()) + with open(file_path, "w") as file: + json.dump( + { + "error": "Success", + "response": {"exechistories": deduplicated}, + }, + file, + ) + + json_output["response"]["exechistories"].clear() + + # Update progress bar based on last valid item + try: + last_date = datetime.datetime.strptime( + history_item["datetime"].replace(" +0000 UTC", ""), # type: ignore + "%Y-%m-%dT%H:%M:%SZ", + ).date() + date_diff = datetime.date.today() - last_date + percentage_diff = ( + ((days + 10) - date_diff.days) / (days + 10) + ) * 100 + pbar.n = round(percentage_diff) + pbar.set_description_str(f"Total of {policy} Complete: ") + pbar.refresh() + except Exception: + pass + + filebar.n = 1 + + # Final output + with open(file_path, "r") as file: + final_output = json.load(file) + os.remove(file_path) + + return json.dumps(final_output) if outputjson else None + + +def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days): + executionhist_policy = pd.DataFrame() + exehist = pullPolicyExechistories(api, policy, type, days, True) + if exehist is not None: + data = json.loads(exehist) + executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) + if not executionhist_policy.empty: + executionhist_policy = executionhist_policy[ + [ + "datetime", + "sha256", + "publisher", + "filename", + "hostname", + "username", + "pprocess", + "gprocess", + "commandline", + ] + ] + executionhist_policy["policy"] = policy # Add policy column here + executionhist_policy = executionhist_policy.drop_duplicates( + subset=["sha256", "filename", "hostname"] + ) + executionhist_policy = executionhist_policy.sort_values( + by=["sha256", "filename"] + ) + logger.debug( f"Staging of Execution history for policy: {policy} is complete") + print( + colorText( + f"Staging of Execution history for policy: {policy} is complete", + "green", + ) + ) + del data + del exehist + gc.collect() + return executionhist_policy + + +def skipback(days): + """ + Generate a MongoDB ObjectId for a given number of days ago from today. + """ + adjusted_days = days + date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta( + days=adjusted_days + ) + timestamp = int(date_days_ago.timestamp()) + hex_timestamp = format(timestamp, "08x") + objectid_hex = hex_timestamp + "0000000000000000" + return ObjectId(objectid_hex) + + +def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper): + policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + for enforcement_policy, audit_policy in policy_relationship_map.items(): + api.policy_clone(enforcement_policy, audit_policy) + api.policy_set_auditmode(audit_policy, "1") + + +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 diff --git a/services/security.py b/services/security.py new file mode 100644 index 0000000..aa95929 --- /dev/null +++ b/services/security.py @@ -0,0 +1,172 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import base64 +from getpass import getpass +import logging +import os +import platform +import re +import sys + +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +import keyring + +# Constants +KDF_ITERATIONS = 200_000 +SALT_SIZE = 16 # 128-bit Salt +NONCE_SIZE = 12 # AES-GCM +KEY_SIZE = 32 # AES-256 + +logger = logging.getLogger(__name__) + + +def _derive_key(password: bytes, salt: bytes) -> bytes: + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=KEY_SIZE, + salt=salt, + iterations=KDF_ITERATIONS, + ) + return kdf.derive(password) + + +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}") + + +def store_api_key(service: str, username: str, api_key: str, password: str): + configure_keyring_backend() + salt = os.urandom(SALT_SIZE) + key = _derive_key(password.encode(), salt) + aesgcm = AESGCM(key) + nonce = os.urandom(NONCE_SIZE) + ct = aesgcm.encrypt(nonce, api_key.encode(), associated_data=None) + blob = salt + nonce + ct + b64 = base64.b64encode(blob).decode() + keyring.set_password(service, username, b64) + + + 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...") + + try: + _ = input() + except Exception: + pass + + _ = None + sys.exit(0) + + +def retrieve_api_key(service: str, username: str, password: str) -> str: + configure_keyring_backend() + b64 = keyring.get_password(service, username) + if b64 is None: + 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:] + key = _derive_key(password.encode(), salt) + aesgcm = AESGCM(key) + pt = aesgcm.decrypt(nonce, ct, associated_data=None) + return pt.decode() + + +def api_key_exists(service: str, username: str) -> bool: + configure_keyring_backend() + return keyring.get_password(service, username) is not None + + +def check_password_complexity(password: str) -> bool: + if len(password) < 12: + return False + if not re.search(r"[A-Z]", password): + return False + if not re.search(r"[a-z]", password): + return False + if not re.search(r"[0-9]", password): + return False + if not re.search(r"[^A-Za-z0-9]", password): + return False + return True + + +def getAPI(USERNAME, SERVICE_NAME): + logging.debug( + f"Checking for stored API key for user '{USERNAME}' in service '{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: ") + try: + apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password) + logging.debug("API key successfully retrieved.") + return apikey + except Exception as e: + logging.warning(f"Attempt {attempt} failed: {str(e)}") + 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") + + while True: + password = getpass("Create a password to encrypt your API key: ") + confirm_password = getpass("Confirm your password: ") + + if password != confirm_password: + logging.warning("Passwords do not match. Try again.") + continue + + if check_password_complexity(password): + try: + store_api_key(SERVICE_NAME, USERNAME, api_key, password) + logging.info("API key stored securely.") + break + except Exception as e: + logging.error(f"Failed to store API key: {e}") + break + else: + logging.warning("Password does not meet complexity requirements. Try again.") + + +class APIKeyManager: + _api_key = None + + @classmethod + def load(cls, service: str, username: str, password: str): + cls._api_key = retrieve_api_key(service, username, password) + + @classmethod + 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 diff --git a/utils/allowlist.py b/utils/allowlist.py deleted file mode 100644 index 0f02958..0000000 --- a/utils/allowlist.py +++ /dev/null @@ -1,155 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import datetime -import requests -import json -import os -import utils.pretty as ct -import ijson -import os -from bson import ObjectId -import datetime -import tqdm -import sys - -def pullPolicyExechistories(url, policiesnames, days, outputjson: bool): - file_path = 'chunkinator.json' - if not os.path.exists(file_path): - with open(file_path, 'w') as file: - json.dump({'error': 'Success', 'response': {'exechistories': []}}, file) - print(f"File '{file_path}' has been crated.") - else: - print(f"File '{file_path}' already exists.") - headers = {"X-APIKey": os.getenv('APIKEY')} - checkpoint = str(skipback(days)) - json_output = {'error': 'Success', 'response': {'exechistories': []}} - with tqdm.tqdm(file=sys.stdout, leave=True, total=10000, desc=f"Checkpoint Progess: {checkpoint}", colour="blue", initial=1) as filebar: - with tqdm.tqdm(file=sys.stdout, leave=True, total=100, desc=f"Total of {policiesnames} Complete: ") as pbar: - while True: - json_response_data = checkpoint_stomper(checkpoint, url, policiesnames, headers) - histories = json_response_data['response']['exechistories'] - filebar.total=len(histories) - if not histories: - break - match_found = True - if match_found == True: - for index, item in enumerate(histories): - if index == len(histories) - 1: - checkpoint = item['checkpoint'] - filebar.desc = f"Checkpoint Progress: {checkpoint}" - break - else: - if (datetime.date.today() - datetime.timedelta(days=days) > datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date()): - pass - else: json_output['response']['exechistories'].append(item) - filebar.update(1) - filebar.refresh() - seen = {} - if os.path.exists(file_path): - with open(file_path, 'r') as file: - existing_data = json.load(file) - combined = existing_data['response']['exechistories'] + json_output['response']['exechistories'] - else: - combined = json_output['response']['exechistories'] - for item in combined: - key = (item.get('sha256'), item.get('filename'), item.get('hostname')) - seen[key] = item - deduplicated = list(seen.values()) - with open(file_path, 'w') as file: - json.dump({'error': 'Success', 'response': {'exechistories': deduplicated}}, file) - json_output['response']['exechistories'].clear() - date_diff = datetime.date.today() - datetime.datetime.strptime(item['datetime'].replace(' +0000 UTC', ''), '%Y-%m-%dT%H:%M:%SZ').date() - percentage_diff = (((days + 10) - date_diff.days) / (days + 10)) * 100 - pbar.n = round(percentage_diff) - pbar.set_description_str(f"Total of {policiesnames} Complete: ") - pbar.refresh() - filebar.n = 1 - with open(file_path, 'r') as file: - final_output = json.load(file) - os.remove(file_path) - return json.dumps(final_output) if outputjson else None - -def checkpoint_stomper(checkpoint, url, policy, headers): - json_output = {'error': 'Success', 'response': {'exechistories': []}} - endpoint = url + '/v1/logging/exechistories' - payload_dict = { - "type":[1,2,6,7], - "checkpoint": checkpoint, - "policy": [policy] - } - payload = json.dumps(payload_dict) - with requests.request("POST", endpoint, headers=headers, data=payload, verify=False, stream=True) as response: - parser = ijson.items(response.raw, 'response.exechistories.item') - for item in parser: - key = (item.get('sha256'), item.get('hostname')) - if key not in json_output: - json_output['response']['exechistories'].append(item) - parse_text = json.loads(json.dumps(json_output)) - return parse_text - -def listPolicies(url): - endpoint = url + '/v1/group' - print(ct.colorText("[+] Grabbing All Policies", "cyan")) - payload = {} - headers = { - "X-APIKey": os.getenv('APIKEY') - } - response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) - parse_text = json.loads(response.text) - policiesnames = [] - policyids = [] - for index, list in enumerate(parse_text['response']['groups'], start=1): - print(ct.colorText(f"{index}. {list['name']}", "yellow")) - policiesnames.append(list['name']) - policyids.append(list['groupid']) - choice = input(ct.colorText("Select Policy Group: ", "white")) - choice = int(choice) - 1 - return choice, policiesnames, policyids - -def listAllowlists(url): - endpoint = url + '/v1/application' - print(ct.colorText("[+] Grabbing All Allowlists", "cyan")) - payload = {} - headers = { - "X-APIKey": os.getenv('APIKEY') - } - response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) - parse_text = json.loads(response.text) - policiesnames = [] - policyids = [] - for index, list in enumerate(parse_text['response']['applications'], start=1): - if index >= 38: - print(ct.colorText(f"{index}. {list['name']}", "yellow")) - policiesnames.append(list['name']) - policyids.append(list['applicationid']) - choice = int(input(ct.colorText("Select allowlist: ", "white"))) - if choice < 38: - print(ct.colorText("Please only choose an allowlist designed for this use - '38+'","red")) - elif choice >= 38: - choice = choice - 38 - return choice, policiesnames, policyids - #Need else and catch for upper bound - -def skipback(days): - """ - Generate a MongoDB ObjectId for a given number of days ago from today. - Adds 1 extra day to the input to look further back. - """ - adjusted_days = days + 10 - date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=adjusted_days) - timestamp = int(date_days_ago.timestamp()) - hex_timestamp = format(timestamp, '08x') - objectid_hex = hex_timestamp + '0000000000000000' - return ObjectId(objectid_hex) \ No newline at end of file diff --git a/utils/configmanager.py b/utils/configmanager.py new file mode 100644 index 0000000..62c2bdf --- /dev/null +++ b/utils/configmanager.py @@ -0,0 +1,130 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import json +import logging +import os +from pathlib import Path +import sys +from typing import Callable, Optional, TypeVar + +T = TypeVar("T") +logger = logging.getLogger(__name__) + +PROTECTED_KEYS = [ + "APPNAME", + "LOG_LEVEL", + "PATH_EXCLUSION_CONST", + "MIN_FILES_FOR_PATH", + "VT_THREAT_TOLERANCE", + "POLICY_MAP_ENF_AUD" +] + +_protected_config = {} + +def get_system_config_path() -> Path: + # Check inside bundled EXE directory first + bundled_dir = Path(getattr(sys, '_MEIPASS', '')) + bundled_path = bundled_dir / "system_config.json" + if bundled_path.exists(): + return bundled_path + + # Fallback to external location + return Path(__file__).parent.parent / "system_config.json" + +def load_protected_config() -> dict: + global _protected_config + try: + with open(get_system_config_path(), "r") as f: + system_config = json.load(f) + except FileNotFoundError: + logging.warning("āš ļø system_config.json not found. Using built-in defaults.") + system_config = { + "APPNAME": "AirlockTools", + "PATH_EXCLUSION_CONST": 4, + "MIN_FILES_FOR_PATH": 4, + "VT_THREAT_TOLERANCE": 4, + "POLICY_MAP_ENF_AUD": { + "enforced_id": "audit_id" + } + } + + _protected_config = {key: system_config[key] for key in PROTECTED_KEYS} + return _protected_config + +def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]: + value = _protected_config.get(key) + if value is None: + logging.warning(f"Protected config key '{key}' not found.") + return default + try: + if isinstance(value, str): + value = value.strip("'\"") + return cast_type(value) + except (ValueError, TypeError): + logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.") + return default + +def get_protected_json(key: str, default: str = "{}") -> dict: + raw = _protected_config.get(key, default) + if isinstance(raw, dict): + return raw + try: + return json.loads(raw) + except json.JSONDecodeError: + try: + escaped = raw.encode('unicode_escape').decode('utf-8') + return json.loads(escaped) + except Exception as e: + logging.error(f"Failed to parse protected JSON key '{key}': {e}") + return json.loads(default) + + + + +def load_env_json(key: str, default: str): + raw = os.getenv(key, default) + try: + return json.loads(raw) + except json.JSONDecodeError: + try: + escaped = raw.encode('unicode_escape').decode('utf-8') + return json.loads(escaped) + except Exception as e: + logging.error(f"Failed to parse {key}: {e}") + return json.loads(default) + +def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]: + """ + Safely retrieves an environment variable and casts it to the desired type. + + Parameters: + key (str): The name of the environment variable. + cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str. + default (Optional[T], optional): Default value if the variable is not set or invalid. + + Returns: + Optional[T]: The casted value or the default. + """ + value = os.getenv(key) + if value is None: + logger.warning(f"Environment variable '{key}' not set.") + return default + try: + value = value.strip("'\"") # Strip surrounding quotes + return cast_type(value) + except (ValueError, TypeError): + logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.") + return default \ No newline at end of file diff --git a/utils/getdeviceevents.py b/utils/getdeviceevents.py deleted file mode 100644 index fbb5d66..0000000 --- a/utils/getdeviceevents.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import datetime -import requests -import json -import os -import utils.pretty as ct - -def devicehistory(url, outputjson: bool): - endpoint = url + '/v1/getexechistory' - print("\n") - print(ct.colorText("1. Today", "yellow")) - print(ct.colorText("2. Last 24 Hours", "yellow")) - print(ct.colorText("3. Past 7 Days", "yellow")) - print(ct.colorText("4. Past 30 Days", "yellow")) - print(ct.colorText("5. Custom Date Range","yellow")) - choice = input(ct.colorText("\nSelect Date Range: ", "white")) - today = datetime.date.today() - today = today.strftime("%Y-%m-%d") - if choice == '1': - date_selected = today - elif choice == '2': - date_selected = datetime.date.today() - datetime.timedelta(days=1) - date_selected = date_selected.strftime('%Y-%m-%d') - elif choice == '3': - date_selected = datetime.date.today() - datetime.timedelta(days=7) - date_selected = date_selected.strftime('%Y-%m-%d') - elif choice == '4': - date_selected = datetime.date.today() - datetime.timedelta(days=30) - date_selected = date_selected.strftime('%Y-%m-%d') - elif choice == "5": - print(ct.colorText("Please Input Dates as YYYY-MM-DD", "cyan")) - date_selected = input(ct.colorText("From: ", "white")) - today = input(ct.colorText("Date To: ", "white")) - print(ct.colorText("WARNING: Device Name is Case Sensitive", "red")) - device = input(ct.colorText("Enter Device Name: ", "white")) - payload_dict = { - "datefrom": date_selected, - "dateto": today, - "hostname": device - } - payload = json.dumps(payload_dict) - print(ct.colorText(payload, "green")) - headers = { - "X-APIKey": os.getenv('APIKEY') - } - - response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) - - if outputjson: - return response - - parse_text = json.loads(response.text) - - # Safely get exechistory - exechistory = parse_text.get('response', {}).get('exechistory') - - if isinstance(exechistory, list): - for block in exechistory: - print(ct.colorText(f"Command: {block.get('commandline', 'N/A')}", "green")) - print(ct.colorText(f"Date: {block.get('datetime', 'N/A')}", "green")) - print(ct.colorText(f"Filename: {block.get('filename', 'N/A')}", "green")) - print(ct.colorText(f"Policy Name: {block.get('policyname', 'N/A')}", "green")) - print(ct.colorText(f"Hostname: {block.get('hostname', 'N/A')}", "green")) - print(ct.colorText(f"Hash: {block.get('sha256', 'N/A')}", "green")) - print("\n") - else: - print(ct.colorText("No execution history found or data is not in expected format.", "red")) \ No newline at end of file diff --git a/utils/hashfunctions.py b/utils/hashfunctions.py deleted file mode 100644 index d617118..0000000 --- a/utils/hashfunctions.py +++ /dev/null @@ -1,375 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import gc -import json -import os -import pandas as pd -import requests -import utils.pathfunctions as pathf -import utils.hashfunctions as hashf -import utils.pretty as ct -from AirlockTools import tryToReadCSV - -def aggregateHashes(executions_json) -> pd.DataFrame: - """ - Takes the executions, aggregates all the data with sha256 as primary, then returns aggregated dataframe - """ - data = json.loads(executions_json) - df = pd.DataFrame(data["response"]["exechistories"]) - - if df.empty: - return df - print(df) - # Aggregate by sha256, deduplicate lists, and preserve order - agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index() - - # Add a column for the number of unique hostnames - agg_df["num_devices"] = agg_df["hostname"].apply(len) - - # Sort by num_devices in descending order - agg_df = agg_df.sort_values("num_devices", ascending=False) - - return agg_df - -def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame: - """ - Takes output of aggregatedHashes, queries API for those hashes, flattens response while keeping one row per hash, - aggregate applications and baselines into lists, then merges results back into agg_df to create a - """ - if 'sha256' not in agg_df.columns or agg_df.empty: - print("āš ļø 'sha256' column missing or DataFrame is empty. Skipping API query.") - return agg_df.copy() # Return as-is to avoid breaking downstream logic - - endpoint = url + '/v1/hash/query' - payload = { - "hashes": agg_df['sha256'].tolist() - } - - headers = {"X-APIKey": os.getenv('APIKEY')} - payload = json.dumps(payload) - - response = requests.post(endpoint, headers=headers, data=payload, verify=False) - data = response.json() - results = data.get("response", {}).get("results", []) - - rows = [] - for res in results: - row = {"sha256": res.get("sha256"), "result": res.get("result")} - - if "data" in res: - d = res["data"] - for key in ["filename", "filepath", "description", "filesize", "md5", - "productname", "productversion", "publisher", "createtime", "modtime", - "sha128", "sha384", "sha512", "datetime"]: - row[key] = d.get(key) - - row["applications"] = d.get("applications", []) - row["baselines"] = d.get("baselines", []) - - reputation = d.get("reputation", {}) - for k, v in reputation.items(): - row[f"reputation_{k}"] = v - - rows.append(row) - - df_api = pd.DataFrame(rows) - - if 'sha256' not in df_api.columns: - print("āš ļø API response missing 'sha256'. Skipping merge.") - return agg_df.copy() - - df = agg_df.merge(df_api, on="sha256", how="left") - - # Only include columns that exist to avoid KeyErrors - expected_columns = ['sha256', 'filename_x', 'description', 'productname', 'productversion', - 'publisher_y', 'publisher_x', 'netdomain', 'hostname', 'username', - 'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount', - 'reputation_status', 'reputation_threatlevel', 'reputation_threatname', - 'reputation_timestamp', 'pprocess', 'gprocess', 'commandline'] - - available_columns = [col for col in expected_columns if col in df.columns] - aug_df = df[available_columns] - - return aug_df - -def categorizeHashes(first_policy, second_policy, df: pd.DataFrame, threat_tolerance: int, untrusted_publishers, pups: list): - if untrusted_publishers is None: untrusted_publishers = [] - if pups is None: pups = [] - - def reputationtool(row): - val = row["reputation_scannermatch"] - if pd.isna(val) or val == "N/A": - return row["publisher"] == "Not Signed" - try: - return int(val) > threat_tolerance - except (ValueError, TypeError): - return row["publisher"] == "Not Signed" - - df["reputation_flag"] = df.apply(reputationtool, axis=1) - - mask_needsreview = ( - ((df["publisher"] == "Not Signed") & df["reputation_flag"]) | - (df["reputation_status"] == "UNKNOWN") - ) - - mask_approved = ( - ( - (df["publisher"] != "Not Signed") & - ~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) & - ~df["reputation_status"].isna() & - ~df["description"].str.contains(pathf.regulator(pups), case=False, na=False) - ) | - ( - (df["publisher"] == "Not Signed") & - ~df["reputation_flag"] & - ~df["publisher"].str.contains(pathf.regulator(untrusted_publishers), case=False, na=False) & - ~df["reputation_status"].isna() & - ~df["description"].str.contains(pathf.regulator(pups), case=False, na=False) - ) - ) - - needsreview_df = df[mask_needsreview] - approved_df = df[mask_approved] - unapproved_df = df[~(mask_needsreview | mask_approved)] - - needsreview_df.to_parquet(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet", index=False) - approved_df.to_parquet(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet", index=False) - unapproved_df.to_parquet(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet", index=False) - - del needsreview_df - del approved_df - del unapproved_df - gc.collect() - -def explode_and_deduplicate(df): - df['sha256'] = df['sha256'].str.split(',') - df = df.explode('sha256') - return df.drop_duplicates().reset_index(drop=True) - -def clean_sha256(df, column='sha256'): - """Discard quotes, brackets, and whitespace from sha256 values.""" - df[column] = df[column].astype(str).str.strip("'[]\" ") - return df - -def destinationHashes( - df_approved_paths: pd.DataFrame, - df_approved_hashes: pd.DataFrame, - df_hashes_auto_approved: pd.DataFrame, - df_hashes_manually_approved: pd.DataFrame, -): - # Deduplicate and explode all input DataFrames - df_approved_paths = explode_and_deduplicate(df_approved_paths) - df_approved_hashes = explode_and_deduplicate(df_approved_hashes) - df_hashes_auto_approved = explode_and_deduplicate(df_hashes_auto_approved) - df_hashes_manually_approved = explode_and_deduplicate(df_hashes_manually_approved) - - # Clean sha256 values in all relevant DataFrames - df_approved_hashes = clean_sha256(df_approved_hashes) - df_hashes_auto_approved = clean_sha256(df_hashes_auto_approved) - df_hashes_manually_approved = clean_sha256(df_hashes_manually_approved) - - # Create sets for faster lookup - auto_approved_sha256 = set(df_hashes_auto_approved['sha256'].values) - manually_approved_sha256 = set(df_hashes_manually_approved['sha256'].values) - - # Debug: Print unmatched hashes - unmatched = set(df_approved_hashes['sha256']) - (auto_approved_sha256 | manually_approved_sha256) - print(f"Unmatched hashes: {unmatched}") - - # Process df_approved_paths - df_paths = df_approved_paths.assign(destination='Path Exclusion') - df_paths = df_paths[['sha256', 'description', 'destination', 'grouped_directory', 'filename']] - - # Process df_approved_hashes - df_hashes = df_approved_hashes.copy() - df_hashes['destination'] = df_hashes['sha256'].apply( - lambda x: 'Parent Policy Baseline' if x in auto_approved_sha256 - else ('Child Policy Allowlist' if x in manually_approved_sha256 else None) - ) - df_hashes = df_hashes.dropna(subset=['destination']) - df_hashes = df_hashes.assign(grouped_directory=None) - - # Use 'filename_x' only if it exists, otherwise fallback to 'filename' - filename_col = 'filename_x' if 'filename_x' in df_hashes.columns else 'filename' - selected_cols = ['sha256', 'description', 'destination', 'grouped_directory', filename_col] - df_hashes = df_hashes[selected_cols] - - # Concatenate results - df_hashdestination = pd.concat([df_paths, df_hashes], ignore_index=True) - return df_hashdestination - -def combineHashAndHist(path, first_policy, second_policy): - - condensed_combo = pd.read_parquet(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet") - df = pd.read_parquet(path) - - #Pull hash info for the entries in the needs approval table - df = pd.merge(condensed_combo, df, on='sha256', how='inner') - - #Rename Publisher, Keep and reorder columns we want - df = df.rename(columns={'publisher_x': 'publisher'}) - df = df[['sha256', 'publisher', 'description', 'filename', 'hostname', 'username', 'productname', 'productversion','reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount','reputation_status', 'reputation_threatlevel', 'reputation_threatname','reputation_timestamp', 'pprocess', 'gprocess', 'commandline']] - df = df.sort_values(by='filename') - - df.to_parquet(path, index=False) - del df - del condensed_combo - gc.collect() - -def combineHashes(url, first_policy, second_policy): - combined_hashes = pd.DataFrame(columns=['sha256', 'publisher']) - hashes = [] - try: - hash1 = pd.read_parquet(f"parquet\\execution_history_{first_policy}.parquet", columns=['sha256', 'publisher']) - pathf.inspect_parquet(f"parquet\\execution_history_{first_policy}.parquet") - if not hash1.empty: - hashes.append(hash1) - else: - print("āš ļø First dataframe is empty.") - except Exception as e: - print(f"āŒ Error reading first Parquet file: {e}") - - try: - hash2 = pd.read_parquet(f"parquet\\execution_history_{second_policy}.parquet", columns=['sha256', 'publisher']) - pathf.inspect_parquet(f"parquet\\execution_history_{second_policy}.parquet") - if not hash2.empty: - hashes.append(hash2) - else: - print("āš ļø Second dataframe is empty.") - except Exception as e: - print(f"āŒ Error reading second Parquet file: {e}") - - if hashes: - combined_hashes = pd.concat(hashes, ignore_index=True) - print(f"āœ… Combined {len(combined_hashes)} hashes.") - else: - print("āš ļø No valid dataframes to combine.") - - combined_hashes = combined_hashes.drop_duplicates(subset=['sha256']) - augmented_combo = hashf.augmentAggregatedHashes(url, combined_hashes) - - numeric_reputation_cols = [ - 'reputation_scannermatch', - 'reputation_scannercount', - 'reputation_threatlevel' - ] - - for col in numeric_reputation_cols: - if col in augmented_combo.columns: - augmented_combo[col] = pd.to_numeric(augmented_combo[col].replace('N/A', pd.NA), errors='coerce') - - augmented_combo = augmented_combo.rename(columns={'publisher_x': 'publisher'}) - augmented_combo = augmented_combo[['sha256', 'publisher', 'description', 'productname', 'productversion', - 'reputation_lastseen', 'reputation_scannermatch', 'reputation_scannercount', - 'reputation_status', 'reputation_threatlevel', 'reputation_threatname', - 'reputation_timestamp']] - augmented_combo = augmented_combo.sort_values(by=['publisher', 'description', 'productname']) - augmented_combo.to_parquet(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet", index=False) - - del combined_hashes - del augmented_combo - gc.collect() - print(ct.colorText("Hash reputation info added to dataframe", "green")) - -def condenseExecutions(first_policy,second_policy): - exe1 = pd.DataFrame() - exe2 = pd.DataFrame() - condensed_combo = pd.DataFrame() - - try: - exe1 = pd.read_parquet(f"parquet\\execution_history_{first_policy}.parquet") - pathf.inspect_parquet(f"parquet\\execution_history_{first_policy}.parquet") - if not exe1.empty: - print() - else: - print("āš ļø First dataframe is empty.") - except Exception as e: - print(f"āŒ Error reading first Parquet file: {e}") - - try: - exe2 = pd.read_parquet(f"parquet\\execution_history_{second_policy}.parquet") - pathf.inspect_parquet(f"parquet\\execution_history_{second_policy}.parquet") - if not exe2.empty: - print() - else: - print("āš ļø Second dataframe is empty.") - except Exception as e: - print(f"āŒ Error reading second Parquet file: {e}") - - if not exe1.empty and not exe2.empty: - condensed_combo = pd.concat([exe1, exe2], ignore_index=True) - - print(f"āœ… Combined {len(condensed_combo)} hashes.") - elif exe1.empty: - condensed_combo = exe2 - elif exe2.empty: - condensed_combo = exe1 - else: - print("āš ļø No valid dataframes to combine.") - - condensed_combo.to_parquet(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet", index=False) - del condensed_combo - gc.collect() - -def divideSortedHashExecutions(first_policy,second_policy, pups): - - combineHashAndHist(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet", first_policy, second_policy) - combineHashAndHist(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet", first_policy, second_policy) - combineHashAndHist(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet", first_policy, second_policy) - - unknown = pd.read_parquet(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") - good = pd.read_parquet(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") - bad = pd.read_parquet(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet") - - # Build regex pattern once - pattern = pathf.regulator(pups) - - # Move matching rows from unknown and good to bad - bad = pd.concat([ - bad, - unknown[unknown["filename"].str.contains(pattern, na=False)], - good[good["filename"].str.contains(pattern, na=False)] - ], ignore_index=True) - - # Remove matching rows from unknown and good - unknown = unknown[~unknown["filename"].str.contains(pattern, na=False)] - good = good[~good["filename"].str.contains(pattern, na=False)] - - unknown.to_csv(f"needs_approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv",index=False) - good.to_csv(f"needs_approved\\hashes_rep_good_{first_policy}_{second_policy}.csv",index=False) - bad.to_csv(f"needs_approved\\hashes_rep_bad_{first_policy}_{second_policy}.csv",index=False) - - ct.style_dataframe_dark(unknown, f"needs_approved\\hashes_rep_unknown_{first_policy}_{second_policy}.html") - ct.style_dataframe_dark(good, f"needs_approved\\hashes_rep_good_{first_policy}_{second_policy}.html") - ct.style_dataframe_dark(bad, f"needs_approved\\hashes_rep_bad_{first_policy}_{second_policy}.html") - -def generatePreflights(first_policy, second_policy): - allhashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet") - - pathexclusions = tryToReadCSV(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv") - pathexclusions.to_parquet(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet", index=False) - - allowbyhash = allhashes[~allhashes['sha256'].isin(pathexclusions['sha256'])] - - allowbyhash.to_parquet(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet", index=False) - - allowbyhash.sort_values(by=["filename"]) - - ct.style_dataframe_dark(allowbyhash, f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html") - ct.style_dataframe_dark(pathexclusions, f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html") - - del allowbyhash - del pathexclusions - gc.collect() \ No newline at end of file diff --git a/utils/pathfunctions.py b/utils/pathfunctions.py deleted file mode 100644 index e1ddfbb..0000000 --- a/utils/pathfunctions.py +++ /dev/null @@ -1,190 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import ast -import gc -import os -import pandas as pd -import re -import utils.pathfunctions as pathf -import utils.pretty as ct -from AirlockTools import tryToReadCSV - - -def split_filepaths_grouped(df, col="filename", group_parts=4, min_parts=4): - def clean_split(path): - parts = os.path.normpath(path).split(os.sep) - # Remove leading empty strings caused by UNC paths - parts = [p for p in parts if p] - return parts - - df = df.copy() - split_paths = df[col].apply(clean_split) - - # Filter out paths with fewer than `min_parts` components - df = df[split_paths.apply(lambda parts: len(parts) >= min_parts)].copy() - split_paths = split_paths[df.index] # Update split_paths to match filtered df - - df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:group_parts])) - grouped = df.groupby("group_key") - new_rows = [] - - for _, group_df in grouped: - paths = group_df[col].tolist() - split_parts = [clean_split(p) for p in paths] - - def longest_common_prefix(paths): - if not paths: - return [] - prefix = paths[0] - for path in paths[1:]: - prefix = [a for a, b in zip(prefix, path) if a == b] - if not prefix: - break - return prefix - - common_prefix = longest_common_prefix(split_parts) - prefix_str = os.sep.join(common_prefix) - - for i, parts in enumerate(split_parts): - filename = parts[-1] - middle = os.sep.join(parts[len(common_prefix):-1]) if len(parts) > len(common_prefix) + 1 else "" - row = group_df.iloc[i].copy() - row["longestcfp"] = prefix_str - row["middle"] = middle - row["filename_only"] = filename - new_rows.append(row) - - return pd.DataFrame(new_rows).drop(columns=["group_key"]) - -def mask_from_csv(df, csv_path, filepath_col): - """ - Reads reviewed CSV of groups, keeps only files in approved groups. - """ - review_df = pd.read_csv(csv_path) - - def parse_paths(val): - if isinstance(val, str): - try: - # Try to parse as a list - parsed = ast.literal_eval(val) - # If it's not a list, wrap it - return parsed if isinstance(parsed, list) else [parsed] - except (ValueError, SyntaxError): - # If parsing fails, treat it as a single path - return [val] - return [val] - - review_df[filepath_col] = review_df[filepath_col].apply(parse_paths) - - # Flatten all approved file paths into a set for masking - approved_files = set() - for paths in review_df[filepath_col]: - approved_files.update(paths) - - # Keep only rows in df that are in approved_files - masked_df = df[df[filepath_col].isin(approved_files)].copy() - remainder = df[~df[filepath_col].isin(approved_files)].copy() - return remainder - -def filter_and_drop(approved, eligiblepaths, min_hashes): - """ - Filters eligiblepaths to rows where all hashes are in approved, - then drops rows with fewer than min_hashes hashes. - """ - approved_hashes = set(approved['sha256']) - - def all_hashes_approved(row): - return all(h in approved_hashes for h in row['sha256']) - - filtered = eligiblepaths[eligiblepaths.apply(all_hashes_approved, axis=1)] - filtered = filtered[filtered['sha256'].apply(len) >= min_hashes] - - return filtered - -def inspect_parquet(path): - try: - df = pd.read_parquet(path) - print(f"āœ… Successfully read: {path}") - print(f"šŸ“„ Columns: {df.columns.tolist()}") - print(f"šŸ”¢ Rows: {len(df)}") - return df - except Exception as e: - print(f"āŒ Error reading {path}: {e}") - return pd.DataFrame() - - -def regulator(paths, case_insensitive=True): - """ - Build a regex pattern that matches any of the given Windows path fragments. - """ - escaped = [re.escape(p) for p in paths] - pattern = "(?:" + "|".join(escaped) + ")" - if case_insensitive: - pattern = "(?i)" + pattern # Add inline case-insensitive flag - print(f"Regulator is providing: {pattern}") - return pattern - -def generatePathReview(first_policy, second_policy, badpathparts, min_files_for_path): - - if not os.path.exists(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet"): - - df1 = tryToReadCSV(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv") - df2 = tryToReadCSV(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv") - - all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename']) - - - - print(ct.colorText(f"Approved hash lists have been combined","green")) - - all_approved_hashes.to_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet", index=False) - del all_approved_hashes - gc.collect() - - if not os.path.exists(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet"): - all_approved_hashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet") - print(ct.colorText(f"Beginning calculating longest common filepaths for path exceptions","green")) - - haslcp = pathf.split_filepaths_grouped(all_approved_hashes) - haslcp.drop_duplicates() - - forbidden = pathf.regulator(badpathparts, True) - forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False) - - - print(ct.colorText("Removing forbidden filepaths for path exceptions", "green")) - - # Make a real DataFrame copy before modifying - lcp_not_forbidden = haslcp[~forbidden_lcfp].copy() - - #For the review, drop down to only the columns we care, and then group by the commmon file path, consolidating and dropping dupes - lcp_not_forbidden_review = lcp_not_forbidden[['longestcfp', 'middle', 'filename_only', 'sha256']] - - # Count unique sha256 per longestcfp - unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index() - unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count'] - - # Merge the count back into the original DataFrame - lcp_not_forbidden_review = lcp_not_forbidden_review.merge(unique_sha_counts, on='longestcfp', how='left') - lcp_not_forbidden_review = lcp_not_forbidden_review[lcp_not_forbidden_review['unique_sha256_count'] >= min_files_for_path] - - lcp_not_forbidden_review.to_parquet(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet",index=False) - lcp_not_forbidden_review.to_csv(f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.csv",index=False) - ct.style_dataframe_dark(lcp_not_forbidden_review,f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.html", True) - - del lcp_not_forbidden - del unique_sha_counts - del lcp_not_forbidden_review - diff --git a/utils/policyfunctions.py b/utils/policyfunctions.py deleted file mode 100644 index ccacdec..0000000 --- a/utils/policyfunctions.py +++ /dev/null @@ -1,123 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import gc -import json -import os -import pandas as pd -import re -import requests -import utils.pretty as ct -import utils.allowlist - - - - -def addHash(policy, hash): - print(f"Adding the following hashes to {policy}:") - for p in hash: - print(p) - - -def addPath(policy, hash): - print(f"Adding the following Path Exclusions to {policy}:") - for p in hash: - print(p) - -def addHashReal(url, allowlistID, hashlist): - endpoint = url + '/v1/hash/application/add' - print(ct.colorText("[+] Grabbing All Categories", "cyan")) - payload = { - "applicationid" : allowlistID, - "hashes" : hashlist - } - headers = { - "X-APIKey": os.getenv('APIKEY') - } - payload = json.dumps(payload) - response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) - response.raise_for_status() # Raise an error for bad status codes - parse_text = json.loads(response.text) - print(parse_text) - - -def addPathReal(url, grouplistID, pathlist): - endpoint = url + '/v1/group/path/add' - print(ct.colorText("[+] Grabbing All Categories", "cyan")) - payload = { - "groupid" : grouplistID, - "path" : pathlist - } - headers = { - "X-APIKey": os.getenv('APIKEY') - } - print(payload) - payload = json.dumps(payload) - response = requests.request("POST", endpoint, headers=headers, data=payload, verify=False) - print(response.text) - -def getPolicyInfo(url, policy, days): - executionhist_policy = pd.DataFrame() - exehist = utils.allowlist.pullPolicyExechistories(url, policy, days, True) - data = json.loads(exehist) - executionhist_policy = pd.DataFrame(data["response"]["exechistories"]) - if not executionhist_policy.empty: - executionhist_policyxecutionhist_policy = executionhist_policy[['sha256', 'publisher', 'filename', 'hostname', 'username', 'pprocess', 'gprocess', 'commandline']] - executionhist_policy = executionhist_policy.drop_duplicates(subset=['sha256', 'filename', 'hostname']) - executionhist_policy = executionhist_policy.sort_values(by=['sha256', 'filename']) - executionhist_policy.to_parquet(f"parquet\\execution_history_{policy}.parquet", index=False) - print(ct.colorText(f"Staging of Execution history for policy: {policy} is complete", "green")) - del data - del exehist - gc.collect() - return executionhist_policy - -def sendToPolicy(url, first_policy, second_policy, destination_name, destination_id, allowlist_parent_name, allowlist_parent_id, allowlist_child_name, allowlist_child_id): - pathexclusions = pd.read_parquet(f"parquet\\final_path_exclusions_{first_policy}_{second_policy}.parquet") - allowbyhash = pd.read_parquet(f"parquet\\final_hash_approvals_{first_policy}_{second_policy}.parquet") - - ct.areYouSure() - confirmation = input(ct.colorText("Type 'I AGREE' to continue: ","white")) - - if confirmation.strip().upper() == "I AGREE": - print(ct.colorText("Proceeding with the code...", "yellow")) - print(ct.colorText(f"Adding path exclusions to {destination_name}", "yellow")) - pathexcludelist = pathexclusions['longestcfp'].unique().tolist() - - # Regex to match a Windows drive letter at the start (e.g., C:\) - drive_letter_pattern = re.compile(r'^[a-zA-Z]:\\') - - # Processed list - processed_paths = [ - (path if drive_letter_pattern.match(path) else f"\\\\{path}") + "**" - for path in pathexcludelist -] - addPath(url, destination_id,processed_paths) - - print(ct.colorText(f"Adding hashes to {allowlist_parent_name}", "yellow")) - - allowlist_parenthashlist = allowbyhash[allowbyhash['reputation_status'] == 'KNOWN']['sha256'].unique().tolist() - addHash(url, allowlist_parent_id,allowlist_parenthashlist) - - print(ct.colorText(f"Adding hashes to {allowlist_child_name}", "yellow")) - allowlist_childhashlist = allowbyhash[allowbyhash['reputation_status'] == 'UNKNOWN']['sha256'].unique().tolist() - addHash(url, allowlist_child_id, allowlist_childhashlist) - - ct.locked() - - exit() - - else: - print(ct.colorText("Operation aborted. You MUST EXPLICITLY AGREE to proceed.", "red")) - \ No newline at end of file diff --git a/utils/selector.py b/utils/selector.py new file mode 100644 index 0000000..714e197 --- /dev/null +++ b/utils/selector.py @@ -0,0 +1,334 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +import logging +from typing import Any, Callable, List, Optional, Union + +import pandas as pd + +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]: + return sorted(items, key=lambda item: label_func(item).lower()) + + @staticmethod + def _display_choices( + items: List[Any], + label_func: Callable[[Any], str], + num_columns: int = 4, + 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 + + rows = (len(items) + num_columns - 1) // num_columns + print(f"\n{header}") + for row in range(rows): + line = "" + for col in range(num_columns): + idx = row + col * rows + if idx < len(items): + label = label_func(items[idx]) + line += f"{idx + 1}: {label:<30}" + print(line) + + @staticmethod + def _display_selected_items( + selected: List[Any], + label_func: Callable[[Any], str], + num_columns: int = 4 + ) -> None: + print(colorText("\nCurrent selections:", "cyan")) + if not selected: + print(" (none)") + return + sorted_selected = sorted(selected, key=lambda item: label_func(item).lower()) + rows = (len(sorted_selected) + num_columns - 1) // num_columns + for row in range(rows): + line = "" + for col in range(num_columns): + idx = row + col * rows + if idx < len(sorted_selected): + label = label_func(sorted_selected[idx]) + line += f"{label:<30}" + print(line) + + @staticmethod + def _parse_selection_input(input_str: str, max_index: int) -> List[int]: + selections = [] + for part in input_str.split(","): + part = part.strip() + if "-" in part: + try: + start, end = map(int, part.split("-")) + selections.extend(range(start, end + 1)) + except ValueError: + continue + elif part.isdigit(): + selections.append(int(part)) + return [i for i in selections if 1 <= i <= max_index] + + @staticmethod + def _select_from_list( + items: List[Any], + label_func: Callable[[Any], str], + allow_multiple: bool = False, + prompt_each: bool = False, + header: str = "Available Choices:", + num_columns: int = 4 + ) -> Union[Optional[Any], List[Any]]: + if not items: + logger.warning("No items available for selection.") + return None + + full_sorted_items = Selector._get_sorted_items(items, label_func) + remaining_items = full_sorted_items.copy() + selected = [] + + 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() + if choice == "q": + break + elif choice == "r": + selected.clear() + remaining_items = full_sorted_items.copy() + print(colorText("šŸ”„ Selections reset.", "yellow")) + continue + indices = Selector._parse_selection_input(choice, len(remaining_items)) + newly_selected = [] + for index in indices: + item = remaining_items[index - 1] + if item not in selected: + selected.append(item) + newly_selected.append(item) + if prompt_each: + 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] + return selected if selected else None + else: + 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): + selected_item = full_sorted_items[choice - 1] + logger.info(f"Selected: {label_func(selected_item)}") + return selected_item + else: + logger.warning("Selection out of range.") + except ValueError: + logger.warning("Invalid input.") + return None + + @staticmethod + def select_with_mode( + items: List[Any], + label_func: Callable[[Any], str], + header: str = "Available Choices:" + ) -> List[Any]: + 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 + selected = Selector._select_from_list( + items, + label_func=label_func, + allow_multiple=True, + prompt_each=False, + header=header + ) + if not selected: + return items + if mode == "i": + print(colorText(f"āœ… Included {len(selected)} item(s).", "green")) + return selected + elif mode == "e": + print(colorText(f"šŸš« Excluded {len(selected)} item(s).", "yellow")) + return [item for item in items if item not in selected] + else: + print(colorText("āš ļø Invalid mode. Returning all items.", "yellow")) + return items + + @staticmethod + def select_objects( + 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:" + ) + + @staticmethod + def select_string( + 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:" + ) + + @staticmethod + def select_int( + 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:" + ) + + @staticmethod + def select_value( + prompt: str, + value_type: type = int, + valid_range: Optional[tuple] = None, + allow_quit: bool = False + ) -> Optional[Any]: + while True: + user_input = get_sanitized_input(prompt).strip().lower() + if allow_quit and user_input == "q": + logger.info("User opted to quit value selection.") + return None + try: + value = value_type(user_input) + if valid_range: + min_val, max_val = valid_range + if not (min_val <= value <= max_val): + logger.warning(f"Value out of range ({min_val}ā€“{max_val}).") + continue + logger.info(f"User selected value: {value}") + return value + except ValueError: + logger.warning(f"Invalid input. Expected a {value_type.__name__}.") + + @staticmethod + def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool: + while True: + response = get_sanitized_input(prompt).strip().lower() + if response in ["y", "yes"]: + logger.info("User confirmed action.") + return True + elif response in ["n", "no"]: + logger.info("User declined action.") + return False + else: + logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.") + + @staticmethod + def select_dataframe_rows( + df: pd.DataFrame, + columns: Optional[List[str]] = None, + allow_multiple: bool = False, + prompt_each: bool = False, + header: str = "Available Rows:" + ) -> List[pd.Series]: + if df.empty: + print("DataFrame is empty.") + return [] + + if columns: + df = df[columns] + + items = [row for _, row in df.iterrows()] + label_func = lambda row: str(row.to_dict()) + + result = Selector._select_from_list( + items, + label_func=label_func, + allow_multiple=allow_multiple, + prompt_each=prompt_each, + header=header + ) + + if isinstance(result, pd.Series): + return [result] + elif isinstance(result, list): + return result + else: + return [] + + @staticmethod + def select_dataframe_with_mode( + df: pd.DataFrame, + columns: Optional[List[str]] = None, + header: str = "Available Rows:" + ) -> List[pd.Series]: + if df.empty: + print("āš ļø DataFrame is empty.") + return [] + + # Filter columns if specified + if columns: + df = df[columns] + + items = df.to_dict("records") + label_func = lambda row: " | ".join(str(row[col]) for col in df.columns) + + # Show rows first + print(colorText(header, "cyan")) + for i, row in enumerate(items): + 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")) + mode = get_sanitized_input("").strip().lower() + + if mode == "a": + return [pd.Series(row) for row in items] + + # Prompt for selection only once + selected = Selector._select_from_list( + items, + label_func=label_func, + allow_multiple=True, + prompt_each=False, + header=header + ) + + if not selected: + return [pd.Series(row) for row in items] + + if mode == "i": + print(colorText(f"āœ… Included {len(selected)} row(s).", "green")) + return [pd.Series(row) for row in selected] + elif mode == "e": + print(colorText(f"šŸš« Excluded {len(selected)} row(s).", "yellow")) + 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 diff --git a/utils/setup.py b/utils/setup.py new file mode 100644 index 0000000..f7ba7f2 --- /dev/null +++ b/utils/setup.py @@ -0,0 +1,203 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import json +import logging +import logging.config +import logging.handlers +import os +from pathlib import Path +import platform +import sys + +from dotenv import load_dotenv, set_key + +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" + else: + return home / '.local' / 'share' / "AirlockTools" + + +def configure_logging(log_dir: Path, log_level: str = "DEBUG"): + log_file = log_dir / "airlocktools.log" + + config = { + "version": 1, # Required key for dictConfig format version + "disable_existing_loggers": False, # Keeps existing loggers active + "formatters": { + "detailed": { + "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + # Includes timestamp, logger name, level, and message + }, + "simple": { + "format": "%(levelname)s - %(message)s" + # Minimal format for console output + }, + }, + "handlers": { + "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 + }, + "console": { + "class": "logging.StreamHandler", + "level": log_level.upper(), # Configurable log level + "formatter": "simple", # Use simple format + }, + }, + "root": { + "level": "DEBUG", # Root logger level + "handlers": ["file", "console"], # Attach both handlers + }, + } + + # Add Windows Event Log handler if on Windows + if platform.system() == "Windows": + 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 + } + config["root"]["handlers"].append("eventlog") + except Exception as e: + logging.warning(f"Could not attach Windows Event Log handler: {e}") + + # Apply the logging configuration + 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__)))) + return base_path.parent / "system_config.json" + + +def load_system_config() -> dict: + try: + config_path = get_system_config_path() + with open(config_path, "r") as f: + return json.load(f) + except FileNotFoundError: + logging.warning("āš ļø system_config.json not found. Using built-in defaults.") + return { + "APPNAME": "AirlockTools", + "LOG_LEVEL": "DEBUG", + "PATH_EXCLUSION_CONST": 4, + "MIN_FILES_FOR_PATH": 4, + "VT_THREAT_TOLERANCE": 4, + "POLICY_MAP_ENF_AUD": { + "enforced_id": "audit_id" + } + } + +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" + } + 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) + 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', + } + + for name, path in dirs.items(): + path.mkdir(parents=True, exist_ok=True) + logging.debug(f"{name.capitalize()} directory ensured at: {path}") + + system_config = load_system_config() + configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG")) + + env_path = base_dir / ".env" + if not env_path.exists(): + env_path.touch() + load_dotenv(dotenv_path=env_path, override=True) + + working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data")) + working_dir.mkdir(parents=True, exist_ok=True) + set_key(env_path, "WORKING_DIR", str(working_dir)) + os.environ["WORKING_DIR"] = str(working_dir) + logging.debug(f"Working directory set to: {working_dir}") + + folders_structure = { + "Approved": [], + "Needs_Review": ["Review_First", "Review_Second", "HTML"], + "Preflight": ["HTML"], + "Archived": [] + } + + for folder_name, subfolders in folders_structure.items(): + folder_path = working_dir / folder_name + folder_path.mkdir(parents=True, exist_ok=True) + logging.debug(f"'{folder_name}' folder ensured at: {folder_path}") + for subfolder in subfolders: + subfolder_path = folder_path / subfolder + subfolder_path.mkdir(parents=True, exist_ok=True) + logging.debug(f" ā””ā”€ '{subfolder}' subfolder created at: {subfolder_path}") + + user_config = load_user_config(dirs['config']) + merged_config = {**system_config, **user_config} + + protected_config = load_protected_config() + merged_config.update(protected_config) + + # āœ… URL resolution order: system_config ā†’ .env ā†’ user prompt + url = system_config.get("URL") + if not url: + url = os.getenv("URL") + if not url: + url = input("šŸŒ Enter the service URL (e.g., https://example.com/api): ").strip() + merged_config["URL"] = url + set_key(env_path, "URL", url) + os.environ["URL"] = url + logging.debug(f"Service URL set to: {url}") + + write_config_to_env(merged_config, env_path) \ No newline at end of file diff --git a/utils/tui.py b/utils/tui.py new file mode 100644 index 0000000..3f1a64b --- /dev/null +++ b/utils/tui.py @@ -0,0 +1,468 @@ +import logging +import os +import sys + +import dotenv +from dotenv import set_key +from textual.app import App, ComposeResult +from textual.containers import Horizontal, Vertical +from textual.reactive import reactive +from textual.screen import Screen +from textual.widgets import ( + Button, + DirectoryTree, + Footer, + Header, + Static, + Tab, + Tabs, +) + +from flows.otp import otp_activities_by_agent, otp_generate, otp_revoke +from flows.prepPolicy import menu_policy_enforce +from flows.quietAgent import findQuietAgents +from services.agenthandler import findAgents, moveAgents, toggleEnforcement +from services.API import AirlockAPIWrapper +from services.policyhandler import confirmUpdateAfromE +from utils.configmanager import load_env +from utils.setup import get_base_directory, load_user_config +from utils.utils import open_directory # we will actually use this now + +dotenv.load_dotenv() + +# --------------------------------------------------------------------------- +# GLOBAL STASH +# --------------------------------------------------------------------------- +# ("legacy", func, args, kwargs) OR ("restart",) +_PENDING_JOB = None + +logger = logging.getLogger(__name__) + + +ASCII_ART = r""" + _____ .__ .__ __ ___________ .__ + / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ + / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ +/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ +\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > + \/ \/ \/ \/ +""" + + +# --------------------------------------------------------------------------- +# helper to persist TEXTUAL_THEME to *user* config and mirror to .env +# --------------------------------------------------------------------------- +def _persist_user_theme(theme_name: str) -> None: + """ + Store the chosen Textual theme in the user's config: + /config/user_config.json + and also mirror to /.env so load_env(...) sees it. + """ + base_dir = get_base_directory() + config_dir = base_dir / "config" + user_config_path = config_dir / "user_config.json" + env_path = base_dir / ".env" + + # ensure dirs / files exist similarly to setup() + 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") + + # load existing user config + user_conf = load_user_config(config_dir) + user_conf["TEXTUAL_THEME"] = theme_name + + # write it back + user_config_path.write_text( + # pretty print so it stays human-readable + __import__("json").dumps(user_conf, indent=4), + encoding="utf-8", + ) + logger.debug("Updated user_config.json with TEXTUAL_THEME=%s", theme_name) + + # mirror to .env (like setup.write_config_to_env does) + env_path.parent.mkdir(parents=True, exist_ok=True) + if not env_path.exists(): + env_path.touch() + try: + set_key(str(env_path), "TEXTUAL_THEME", theme_name) + except Exception as exc: # keep going even if .env write fails + logger.warning("Failed to mirror TEXTUAL_THEME to .env: %s", exc) + + # reload so load_env(...) sees the new value right now + dotenv.load_dotenv(dotenv_path=env_path, override=True) + logger.debug("Reloaded .env from %s", env_path) + + +# --------------------------------------------------------------------------- +# 1) SCREEN +# --------------------------------------------------------------------------- +class MainMenuScreen(Screen): + current_tab = reactive("") + + BUTTON_DEFS = { + "find": [ + ("šŸ” - Device Search", "find_device_button"), + ("šŸ”‡ - Find Quiet Hosts", "find_quiet_button"), + ], + "move": [ + ("āœ… - Move to local approval", "move_local_button"), + ("šŸ”„ - Move to Audit/Enforcement", "move_audit_button"), + ("šŸ”€ - Move - Other", "move_other_button"), + ], + "otp": [ + ("šŸ” - Generate OTPs", "otp_generate_button"), + ("šŸ“Š - OTP Activities By Agent", "otp_activities_button"), + ("āŒ - Revoke OTPs", "otp_revoke_button"), + ], + "policy": [ + ("šŸ”’ - Prepare Policy For Enforcement", "policy_prep_button"), + ("šŸ”„ - Update Audit Policies", "policy_audit_update_button"), + ], + } + + # textual themes to expose + THEME_BUTTONS = [ + ("textual-dark", "textual-dark"), + ("textual-light", "textual-light"), + ("nord", "nord"), + ("gruvbox", "gruvbox"), + ("catppuccin-mocha", "catppuccin-mocha"), + ("dracula", "dracula"), + ("tokyo-night", "tokyo-night"), + ("monokai", "monokai"), + ("flexoki", "flexoki"), + ("catppuccin-latte", "catppuccin-latte"), + ("solarized-light", "solarized-light"), + ] + + def __init__(self) -> None: + super().__init__() + self.extras = load_env("EXTRAS") + wd = load_env("WORKING_DIR") or os.getcwd() + if not os.path.isdir(wd): + wd = os.getcwd() + self.working_dir = wd + + + def _make_buttons_for(self, tab_id: str) -> Vertical: + defs = self.BUTTON_DEFS.get(tab_id, []) + buttons = [] + for label, btn_id in defs: + btn = Button(label, id=btn_id) + btn.styles.width = "100%" # Make button span full width of parent + buttons.append(btn) + return Vertical(*buttons) + + + + + def compose(self) -> ComposeResult: + yield Header() + yield Static(ASCII_ART, id="logo") + + tabs = [ + Tab("Device Search", id="find"), + Tab("Move Agent", id="move"), + Tab("OTP", id="otp"), + Tab("Directory", id="dir"), + Tab("Settings", id="settings"), + ] + + if self.extras == "POLICYPREP": + tabs.insert(3, Tab("Policy Prep", id="policy")) + + yield Tabs(*tabs, id="tabs") + yield Vertical(id="content") + yield Footer() + + def on_mount(self) -> None: + self.switch_tab("find") + + # focus helpers + def _get_content_buttons(self) -> list[Button]: + content = self.query_one("#content", Vertical) + return list(content.query(Button)) + + def _focus_first_button(self) -> None: + buttons = self._get_content_buttons() + if buttons: + buttons[0].focus() + + def _focus_tabs(self) -> None: + tabs = self.query_one("#tabs", Tabs) + tabs.focus() + + def _focus_nearby_button(self, direction: int) -> None: + buttons = self._get_content_buttons() + if not buttons: + return + + try: + current = next(i for i, b in enumerate(buttons) if b.has_focus) + except StopIteration: + if direction > 0: + buttons[0].focus() + else: + buttons[-1].focus() + return + + if direction < 0 and current == 0: + self._focus_tabs() + return + + 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) + content.remove_children() + + if tab_id in self.BUTTON_DEFS: + content.mount(self._make_buttons_for(tab_id)) + self.call_later(self._focus_first_button) + elif tab_id == "dir": + content.mount(DirectoryTree(self.working_dir, id="dir_tree")) + elif tab_id == "settings": + # Create and mount the horizontal container + horizontal_container = Horizontal(id="settings_grid") + horizontal_container.styles.layout = "horizontal" + horizontal_container.styles.height = "auto" + content.mount(Static("Theme Options")) + content.mount(horizontal_container) # Mount the horizontal container first + + # Create 3 columns + for i in range(1): + column = Vertical() + column.styles.width = "1fr" + column.styles.height = "auto" + horizontal_container.mount(column) # Mount each column + + for j in range(i, len(self.THEME_BUTTONS), 1): + 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%" + column.mount(button) # Mount each button + + else: + content.mount(Static(f"Unknown tab: {tab_id}")) + + def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: + self.switch_tab(event.tab.id) + + def on_key(self, event) -> None: + key = event.key + logger.debug("KEY: %r", key) + + if self.current_tab == "dir": + return + + if key in ("down", "j"): + self._focus_nearby_button(+1) + event.stop() + return + if key in ("up", "k"): + self._focus_nearby_button(-1) + event.stop() + return + + if key in ("left", "right"): + tabs = self.query_one("#tabs", Tabs) + if not tabs.has_focus: + tabs.focus() + if key == "left": + tabs.action_previous_tab() + else: + tabs.action_next_tab() + event.stop() + return + return + + def on_directory_tree_file_selected(self, event: DirectoryTree.FileSelected) -> None: + path = event.path + logger.debug("Directory file selected: %s", path) + try: + open_directory(str(path)) + except Exception as exc: + logger.error("Failed to open %s: %s", path, exc) + self.app.bell() + + def on_button_pressed(self, event: Button.Pressed) -> None: + global _PENDING_JOB + button_id = event.button.id + logger.debug("Button pressed: %s", button_id) + + # theme selection ā†’ user config + if button_id.startswith("set_theme_"): + theme_name = button_id.replace("set_theme_", "") + _persist_user_theme(theme_name) + _PENDING_JOB = ("restart",) + self.app.exit() + return + + match button_id: + case "find_device_button": + _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) + case "find_quiet_button": + _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) + case "move_local_button": + _PENDING_JOB = ( + "legacy", + print, + ("Move to local approval (placeholder)",), + {}, + ) + case "move_audit_button": + _PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {}) + case "move_other_button": + _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {}) + case "otp_generate_button": + _PENDING_JOB = ("legacy", otp_generate, (self.app.api,), {}) + case "otp_activities_button": + _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) + case "otp_revoke_button": + _PENDING_JOB = ("legacy", otp_revoke, (self.app.api,), {}) + case "policy_prep_button": + _PENDING_JOB = ("legacy", menu_policy_enforce, (self.app.api,), {}) + case "policy_audit_update_button": + _PENDING_JOB = ("legacy", confirmUpdateAfromE, (self.app.api,), {}) + case _: + self.app.bell() + logger.warning("Unknown button pressed: %s", button_id) + return + + logger.debug("Set _PENDING_JOB = %r", _PENDING_JOB) + self.app.exit() + + + + + + +# --------------------------------------------------------------------------- +# 2) APP +# --------------------------------------------------------------------------- +class AirlockTools(App): + CSS = """ + #logo { + width: 100%; + content-align: center middle; + text-align: center; + } + """ + + BINDINGS = [ + ("q", "quit", "Quit"), + ("d", "open_dir", "Open Directory"), + ] + + def __init__(self, api: AirlockAPIWrapper): + self._textual_theme = load_env("TEXTUAL_THEME") or "textual-dark" + super().__init__() + self.api = api + wd = load_env("WORKING_DIR") or os.getcwd() + if not os.path.isdir(wd): + wd = os.getcwd() + self.working_dir = wd + + def on_mount(self) -> None: + self.theme = self._textual_theme + self.push_screen(MainMenuScreen()) + + def action_quit(self) -> None: + global _PENDING_JOB + _PENDING_JOB = None + self.exit() + + def action_open_dir(self) -> None: + screen = self.screen_stack[-1] + if isinstance(screen, MainMenuScreen): + if screen.current_tab != "dir": + screen.switch_tab("dir") + + + +# --------------------------------------------------------------------------- +# 3) TERMINAL + LEGACY +# --------------------------------------------------------------------------- +def _restore_terminal_for_legacy() -> None: + sys.stdout.write("\033[?1049l") + sys.stdout.write("\033[?25h") + sys.stdout.write("\033[0m") + sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l") + sys.stdout.write("\033[2J\033[H") + sys.stdout.flush() + + if os.name == "nt": + try: + import ctypes + kernel32 = ctypes.windll.kernel32 + handle = kernel32.GetStdHandle(-11) + mode = ctypes.c_ulong() + if kernel32.GetConsoleMode(handle, ctypes.byref(mode)): + kernel32.SetConsoleMode(handle, mode.value | 0x0004) + except Exception as exc: + logger.debug("VT enable on Windows failed: %s", exc) + + +def _run_legacy_job(func, args, kwargs) -> None: + logger.debug("Running legacy job: %s", getattr(func, "__name__", func)) + _restore_terminal_for_legacy() + + try: + func(*args, **kwargs) + finally: + try: + input("\nPress Enter to return to the UI...") + except EOFError: + pass + + +# --------------------------------------------------------------------------- +# 4) PUBLIC ENTRYPOINT +# --------------------------------------------------------------------------- +def run_AirlockTools(api: AirlockAPIWrapper) -> None: + global _PENDING_JOB + + while True: + base_dir = get_base_directory() + env_path = base_dir / ".env" + dotenv.load_dotenv(dotenv_path=env_path, override=True) + + _PENDING_JOB = None + app = AirlockTools(api) + + try: + app.run() + except SystemExit as exc: + logger.debug("Caught SystemExit from Textual: %s", exc) + + job = _PENDING_JOB + logger.debug("After app.run(), _PENDING_JOB = %r", job) + + if not job: + break + + if job[0] == "legacy": + _, func, args, kwargs = job + _run_legacy_job(func, args, kwargs) + continue + + if job[0] == "restart": + # just loop again; fresh .env was already loaded at the top + continue + + break + + +# --------------------------------------------------------------------------- +# 5) DEV +# --------------------------------------------------------------------------- +if __name__ == "__main__": + api = AirlockAPIWrapper() + run_AirlockTools(api) diff --git a/utils/pretty.py b/utils/utils.py similarity index 54% rename from utils/pretty.py rename to utils/utils.py index aaba329..539a588 100644 --- a/utils/pretty.py +++ b/utils/utils.py @@ -12,9 +12,377 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -import os -def colorText(text: str, color: str) -> str: + +import logging +import os +import platform +import re +import subprocess +import tempfile +import tkinter as tk +from tkinter import filedialog + +import pandas as pd + +logger = logging.getLogger(__name__) + + + + +def import_to_dataframe(file_path: str) -> pd.DataFrame: + df = pd.DataFrame() + + try: + if not os.path.exists(file_path): + print(colorText(f"Error: File '{file_path}' does not exist.", "red")) + return df + + ext = os.path.splitext(file_path)[1].lower() + + if ext == ".csv": + df = pd.read_csv(file_path) + elif ext == ".parquet": + df = pd.read_parquet(file_path) + else: + print(colorText(f"Error: Unsupported file extension '{ext}'.", "red")) + return df + + if df.empty: + print(colorText("Error: File has headers but no data rows.", "red")) + else: + print(colorText(f"Data loaded successfully from {file_path}", "green")) + + return df + + except pd.errors.EmptyDataError: + print( + colorText( + "Notice: CSV file is completely empty, falling back to empty frame", + "white", + ) + ) + return pd.DataFrame() + + except Exception as e: + print(colorText(f"Error reading file: {e}", "red")) + return pd.DataFrame() + + +def choose_directory(): + root = tk.Tk() + root.withdraw() # Hide the main window + directory = filedialog.askdirectory(title="Select a Directory") + print("Selected directory:", directory) + return directory + + +def choose_file(initial_directory=None, required_substring=None): + """Open a file dialog and ensure the selected file contains a required substring.""" + while True: + root = tk.Tk() + root.withdraw() # Hide the main window + file_path = filedialog.askopenfilename(initialdir=initial_directory) + + if not file_path: + print("No file selected.") + return None + + if required_substring and required_substring not in file_path: + print( + f"The selected file must contain '{required_substring}' in its path or name. Please try again." + ) + else: + 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()): + return user_input + else: + print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.") + + +def regulator(paths, case_insensitive=True): + """ + Build a regex pattern that matches any of the given Windows path fragments. + """ + escaped = [re.escape(p) for p in paths] + pattern = "(?:" + "|".join(escaped) + ")" + if case_insensitive: + pattern = "(?i)" + pattern # Add inline case-insensitive flag + print(f"Regulator is providing: {pattern}") + return pattern +def irtang(): + print( + colorText( + r""" + ā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆ ā–‘ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–“ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–‘ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–“ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–’ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–’ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆ +""", + "yellow", + ) + ) +def displayIntro(): + + print( + colorText( + r""" + _____ .__ .__ __ ___________ .__ + / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ + / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ +/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ +\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > + \/ \/ \/ \/ +""", + "cyan", + ) + ) +def welcome(): + print( + colorText( + "=================================================================================", + "cyan", + ) + ) + print( + colorText( + "======================== Welcome to the Airlock API Tool ========================", + "cyan", + ) + ) + print( + colorText( + "=================================================================================", + "cyan", + ) + ) + +def section_header(title): + print(colorText("\n --------------------------------------------------------------------", "cyan")) + print(colorText(f" ------------- {title} -------------", "cyan")) + print(colorText(" --------------------------------------------------------------------", "cyan")) + + + +def areYouSure(): + print( + colorText( + "šŸ›‘****************************************************************************************************************************************šŸ›‘", + "red", + ) + ) + print( + colorText( + "āš ļø=========================================================================================================================================āš ļø", + "yellow", + ) + ) + print( + colorText( + "šŸ›‘========================================================================================================================================šŸ›‘", + "red", + ) + ) + print( + colorText( + "āš ļø-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------āš ļø", + "yellow", + ) + ) + print( + colorText( + "šŸ›‘========================================================================================================================================šŸ›‘", + "red", + ) + ) + print( + colorText( + "āš ļø=========================================================================================================================================āš ļø", + "yellow", + ) + ) + print( + colorText( + "šŸ›‘****************************************************************************************************************************************šŸ›‘", + "red", + ) + ) + + +def locked(): + print( + colorText( + r""" + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆ ā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–’ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆ ā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–’ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–’ + ā–’ā–ˆā–ˆā–ˆā–ˆ + ā–’ā–ˆā–ˆā–ˆā–ˆ + ā–“ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ + ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–‘ +""", + "yellow", + ) + ) + + +def printDeviceEnforceChecklist(): + print( + colorText( + "\n --------------------------------------------------------------------", + "cyan", + ) + ) + print( + colorText( + " ------------- šŸ› ļø šŸ”’ Prepare to Enforce Policy šŸ› ļø šŸ”’ ------------------", + "cyan", + ) + ) + print( + colorText( + " --------------------------------------------------------------------", + "cyan", + ) + ) + print( + colorText( + "\nSequentually follow these steps to prepare a policy for enforcement:", + "white", + ) + ) + + print( + colorText( + "\n1. Choose which originating policy or policies to move to enforcement", + "cyan", + ) + ) + print( + colorText( + "2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", + "cyan", + ) + ) + print(colorText("3. Manually review the files:", "cyan")) + print( + colorText( + " 'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'", + "cyan", + ) + ) + print( + colorText( + " Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", + "cyan", + ) + ) + print( + colorText( + " If metarules need to be created, please make note of them, and remove the row from the csv.", + "cyan", + ) + ) + print( + colorText( + " When complete, save both csv files to the directory 'approved' and choose this option.", + "cyan", + ) + ) + print( + colorText( + " This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", + "cyan", + ) + ) + print( + colorText( + "4. Manually review the file 'needs_approved\\paths_needing_review.csv'", + "cyan", + ) + ) + print( + colorText( + " Remove the rows containing path exclusions you do not approve of", + "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", + "cyan", + ) + ) + print(colorText(" Preflight Lists will be generated", "cyan")) + + print(colorText("5. Choose the destination policy and parent and child allow list", "cyan")) + + print(colorText("6. Test ------------------------------------------------------", "cyan")) + print(colorText(" Print rather than apply selected data.", "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( + "R. Remove/Reset Generated data - will prompt to allow keeping execution history", + "cyan", + ) + ) + + print(colorText("B. Back", "cyan")) + + +def colorText(text, color): colors = { "red": "\033[91m", "green": "\033[92m", @@ -23,17 +391,17 @@ def colorText(text: str, color: str) -> str: "magenta": "\033[95m", "cyan": "\033[96m", "white": "\033[97m", - "reset": "\033[0m" + "reset": "\033[0m", } - return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}" -def style_dataframe_dark(df, output_html_path=None, overwrite=True): + +def formatHTML(df, output_html_path=None, overwrite=True): from datetime import datetime # Get current date and filename for subtitle today = datetime.now().strftime("%d %B %Y") # Changed to "Day Month Year" - filename = output_html_path.replace('.html', '') if output_html_path else "Report" + filename = output_html_path.replace(".html", "") if output_html_path else "Report" dark_css = """