diff --git a/.gitignore b/.gitignore index 5145354..9cd26e9 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 +Devel_unused/ +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..08a46d9 --- /dev/null +++ b/AirlockTools_Client.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 . + + +#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 + +import utils.menus as menus +from services.API import AirlockAPIWrapper +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__) + logger.debug("🔍 Logging test: this should appear in both console and file.") + 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"), + ) + + + menus.menu_main(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/README.md b/README.md index 86edfdb..6207bdc 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,9 @@ Python based Carbon Black App Control feature implementation for Airlock - "Local Approval Initialization" This programmatically scans devices in audit mode within Airlock and subsequently adds the identified blocks to a user-specified whitelist. +- KDE Wallet must be running on Linux (`kwalletd`) + + ## License **AirlockTools** is licensed under the **GNU Affero General Public License v3.0**. diff --git a/Server/scheduler_async.py b/Server/scheduler_async.py new file mode 100644 index 0000000..e20c123 --- /dev/null +++ b/Server/scheduler_async.py @@ -0,0 +1,177 @@ +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..4f5d52b --- /dev/null +++ b/flows/otp.py @@ -0,0 +1,69 @@ +# 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 services.agenthandler import selectAgents +from services.API import AirlockAPIWrapper +from utils.selector import Selector +from utils.utils import colorText, get_sanitized_input + +logger = logging.getLogger(__name__) + + + + +def 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 duration_selected: + for agent in agents: + otp_code = api.otp_generate(agent.agentid, duration_selected, purpose) + logger.info(f"Generated OTP for {agent.hostname}: {otp_code}") + otp_dict[agent.hostname] = otp_code + + return otp_dict + +def otp_activities_by_agent(api: AirlockAPIWrapper): + agents = selectAgents(api) + otp_dict = {} + for agent in agents: + otp_info = api.otp_find_by_agent(agent.agentid) + otp_dict[agent.hostname] = otp_info + + return otp_dict + +def revoke(api: AirlockAPIWrapper): + otp_dict = otp_activities_by_agent(api) + list_to_revoke = [entry["otpid"] for entry in otp_dict] + if otp_dict and list_to_revoke: + for revokee in list_to_revoke: + api.otp_revoke(revokee) \ No newline at end of file diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py new file mode 100644 index 0000000..d79d242 --- /dev/null +++ b/flows/prepPolicy.py @@ -0,0 +1,360 @@ +# 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 +from typing import List + +import dotenv +import pandas as pd + +from models.execution import ExecutionHistoryRecord, Hash +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 ( + colorText, + formatHTML, + import_to_dataframe, + 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 + + executions = [] + hashes = [] + + # Pull execution histories for each policy + + policy_executions = ExecutionHistoryRecord.from_policies( + api, selected_policies, type_=type, history_days=history_days + ) + + logger.debug(f"Policy_executions is {policy_executions}") + + executions.extend(policy_executions) + logger.debug(f"Executions contains {executions}") + if executions: + hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows() + ] + + if hashes: + unique_hashes = Hash.deduplicate(hashes) + + needs_review, approved, unapproved = Hash.categorize_hashes( + hashes=unique_hashes + ) + + categories = { + "needs_review": needs_review, + "approved": approved, + "unapproved": unapproved, + } + + for label, category in categories.items(): + csv_path = f"{working_dir}\\Needs_Review\\Review_First\\{label}_executions.csv" + html_path = f"{working_dir}\\Needs_Review\\HTML\\{label}.html" + + ExecutionHistoryRecord.enrich_with_hashes_and_export( + executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label + ) + df = import_to_dataframe(csv_path) + formatHTML(df, html_path) + +def buildPathsandPublishers(split): + working_dir = load_env("WORKING_DIR") + df1 = pd.DataFrame() + df2 = pd.DataFrame() + all_approved_hashes = pd.DataFrame() + path1 = f"{working_dir}\\Approved\\approved_executions.csv" + path2 = f"{working_dir}\\Approved\\needs_review_executions.csv" + + 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_exec" in all_approved_hashes.columns: + all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec") + else: + logger.warning("Warning: 'filename_exec' column not found in concatenated DataFrame.") + + if not all_approved_hashes.empty: + primary_path_exclusions = calculatePath( + all_approved_hashes, + split, + ) + remaining_hashes = all_approved_hashes[ + ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"]) + ] + secondary_path_exclusions = calculatePath( + remaining_hashes, split + ) + remaining_hashes = remaining_hashes[ + ~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"]) + ] + dataframes = { + "primary_Paths": primary_path_exclusions, + "secondary_Paths": secondary_path_exclusions, + "hashes_to_add": remaining_hashes, + } + logger.debug("Preparing to sort dataframes") + for name, df in dataframes.items(): + logger.debug(f" DataFrame headers: {list(df.columns)}") + if name == "hashes_to_add": df.sort_values(by="filename_exec", inplace=True) + else: df.sort_values(by="longestcfp", inplace=True) + + df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False) + formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html") + + if not all_approved_hashes.empty: + # Drop all not signed, only keep unique values + publist = all_approved_hashes[ + all_approved_hashes["publisher_hash"] != "Not Signed" + ].drop_duplicates(subset=["publisher_hash"]) + # Remove Bad publisher if somehow they made it this far + pattern = regulator(load_env_json("BAD_PUBLISHERS","[]")) + publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)] + publist = publist[["publisher_hash"]] + publist.sort_values(by="publisher_hash", inplace=True) + publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False) + +def buildPreflights(): + working_dir = load_env("WORKING_DIR") + + df1 = pd.DataFrame() + df2 = pd.DataFrame() + approved_hashes = pd.DataFrame() + approved_publishers = pd.DataFrame() + + hash = f"{working_dir}\\Approved\\hashes_to_add.csv" + path1 = f"{working_dir}\\Approved\\primary_Paths.csv" + path2 = f"{working_dir}\\Approved\\secondary_Paths.csv" + publishers = f"{working_dir}\\Approved\\publishers.csv" + + if os.path.exists(hash): + approved_hashes = pd.read_csv(hash) + + else: + logger.warning(f"File not found: {hash}") + + 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) + + 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_exec", inplace=True) + elif name == "approved_publishers" : df.sort_values(by="publisher_hash", inplace=True) + + df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False) + formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html") + +def splitFilepathsGrouped(df, col="filename"): + path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int) + min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int) + + def clean_split(path): + if not isinstance(path, (str, bytes, os.PathLike)): + return [] + parts = 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) + + # Filter out paths with fewer than `min_files_for_path` components + df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy() + split_paths = split_paths[df.index] + + 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, 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, "filename_exec") + 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 diff --git a/flows/quietAgent.py b/flows/quietAgent.py new file mode 100644 index 0000000..6e8a642 --- /dev/null +++ b/flows/quietAgent.py @@ -0,0 +1,124 @@ +# 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.selector import Selector +from utils.utils import colorText, load_env + +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), + ) + + # Get execution history as a DataFrame + 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.") + 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") diff --git a/icon.ico b/icon.ico new file mode 100644 index 0000000..eb6d558 Binary files /dev/null and b/icon.ico differ 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..a7466f6 --- /dev/null +++ b/models/execution.py @@ -0,0 +1,398 @@ +# 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 inspect +import json +import logging +import os +import re +from dataclasses import dataclass +from datetime import datetime +from typing import List, Optional + +import dotenv +import pandas as pd + +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() + + +class Hash: + """ + Hash model representing Hash data + """ + + def __init__( + self, + sha256, + applications=None, + baselines=None, + blocklists=None, + createtime=None, + datetime=None, + description=None, + filename=None, + filepath=None, + filesize=None, + md5=None, + modtime=None, + origname=None, + productname=None, + productversion=None, + publisher=None, + reputation=None, + sha128=None, + sha384=None, + sha512=None, + ): + self.sha256 = sha256 + self.applications = applications + self.baselines = baselines + self.blocklists = blocklists + self.createtime = createtime + self.datetime = datetime + self.description = description + self.filename = filename + self.filepath = filepath + self.filesize = filesize + self.md5 = md5 + self.modtime = modtime + self.origname = origname + self.productname = productname + self.productversion = productversion + self.publisher = publisher + self.reputation = reputation + self.sha128 = sha128 + self.sha384 = sha384 + self.sha512 = sha512 + + def __repr__(self): + attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + return f"" + + def __eq__(self, other): + if isinstance(other, Hash): + return self.sha256 == other.sha256 + return False + + def __hash__(self): + return hash(self.sha256) + + def to_dict(self): + """Returns a dictionary representation of the hash.""" + return self.__dict__ + + @staticmethod + def safe_int(value, default=0): + """Safely convert a value to int, returning default on failure.""" + try: + return int(value) + except (TypeError, ValueError): + return default + + + @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", "[]")) + + needs_review = [] + approved = [] + unapproved = [] + + 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.") + unapproved.append(hash_obj) + continue + + if re.search(pups_pattern, description, re.IGNORECASE): + logger.debug("Unapproved: Description matches PUP pattern.") + unapproved.append(hash_obj) + continue + + # 2. Approved: signed + if publisher != "Not Signed": + logger.debug("Approved: File is signed and not flagged.") + approved.append(hash_obj) + 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.") + unapproved.append(hash_obj) + else: + logger.debug("Approved: Unsigned file with low threat score.") + approved.append(hash_obj) + except (ValueError, TypeError): + logger.debug("Needs Review: Scannermatch score is missing or invalid.") + needs_review.append(hash_obj) + + + logger.debug(f"Final counts — Needs Review: {len(needs_review)}, Approved: {len(approved)}, Unapproved: {len(unapproved)}") + return needs_review, approved, unapproved + + + @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}") + + +""" +#Example - Convert Dataframe returned by hash query into hash objects +hash_objects = [] +for _, row in df.iterrows(): + try: + parsed_data = ast.literal_eval(row['data']) + hash_obj = Hash(sha256=row['sha256'], **parsed_data) + hash_objects.append(hash_obj) + except Exception as e: + print(f"Error parsing row: {e}") + +# Display the created Hash objects +for obj in hash_objects: + print(obj) + + +# Categorize hashes +needs_review, approved, unapproved = Hash.categorize_hashes( + hashes=hash_objects, + threat_tolerance=3, + untrusted_pattern=untrusted_pattern, + pups_pattern=pups_pattern +) + + # Deduplicate + deduped_hashes = Hash.deduplicate(hash_list) + +# Specify the directory where you want to save the CSV +output_directory = "C:/Users/Brandon/Documents/HashExports" + +# Call the export method +Hash.export_to_csv(hashes_for_export, output_directory) + + +""" + + +@dataclass +class ExecutionHistoryRecord: + # Mandatory fields + 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 + + @staticmethod + def enrich_with_hashes_and_export( + executions: list, hashes: list, directory_path: str, label: str = "enriched" + ): + exec_df = pd.DataFrame([e.__dict__ for e in executions]) + hash_df = pd.DataFrame([h.to_dict() for h in hashes]) + + logger.debug(f"Execution DataFrame columns: {exec_df.columns}") + logger.debug(f"Hash DataFrame columns: {hash_df.columns}") + + if hash_df.empty: + logger.warning(f"hash_df is empty for label: {label}. Skipping merge.") + merged_df = exec_df.copy() + else: + merged_df = pd.merge( + exec_df, + hash_df, + on="sha256", + how="left", # Preserve all executions, enrich where possible + suffixes=("_exec", "_hash") + ) + merged_df.sort_values(by="filename_exec", inplace=True) + logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}") + + filename = f"{label}_executions.csv" + os.makedirs(directory_path, exist_ok=True) + file_path = os.path.join(directory_path, filename) + merged_df.to_csv(file_path, index=False) + + logger.info(f"CSV file saved to: {file_path}") + + + @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"), + ) + + @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 + + def __repr__(self): + attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items()) + return f"" + + +""" +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..5742d0a --- /dev/null +++ b/services/API.py @@ -0,0 +1,274 @@ +# 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_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..290baba --- /dev/null +++ b/services/agenthandler.py @@ -0,0 +1,250 @@ +# 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 +import re +from dataclasses import asdict +from datetime import datetime, timedelta +from typing import List + +import pandas as pd + +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()] + + # Step 2: Create groupid → groupname map + groupid_to_name = {policy.groupid: policy.name for policy in policies} + + # Step 3: Enrich agents + for agent in agents: + agent.enrich(groupid_to_name) + + 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(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 selectAgents(api: AirlockAPIWrapper) -> List[Agent]: + print(colorText("🔍 Device Search", "cyan")) + print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan")) + print(colorText("When you're done, press Enter twice.\n", "cyan")) + print(colorText("Example:", "cyan")) + print(colorText("H00000", "cyan")) + print(colorText("UTN00000", "cyan")) + print(colorText("i-hSuperSecretServer", "cyan")) + print(colorText("u-hVenderBroke\n", "cyan")) + + print(colorText("Paste or type your device names below:", "white")) + policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()] + + device_input_lines = [] + empty_line_count = 0 + + # Regex to validate each line + 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 # Don't validate empty lines + else: + empty_line_count = 0 + + # Validate only non-empty lines + 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")) + + device_names = [name for name in device_input_lines if name] + if not device_names: + logger.debug("No device names entered") + print(colorText("⚠️ No device names entered.", "red")) + return [] + + # Build regex pattern to match hostnames + pattern = "|".join(map(re.escape, device_names)) + regex = re.compile(pattern, re.IGNORECASE) + + # Fetch agents + agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()] + matched_agents = [agent for agent in agents if regex.search(agent.hostname)] + matched_agents.sort(key=lambda agent: agent.hostname.lower()) + + # Show unmatched + unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)] + if unmatched: + logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}") + print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow")) + + if not matched_agents: + logger.debug("❌ No matching devices found.") + print(colorText("❌ No matching devices found.", "red")) + else: + logger.debug(f"✅ Found {len(matched_agents)} matching device(s).") + print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green")) + + # Enrich each agent using its class method + for agent in matched_agents: + agent.enrich_with_policies(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 + + api.agent_move(agent.agentid, target_policy) diff --git a/services/policyhandler.py b/services/policyhandler.py new file mode 100644 index 0000000..f76b969 --- /dev/null +++ b/services/policyhandler.py @@ -0,0 +1,237 @@ +# 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 + +import pandas as pd +import tqdm +from bson import ObjectId + +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 colorText + +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") diff --git a/services/security.py b/services/security.py new file mode 100644 index 0000000..093b58d --- /dev/null +++ b/services/security.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 . + +import base64 +import logging +import os +import platform +import re +import sys +from getpass import getpass + +import keyring +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC + +# Constants +KDF_ITERATIONS = 200_000 +SALT_SIZE = 16 # 128-bit Salt +NONCE_SIZE = 12 # AES-GCM +KEY_SIZE = 32 # AES-256 + + +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) + + +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() + + 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 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..56c8f3c --- /dev/null +++ b/utils/configmanager.py @@ -0,0 +1,115 @@ +import json +import logging +import os +import sys +from pathlib import Path +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/menus.py b/utils/menus.py new file mode 100644 index 0000000..db878c2 --- /dev/null +++ b/utils/menus.py @@ -0,0 +1,316 @@ +# 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 re + +import dotenv +import pandas as pd + +import services.policyhandler as policyh +from flows.otp import generate, otp_activities_by_agent, revoke +from flows.prepPolicy import ( + buildPathsandPublishers, + buildPreflights, + selectAllowlists, + selectPolicies, + sortHashes, +) +from flows.quietAgent import findQuietAgents +from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents +from services.API import AirlockAPIWrapper +from utils.configmanager import load_env +from utils.selector import Selector +from utils.utils import ( + areYouSure, + colorText, + displayIntro, + get_sanitized_input, + open_directory, + printEnforceChecklist, +) + +logger = logging.getLogger(__name__) + +dotenv.load_dotenv() + +def menu_main(api: AirlockAPIWrapper): + working_dir = load_env("WORKING_DIR") + extras = load_env("EXTRAS") + while True: + displayIntro() + # Add Settings, and give option to change working dir + print(colorText("1. ✅ - Move Device(s) to local approval", "yellow")) + print(colorText("2. 🎫 - OTP", "yellow")) + print(colorText("3. 🔄 - Move to Audit/Enforcement", "yellow")) + print(colorText("4. 🔍 - Device Search", "yellow")) + print(colorText("5. 🔇 - Find Quiet Hosts", "yellow")) + if extras == "POLICYPREP" : print(colorText("6. 🛡️ - Policy Enforcement Tools", "yellow")) + print(colorText("F. 📂 - Open Working Directory", "yellow")) + print(colorText("S. 🛠️ - Settings", "yellow")) + print(colorText("Q. 🔚 - Quit", "yellow")) + + choice = get_sanitized_input("\nEnter Menu Item: ") + if choice == "1": + print("This Feature is still in development") + get_sanitized_input("Press enter to continue") + elif choice == "2": + menu_otp(api) + elif choice == "3": + choices = ["audit", "enforcement"] + print(colorText("Move devices to which state?:", "yellow")) + direction = Selector.select_string(choices, False, False) + devices = selectAgents(api) + print(colorText("Would you like to continue with these devices?","white")) + for device in devices: + print(device.hostname) + confirm = Selector.confirm() + if direction and devices and confirm: + for device in devices: + moveAgentToRelatedPolicy(api,device, direction) + elif choice == "4": + findAgents(api,False) + elif choice == "5": + findQuietAgents(api) + elif choice == "6": + if extras == "POLICYPREP": menu_policymanagment(api) + elif choice.upper() == "F": + open_directory(working_dir) + elif choice.upper() == "S": + menu_settings() + elif choice.upper() == "Q": + break + else: + print(colorText("Invalid choice. Please try again.", "red")) + + + +def menu_policy_enforce(api: AirlockAPIWrapper): + selected_policies = [] + destination_policy = [] + destination_allowlist = [] + processed_paths = [] + processed_hashes = [] + processed_publishers = [] + tested = False + 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": + selected_policies = selectPolicies(api,True) + + elif choice == "2": + 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": + sortHashes( + api, + selected_policies, + type=[1, 2, 6, 7], + ) + + elif choice == "4": + if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"): + buildPathsandPublishers(False) + else: + print("File not found. Please make sure it's saved correctly and try again.") + + elif choice == "5": + if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists( + f"{working_dir}\\Approved\\primary_Paths.csv" + ): + buildPreflights() + else: + print("File not found. Please make sure it's saved correctly and try again.") + + elif choice == "6": + if ( + os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") + and os.path.exists(f"{working_dir}\\Preflight\\approved_hashes.csv") + and destination_policy + and destination_allowlist + ): + print(colorText("These path exclusions would be added to:", "yellow")) + print(destination_policy) + + pathexclusions = pd.read_csv(f"{working_dir}\\Preflight\\approved_paths.csv") + hashes = pd.read_csv(f"{working_dir}\\Preflight\\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) + ] + + print(processed_paths) + print(colorText("These publishers would added", "yellow")) + + if os.path.exists(f"{working_dir}\\Preflight\\approved_publishers.csv"): + publishers = pd.read_csv(f"{working_dir}\\Preflight\\approved_hashes.csv") + if publishers.empty: + print(colorText("The publishers list is empty.", "red")) + else: + processed_publishers = ( + publishers[publishers["publisher_hash"] != "Not Signed"] + ["publisher_hash"] + .drop_duplicates() + .tolist() + ) + print(processed_publishers) + + print(colorText("These hashes would be added to:", "yellow")) + print(destination_allowlist) + + processed_hashes = hashes["sha256"].unique().tolist() + print(processed_hashes) + + if processed_paths and processed_hashes: + tested = True + else: + # Log which condition(s) failed + missing_items = [] + if not os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv"): + missing_items.append("approved_paths.csv not found") + if not os.path.exists(f"{working_dir}\\Preflight\\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": + areYouSure() + confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") + if ( + tested + 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) + else: + logger.error("Confirmation block failed. Reasons:") + if not tested: + logger.error(" - Preflight checks were not completed successfully (`tested` is False).") + 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() == "S": + menu_settings() + elif choice.upper == "B": + break + + + else: + print(colorText("Invalid choice. Please try again.", "red")) + + +def menu_otp(api: AirlockAPIWrapper): + working_dir = load_env("WORKING_DIR") + while True: + + print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan")) + print(colorText("1. 🔐 -Generate OTPs", "cyan")) + print(colorText("2. 📊 -OTP Activities By Agent", "cyan")) + print(colorText("3. ❌ -Revoke OTPs", "cyan")) + print(colorText("F. 📂 - Open Working Directory", "yellow")) + print(colorText("S. 🛠️ - Settings", "yellow")) + print(colorText("B. 🔙 - Back", "yellow")) + + choice = get_sanitized_input("Enter your choice: ") + + if choice == "1": + otp_list = generate(api) + print(colorText(otp_list,"green")) + elif choice == "2": + otp_activities_by_agent(api) + elif choice == "3": + revoke(api) + elif choice.upper() == "F": + open_directory(working_dir) + elif choice.upper() == "S": + menu_settings() + elif choice.upper() == "B": + break + +def menu_policymanagment(api: AirlockAPIWrapper): + working_dir = load_env("WORKING_DIR") + while True: + print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow")) + print(colorText("2. 🔄 - Update Audit Policies from Enforcement Policies", "yellow")) + print(colorText("F. 📂 - Open Working Directory", "yellow")) + print(colorText("S. 🛠️ - Settings", "yellow")) + print(colorText("B. 🔙 - Back", "yellow")) + choice = get_sanitized_input("\n Enter Menu Item: ") + + if choice == "1": + menu_policy_enforce(api) + elif choice == "2": + areYouSure() + confirmation = get_sanitized_input("Type 'I AGREE' to continue: ") + if confirmation.strip() == "I AGREE": + policyh.updateAuditPoliciesFromEnforcementPolices(api) + + elif choice.upper() == "F": + open_directory(working_dir) + elif choice.upper() == "S": + menu_settings() + elif choice.upper() == "B": + break + else: + print(colorText("Invalid choice. Please try again.", "red")) +def menu_settings(): + while True: + print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan")) + print(colorText("This Feature is still in development", "cyan")) + # print(colorText("2. Sub-option B","cyan")) + print(colorText("B. 🔙 - Back", "yellow")) + choice = get_sanitized_input("Enter your choice: ") + + if choice == "1": + pass #TODO ADD CHANGE WORKDIR CODE + + elif choice.upper() == "B": + print("Returning to Main Menu...") + break + else: + print("Invalid choice. Please try again.") 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/pretty.py b/utils/pretty.py deleted file mode 100644 index aaba329..0000000 --- a/utils/pretty.py +++ /dev/null @@ -1,339 +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 os - -def colorText(text: str, color: str) -> str: - colors = { - "red": "\033[91m", - "green": "\033[92m", - "yellow": "\033[93m", - "blue": "\033[94m", - "magenta": "\033[95m", - "cyan": "\033[96m", - "white": "\033[97m", - "reset": "\033[0m" - } - - return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}" - -def style_dataframe_dark(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" - - dark_css = """ - - """ - - header = f""" -
-

Airlock Tools

-

{filename} - {today}

-
- """ - - html_table = df.to_html(index=False, escape=False) - styled_html = ( - f"\n" - f"Airlock Tools Report\n" - f"\n" - f"{dark_css}\n" - f"{header}\n" - f"
\n" - f" {html_table}\n" - f"
\n" - f"\n" - f"" - ) - if output_html_path: - with open(output_html_path, "w", encoding="utf-8") as f: - f.write(styled_html) - print(f"✅ Styled table saved to '{output_html_path}'") - elif overwrite: - import tempfile - temp_path = tempfile.mktemp(suffix=".html") - with open(temp_path, "w", encoding="utf-8") as f: - f.write(styled_html) - print(f"✅ Styled table saved to temporary file: {temp_path}") - else: - return styled_html - - -def displayIntro(): - - print(colorText(r""" - ███ - ████ ░████████ - █████████████ ███████████████ - █████████████████████ █████████████████████ - ███████████████████ ██████████████████████▓ - ███████████████████ ██████████████████████ - █████████████████████ ███████████████████████ - ████████████████████████████████████████████████████████ - █████████ ██ ██ █████████ - █████████ ██ ███ █ █████████ - █████████ ██ ████ █████ █████████████ - █████████ ██ ██████ █████████████ - ████████ ██ ███████ ████████████░ - ███████ ██ ██▓ ██████ ████████████ - ██████ ██ ████ █████ ███████████ - █████████████████████████████████████████████████ - ▒████████████████████ ██████████████████ - ███████████████████ ███████████████▒ - ███████████████ █████████████ - ██████████ ███████████ - ████████ - ████ -""", "yellow")) - print(colorText(r""" - _____ .__ .__ __ ___________ .__ - / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ - / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ -/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ -\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > - \/ \/ \/ \/ -""", "cyan")) - print(colorText("=================================================================================", "cyan")) - print(colorText("======================== Welcome to the Airlock API Tool ========================", "cyan")) - print(colorText("=================================================================================", "cyan")) - -def printEnforceChecklist(first_policy, second_policy, allowlist_child_name, allowlist_parent_name, destination_name): - - 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")) - if first_policy == " " and second_policy == " ": - print(colorText(f" [✗] No policies have been chosen","red")) - elif first_policy != " " and second_policy is first_policy: - print(colorText(f" [✓] {first_policy} has been selected,", "green")) - elif first_policy != " " and second_policy != " ": - print(colorText(f" [✓] {first_policy} has been selected as Policy 1","green")) - print(colorText(f" [✓] {second_policy} has been selected as Policy 2","green")) - - - - print(colorText("2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", "cyan")) - - if os.path.exists(f"parquet\\execution_history_{first_policy}.parquet"): - print(colorText(f" [✓] Execution history has been compiled for {first_policy}","green")) - elif not os.path.exists(f"parquet\\execution_history_{first_policy}.parquet"): - print(colorText(f" [✗] Execution history has not been compiled for {first_policy}","red")) - elif second_policy is not first_policy and os.path.exists(f"parquet\\execution_history_{second_policy}.parquet"): - print(colorText(f" [✓] Execution history has been compiled for {second_policy}","green")) - elif second_policy is not first_policy and not os.path.exists(f"parquet\\execution_history_{second_policy}.parquet"): - print(colorText(f" [✗] Execution history has not been compiled for {second_policy}","red")) - - if os.path.exists(f"parquet\\combined_hashlist_{first_policy}_{second_policy}.parquet"): - print(colorText(f" [✓] Hash Info has been added to the combined execution history", "green")) - else: - print(colorText(f" [✗] Hash Info has not been added to the combined execution history", "red")) - - if os.path.exists(f"parquet\\hashes_rep_unknown_{first_policy}_{second_policy}.parquet") and os.path.exists(f"parquet\\hashes_rep_good_{first_policy}_{second_policy}.parquet") and os.path.exists(f"parquet\\hashes_rep_bad_{first_policy}_{second_policy}.parquet"): - print(colorText(f" [✓] Hashes have been cateogrized", "green")) - else: - print(colorText(f" [✗] Hashes have not been cateogrized", "red")) - - if os.path.exists(f"parquet\\condensed_executions_{first_policy}_{second_policy}.parquet"): - print(colorText(f" [✓] Execution history has been_combined_for {first_policy} and_{second_policy}", "green")) - else: - print(colorText(f" [✗] Execution history has not been_combined_for {first_policy} and_{second_policy}", "red")) - - - print(colorText(f"3. Manually review the files:","cyan")) - print(colorText(" 'needs_approved\\hashes_rep_good_{first_policy}_{second_policy}.csv' and 'needs_approved\\hashes_rep_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")) - - if os.path.exists(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv") and os.path.exists(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv"): - print(colorText(" [✓] Reviewed hashes have been loaded","green")) - else: - print(colorText(" [✗] Reviewed hashes have not been loaded","red")) - - if os.path.exists(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet"): - print(colorText(" [✓] The combined approved hashes list has been generated","green")) - else: - print(colorText(" [✗] The combined approved hashes list has not been generated","red")) - - if os.path.exists(f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.csv"): - print(colorText(" [✓] Path review list created","green")) - else: - print(colorText(" [✗] Path review list has not been created","red")) - - - print(colorText(f"4. Manually review the file 'needs_approved\\paths_needing_review_{first_policy}_{second_policy}.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(" Preflight Lists will be generated", "cyan")) - - if os.path.exists(f"approved\\path_needs_approved_{first_policy}_{second_policy}.csv"): - print(colorText(" [✓] Reviewed path list detected","green")) - else: - print(colorText(" [✗] Path review list has not been detected","red")) - - if os.path.exists(f"preflight\\final_path_exclusions_{first_policy}_{second_policy}.html"): - print(colorText(" [✓] Preflight Path Exclusion List has been generated","green")) - else: - print(colorText(" [✗] Preflight Path Exclusion List has not been generated","red")) - - if os.path.exists(f"preflight\\final_hash_approvals_{first_policy}_{second_policy}.html"): - print(colorText(" [✓] Preflight hash approval list has been generated","green")) - else: - print(colorText(" [✗] Preflight hash approval list has not been generated","red")) - - - print(colorText(f"5. Choose the destination policy and parent and child allow list", "cyan")) - if allowlist_child_name == " " and allowlist_parent_name== " ": - print(colorText(f" [✗] No allowlists have been chosen","red")) - elif allowlist_parent_name != " " and allowlist_child_name != " " and allowlist_parent_name is allowlist_child_name: - print(colorText(f" [✓] [✗] Only {allowlist_parent_name} has been selected this is unusual, but potentially valid case, double check before proceeding,", "yellow")) - elif allowlist_parent_name != " " and allowlist_child_name != " " and allowlist_parent_name is not allowlist_child_name: - print(colorText(f" [✓] {allowlist_parent_name} has been selected as Parent Policy","green")) - print(colorText(f" [✓] {allowlist_child_name} has been selected as Child Policy","green")) - if destination_name == " ": - print(colorText(f" [✗] No destination policy has been chosen","red")) - else: - print(colorText(f" [✓] destination policy is {destination_name}","green")) - - print(colorText(f"6. Liftoff ------------------------------------------------------", "cyan")) - print(colorText(f" Apply path exclusions according to allowed and approved paths", "cyan")) - print(colorText(f" Apply signed or attested hashes to Parent Allow List", "cyan")) - print(colorText(f" Apply approved, but unsigned hashes to the Child Allow List", "cyan")) - - print(colorText("Q. Quit", "cyan")) - - - - -def areYouSure(): - print(colorText(f"*******************************************************************************************************************************************","red")) - print(colorText(f"*=========================================================================================================================================*","yellow")) - print(colorText(f"*=========================================================================================================================================*","red")) - print(colorText(f"*-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------*", "yellow")) - print(colorText(f"*=========================================================================================================================================*","red")) - print(colorText(f"*=========================================================================================================================================*","yellow")) - print(colorText(f"*******************************************************************************************************************************************","red")) - -def locked(): - - print(colorText(r""" - ████████████████████████████████████████████████████████████████ - ███ ██ - ██ ██████ ███ - ██ ████████████ ███ - ██ ████ ███ ███ - ██ ███ ███ ███ - ██ ███ ███ ███ - ██ ▒████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ██████████████████████ ███ - ██ ███ - ███ ███ - ████████████████████████████████████████████████████████████████████ - ▒██████████████████████████████████████████████████████████████████▒ - ▒████ - ▒████ - ▓██████████████████████████████████████████ - █████████████████████████████████████████████░ -""", "yellow")) \ No newline at end of file diff --git a/utils/selector.py b/utils/selector.py new file mode 100644 index 0000000..87f1bfb --- /dev/null +++ b/utils/selector.py @@ -0,0 +1,171 @@ +# 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 + +from utils.utils import get_sanitized_input + +logger = logging.getLogger(__name__) + + +class Selector: + @staticmethod + def _display_choices( + items: List[Any], + label_func: Callable[[Any], str], + num_columns: int = 4, + header: str = "Available Choices:" + ) -> None: + sorted_items = sorted(items, key=lambda item: label_func(item).lower()) + rows = (len(sorted_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(sorted_items): + label = label_func(sorted_items[idx]) + line += f"{idx + 1}: {label:<30}" + print(line) + + @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:" + ) -> Union[Optional[Any], List[Any]]: + if not items: + logger.warning("No items available for selection.") + return None + + Selector._display_choices(items, label_func, header=header) + sorted_items = sorted(items, key=lambda item: label_func(item).lower()) + selected = [] + + if allow_multiple: + while True: + choice = get_sanitized_input("Select an item by number (or Q to finish): ").strip().lower() + if choice == "q": + break + try: + index = int(choice) + if 1 <= index <= len(sorted_items): + item = sorted_items[index - 1] + if item not in selected: + selected.append(item) + if prompt_each: + logger.info(f"Selected: {label_func(item)}") + else: + logger.warning("Item already selected.") + else: + logger.warning("Selection out of range. Try again.") + except ValueError: + logger.warning("Invalid input. Enter a number or 'Q' to quit.") + return selected if selected else None + else: + try: + choice = int(get_sanitized_input("Select one item by number: ")) + if 1 <= choice <= len(sorted_items): + selected_item = 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_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'.") \ No newline at end of file diff --git a/utils/setup.py b/utils/setup.py new file mode 100644 index 0000000..d2d70ef --- /dev/null +++ b/utils/setup.py @@ -0,0 +1,175 @@ +# 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.handlers +import os +import platform +import sys +from pathlib import Path + +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" + logger = logging.getLogger() + logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG)) + + # 🔧 Clear existing handlers + for handler in logger.handlers[:]: + logger.removeHandler(handler) + + file_handler = logging.handlers.RotatingFileHandler( + log_file, maxBytes=5_000_000, backupCount=5, encoding='utf-8' + ) + file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(file_handler) + + console_handler = logging.StreamHandler() + console_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s')) + logger.addHandler(console_handler) + + if platform.system() == "Windows": + try: + event_handler = logging.handlers.NTEventLogHandler("AirlockTools") + event_handler.setLevel(logging.CRITICAL) + event_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s')) + logger.addHandler(event_handler) + except Exception as e: + logger.warning(f"Could not attach Windows Event Log handler: {e}") + + logger.debug("✅ Logging configured.") + +def get_system_config_path() -> Path: + base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))) + return base_path.parent / "system_config.json" + +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() -> Path: + 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) + + return working_dir \ No newline at end of file diff --git a/utils/utils.py b/utils/utils.py new file mode 100644 index 0000000..6803411 --- /dev/null +++ b/utils/utils.py @@ -0,0 +1,708 @@ +# 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 platform +import re +import subprocess +import tempfile +import tkinter as tk +from tkinter import filedialog + +import pandas as pd + +from utils.configmanager import load_env + +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 displayIntro(): + print( + colorText( + r""" + ███ + ████ ░████████ + █████████████ ███████████████ + █████████████████████ █████████████████████ + ███████████████████ ██████████████████████▓ + ███████████████████ ██████████████████████ + █████████████████████ ███████████████████████ + ████████████████████████████████████████████████████████ + █████████ ██ ██ █████████ + █████████ ██ ███ █ █████████ + █████████ ██ ████ █████ █████████████ + █████████ ██ ██████ █████████████ + ████████ ██ ███████ ████████████░ + ███████ ██ ██▓ ██████ ████████████ + ██████ ██ ████ █████ ███████████ + █████████████████████████████████████████████████ + ▒████████████████████ ██████████████████ + ███████████████████ ███████████████▒ + ███████████████ █████████████ + ██████████ ███████████ + ████████ + ████ +""", + "yellow", + ) + ) + print( + colorText( + r""" + _____ .__ .__ __ ___________ .__ + / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ + / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ +/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ +\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > + \/ \/ \/ \/ +""", + "cyan", + ) + ) + print( + colorText( + "=================================================================================", + "cyan", + ) + ) + print( + colorText( + "======================== Welcome to the Airlock API Tool ========================", + "cyan", + ) + ) + print( + colorText( + "=================================================================================", + "cyan", + ) + ) + + +def printEnforceChecklist(selected_policies, destination_policy, destination_allowlist): + working_dir = load_env("WORKING_DIR") + 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", + ) + ) + if not selected_policies: + print(colorText(" [✗] No policies have been chosen", "red")) + else: + print(colorText("The following policies have been choosen:", "green")) + for policy in selected_policies: + print(colorText(f" [✓] {policy.name}", "green")) + + print(colorText("2. Choose the destination policy and allowlist", "cyan")) + + if not destination_policy: + print(colorText(" [✗] No destination policy has been chosen", "red")) + elif destination_policy: + print(colorText(f" [✓] {destination_policy[0].name} has been selected as the destination policy", "green")) + + + + if not destination_allowlist: + print(colorText(" [✗] No allowlist has been chosen", "red")) + elif destination_allowlist: + print( + colorText( + f" [✓] {destination_allowlist[0].name} has been selected as allowlist", + "green", + ) + ) + + + + print( + colorText( + "3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", + "cyan", + ) + ) + if not selected_policies: + print(colorText(" [✗] No policies have been chosen", "red")) + else: + if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"): + print(colorText(" [✓] Data has been fetched", "green")) + else: + print(colorText(" [✗] Data has not been fetched", "red")) + + print(colorText("4. Manually review the files:", "cyan")) + print( + colorText( + " 'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n", + "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", + ) + ) + + if os.path.exists(f"{working_dir}\\Approved\\approved_executions.csv"): + print(colorText(" [✓] Reviewed hashes have been loaded", "green")) + else: + print(colorText(" [✗] Reviewed hashes have not been loaded", "red")) + + if os.path.exists( + f"{working_dir}\\Needs_Review\\Review_Second\\primary_Paths.csv", + ): + print(colorText(" [✓] Path review list created", "green")) + else: + print(colorText(" [✗] Path review list has not been created", "red")) + + print( + colorText( + "5. Manually review the files \n 'prepare_policy\\needs_approved\\primary_Paths.csv'\n 'prepare_policy\\needs_approved\\secondary_Paths.csv'", + "cyan", + ) + ) + print( + colorText( + " Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful", + "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")) + + if os.path.exists( + f"{working_dir}\\Approved\\primary_Paths.csv", + ): + print(colorText(" [✓] Reviewed path list detected", "green")) + else: + print(colorText(" [✗] Path review list has not been detected", "red")) + + if os.path.exists(f"{working_dir}\\Preflight\\approved_paths.csv") and os.path.exists( + f"{working_dir}\\Preflight\\approved_hashes.csv" + ): + print(colorText(" [✓] Preflight Path Exclusion List has been generated", "green")) + else: + print(colorText(" [✗] Preflight Path Exclusion List has not been generated", "red")) + + 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("F. 📂 - Open Working Directory", "cyan")) + print(colorText("Q. 🔚 - Quit", "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", + "yellow": "\033[93m", + "blue": "\033[94m", + "magenta": "\033[95m", + "cyan": "\033[96m", + "white": "\033[97m", + "reset": "\033[0m", + } + return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}" + + +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" + + dark_css = """ + + """ + + header = f""" +
+

Airlock Tools

+

{filename} - {today}

+
+ """ + + html_table = df.to_html(index=False, escape=False) + styled_html = ( + f"\n" + f"Airlock Tools Report\n" + f"\n" + f"{dark_css}\n" + f"{header}\n" + f"
\n" + f" {html_table}\n" + f"
\n" + f"\n" + f"" + ) + if output_html_path: + with open(output_html_path, "w", encoding="utf-8") as f: + f.write(styled_html) + print(f"✅ Styled table saved to '{output_html_path}'") + elif overwrite: + with tempfile.NamedTemporaryFile( + suffix=".html", delete=False, mode="w", encoding="utf-8" + ) as f: + f.write(styled_html) + temp_path = f.name + + print(f"✅ Styled table saved to temporary file: {temp_path}") + else: + return styled_html + + + +def open_directory(path): + system = platform.system() + + if system == "Windows": + os.startfile(path) + elif system == "Linux": + subprocess.run(["xdg-open", path]) + else: + raise OSError(f"Unsupported operating system: {system}") + +