diff --git a/.gitignore b/.gitignore index 9cd26e9..7d2a7e0 100644 --- a/.gitignore +++ b/.gitignore @@ -10,5 +10,5 @@ jobs.json securitytest.py *.toml system_config.json -Devel_unused/ +Development/ AirlockTools_client*/ \ No newline at end of file diff --git a/Development/AdminTools/policycloner.py b/Development/AdminTools/policycloner.py new file mode 100644 index 0000000..a4d642a --- /dev/null +++ b/Development/AdminTools/policycloner.py @@ -0,0 +1,46 @@ +import sys +import os +import logging + + +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +sys.path.append(project_root) + + +from services.API import AirlockAPIWrapper +from services.security import getAPI + +logger = logging.getLogger() +logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') +def main(): + try: + url = "https://172.17.22.240:3129" + 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 + + + if username: + api = AirlockAPIWrapper( + base_url= url, + api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType] + ) + + source = "" + target = "" + response = api.policy_clone(source, target) + print(response) + + +if __name__ == "__main__": + main() diff --git a/Development/AdminTools/policylister.py b/Development/AdminTools/policylister.py new file mode 100644 index 0000000..beb5f6e --- /dev/null +++ b/Development/AdminTools/policylister.py @@ -0,0 +1,44 @@ +import sys +import os +import logging + + +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', "..")) +sys.path.append(project_root) + + +from services.API import AirlockAPIWrapper +from services.security import getAPI + +logger = logging.getLogger() +logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') +def main(): + try: + url = "https://172.17.22.240:3129" + 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 + + + if username: + api = AirlockAPIWrapper( + base_url= url, + api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType] + ) + + response = api.policy_find_all() + response.to_csv("All_Policies.csv", index=False) + + print(response) +if __name__ == "__main__": + main() diff --git a/Development/AdminTools/scriptPolicySetter.py b/Development/AdminTools/scriptPolicySetter.py new file mode 100644 index 0000000..1f32549 --- /dev/null +++ b/Development/AdminTools/scriptPolicySetter.py @@ -0,0 +1,75 @@ +import sys +import os +import logging + + +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +sys.path.append(project_root) + + +from services.API import AirlockAPIWrapper +from services.security import getAPI + +logger = logging.getLogger() +logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') +def main(): + try: + url = "https://172.17.22.240:3129" + 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 + + + if username: + api = AirlockAPIWrapper( + base_url= url, + api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType] + ) + """ + Available script types are: + + "batch", + "powershell", + "command", + "vbscript", + "javascript" + ,"windowsinstaller", + "htmlapplication", + "javaapplication", + "windowsscriptcomponent", + "compiledhtml", + "shellscript", + "dylib", + "python" + + """ + groupid = "5aebf6a0-1d67-47b4-9c5f-2866ffca5671" #AT Testing + script_custom = 1 + scripts_audit = [ + "batch", + "powershell", + "command", + "vbscript", + "javascript", + "windowsinstaller", + "javaapplication", + "python" + ] + + scripts_disabled = ["compiledhtml", "htmlapplication", "shellscript", "dylib","windowsscriptcomponent"] + scripts_respect = [] + + response = api.policy_set_script_custom(groupid, script_custom, scripts_audit, scripts_disabled, scripts_respect) + print(response) +if __name__ == "__main__": + main() diff --git a/Development/AdminTools/testsetup.py b/Development/AdminTools/testsetup.py new file mode 100644 index 0000000..9af12c9 --- /dev/null +++ b/Development/AdminTools/testsetup.py @@ -0,0 +1,39 @@ +import os +import json +from utils.configmanager import ( + load_protected_config, + get_protected_value, + get_protected_json, + PROTECTED_KEYS +) +from utils.setup import setup + +def test_protected_config(): + print("๐Ÿ”’ Testing protected config loading...") + protected = load_protected_config() + assert isinstance(protected, dict), "Protected config should be a dictionary" + for key in PROTECTED_KEYS: + assert key in protected, f"Missing protected key: {key}" + print(f"โœ” {key} = {protected[key]}") + +def test_json_parsing(): + print("\n๐Ÿงช Testing JSON parsing for POLICY_MAP_ENF_AUD...") + policy = get_protected_json("POLICY_MAP_ENF_AUD") + assert isinstance(policy, dict), "POLICY_MAP_ENF_AUD should be a dictionary" + print(f"โœ” POLICY_MAP_ENF_AUD = {json.dumps(policy, indent=2)}") + +def test_setup_env(): + print("\nโš™๏ธ Running setup() to validate environment setup...") + working_dir = setup() + assert working_dir.exists(), "Working directory should exist" + print(f"โœ” Working directory: {working_dir}") + + print("\n๐ŸŒ Checking .env values (excluding protected)...") + for key in os.environ: + if key not in PROTECTED_KEYS: + print(f"๐Ÿ”ง {key} = {os.environ[key]}") + +if __name__ == "__main__": + test_protected_config() + test_json_parsing() + test_setup_env() diff --git a/Development/WIP/Server/datastore.py b/Development/WIP/Server/datastore.py new file mode 100644 index 0000000..df9721f --- /dev/null +++ b/Development/WIP/Server/datastore.py @@ -0,0 +1,170 @@ +""" +datastore.py + +Utility module for saving/loading objects to/from JSON and inserting/updating/reading them in SQLite. + +Supports any class with: +- a `to_dict()` method +- a constructor accepting `**kwargs` +- optionally, a `from_dict()` method + +Author: Brandon Wickline, James Brotosky +License: GNU Affero General Public License v3.0 +""" + +import json +import sqlite3 +import logging +from typing import Type, List, TypeVar + +# Setup logger +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +# Generic type variable for typing +T = TypeVar('T') + +# ------------------------------------------------------------------- +# JSON SAVE +# ------------------------------------------------------------------- +def save_to_json(obj_list: List[T], filepath: str) -> None: + """ + Save a list of objects to a JSON file using their `to_dict()` method. + + Args: + obj_list (List[T]): List of objects to serialize. + filepath (str): Path to the output JSON file. + + Example: + save_to_json(policies, "policies.json") + """ + try: + with open(filepath, 'w', encoding='utf-8') as f: + json.dump([obj.to_dict() for obj in obj_list], f, ensure_ascii=False, indent=4) # pyright: ignore[reportAttributeAccessIssue] + logger.info(f"Saved {len(obj_list)} objects to {filepath}") + except Exception as e: + logger.error(f"Failed to save to {filepath}: {e}") + +# ------------------------------------------------------------------- +# JSON LOAD +# ------------------------------------------------------------------- +def load_from_json(cls: Type[T], filepath: str) -> List[T]: + """ + Load a list of objects from a JSON file and instantiate them using the class constructor + or a `from_dict()` method if available. + + Args: + cls (Type[T]): Class type to instantiate. + filepath (str): Path to the input JSON file. + + Returns: + List[T]: List of instantiated objects. + + Example: + loaded_agents = load_from_json(Agent, "agents.json") + """ + try: + with open(filepath, 'r', encoding='utf-8') as f: + data_list = json.load(f) + logger.info(f"Loaded {len(data_list)} records from {filepath}") + if hasattr(cls, "from_dict"): + return [cls.from_dict(data) for data in data_list] # pyright: ignore[reportAttributeAccessIssue] + return [cls(**data) for data in data_list] + except Exception as e: + logger.error(f"Failed to load from {filepath}: {e}") + return [] + +# ------------------------------------------------------------------- +# SQLITE INSERT OR UPDATE +# ------------------------------------------------------------------- +def insert_or_update_objects_to_sqlite(obj_list: List[T], table_name: str, db_path: str, primary_key: str) -> None: + """ + Insert or update a list of objects into a SQLite table. + Uses `ON CONFLICT(primary_key) DO UPDATE` for upsert behavior. + + Args: + obj_list (List[T]): List of objects with `to_dict()` method. + table_name (str): Name of the SQLite table. + db_path (str): Path to the SQLite database file. + primary_key (str): Field name to use as the primary key. + + Example: + insert_or_update_objects_to_sqlite(hashes, "hashes", "data_store.db", primary_key="sha256") + """ + if not obj_list: + logger.warning("No objects to insert or update into SQLite.") + return + + try: + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + + sample_dict = obj_list[0].to_dict() # pyright: ignore[reportAttributeAccessIssue] + columns = ', '.join(sample_dict.keys()) + placeholders = ', '.join(['?'] * len(sample_dict)) + update_clause = ', '.join([f"{key}=excluded.{key}" for key in sample_dict.keys() if key != primary_key]) + + # Create table if it doesn't exist + create_stmt = f""" + CREATE TABLE IF NOT EXISTS {table_name} ( + {', '.join([f"{key} TEXT" for key in sample_dict.keys()])}, + PRIMARY KEY ({primary_key}) + ) + """ + cursor.execute(create_stmt) + + # Insert or update each object + for obj in obj_list: + values = tuple(str(v) if v is not None else "" for v in obj.to_dict().values()) # pyright: ignore[reportAttributeAccessIssue] + insert_stmt = f""" + INSERT INTO {table_name} ({columns}) VALUES ({placeholders}) + ON CONFLICT({primary_key}) DO UPDATE SET {update_clause} + """ + cursor.execute(insert_stmt, values) + + conn.commit() + conn.close() + logger.info(f"Inserted or updated {len(obj_list)} records into {table_name} table in {db_path}") + except Exception as e: + logger.error(f"Failed to insert or update into SQLite: {e}") + +# ------------------------------------------------------------------- +# SQLITE READ +# ------------------------------------------------------------------- +def read_objects_from_sqlite(cls: Type[T], table_name: str, db_path: str) -> List[T]: + """ + Generic function to read rows from a SQLite table and convert them into class instances. + + Args: + cls (Type[T]): The class to instantiate (e.g., Agent, Policy, Hash). + table_name (str): The name of the table to query. + db_path (str): Path to the SQLite database file. + + Returns: + List[T]: List of class instances. + + Example: + agents = read_objects_from_sqlite(Agent, "agents", "data_store.db") + """ + try: + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute(f"SELECT * FROM {table_name}") + rows = cursor.fetchall() + columns = [desc[0] for desc in cursor.description] + + objects = [] + for row in rows: + data = dict(zip(columns, row)) + if hasattr(cls, "from_dict"): + obj = cls.from_dict(data) # pyright: ignore[reportAttributeAccessIssue] + else: + obj = cls(**data) + objects.append(obj) + + conn.close() + logger.info(f"Read {len(objects)} records from {table_name} table in {db_path}") + return objects + except Exception as e: + logger.error(f"Failed to read from {table_name} in {db_path}: {e}") + return [] \ No newline at end of file diff --git a/Development/WIP/Server/localApproval.py b/Development/WIP/Server/localApproval.py new file mode 100644 index 0000000..789b917 --- /dev/null +++ b/Development/WIP/Server/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/Development/WIP/Server/scheduler.py b/Development/WIP/Server/scheduler.py new file mode 100644 index 0000000..a653c3e --- /dev/null +++ b/Development/WIP/Server/scheduler.py @@ -0,0 +1,352 @@ +# 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 time +from typing import Any, Callable, Dict, List + +import schedule + +logger = logging.getLogger(__name__) + + +# TODO - move this File where all jobs are persisted +JOBS_FILE = f"{os.getenv('WORKING_DIR')}\\scheduling\\jobs.json" + +# Ensure directory exists +os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True) + +# Registry of functions that can be scheduled +FUNCTION_MAP: Dict[str, Callable] = {} + +# ------------------------------- +# Function Registration +# ------------------------------- + + +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 + + +# ------------------------------- +# Persistence Helpers +# ------------------------------- + + +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 _atomic_save(path: str, data: Any): + """Write JSON atomically to avoid partial writes.""" + tmp = f"{path}.tmp" + with open(tmp, "w") as f: + json.dump(data, f, indent=4) + os.replace(tmp, path) + + +def save_jobs(jobs: List[Dict[str, Any]]): + """Save jobs to the JSON file (overwrite).""" + _atomic_save(JOBS_FILE, jobs) + + +# ------------------------------- +# Uniqueness Helpers +# ------------------------------- + + +def job_in_store(job_id: str) -> bool: + """Check if a job id exists in the persisted JSON file.""" + return any(j.get("id") == job_id for j in load_jobs()) + + +def job_in_scheduler(job_id: str) -> bool: + """ + Check if a job with this tag exists in the in-memory scheduler. + Uses schedule.get_jobs(tag=...) if available, otherwise scans tags. + """ + try: + jobs = schedule.get_jobs(tag=job_id) # schedule >= 1.2.0 + return len(jobs) > 0 + except TypeError: + # Fallback for older versions + return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs) + + +def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool: + """ + Ensure the job_id is unique across persistence and in-memory schedule. + on_conflict: + - "error": raise ValueError if exists. + - "skip" : print and return False. + - "replace": remove existing (in-memory + JSON), then continue. + """ + exists = job_in_store(job_id) or job_in_scheduler(job_id) + if not exists: + return True + if on_conflict == "error": + raise ValueError(f"Job id '{job_id}' already exists.") + elif on_conflict == "skip": + logger.info(f"Job '{job_id}' already exists. Skipping creation.") + return False + elif on_conflict == "replace": + # Clear from scheduler + schedule.clear(job_id) + # Remove from persistence + jobs = [j for j in load_jobs() if j.get("id") != job_id] + save_jobs(jobs) + return True + else: + raise ValueError(f"Unsupported on_conflict policy: {on_conflict}") + + +# ------------------------------- +# Internal scheduling (no persistence) +# ------------------------------- + + +def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None): + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + """Executes the job once, then removes it.""" + if func_name not in FUNCTION_MAP: + logger.error(f"Function '{func_name}' is not registered.") + return + FUNCTION_MAP[func_name](*args, **kwargs) + # Remove from persistence + jobs = load_jobs() + jobs = [j for j in jobs if j["id"] != job_id] + save_jobs(jobs) + # Clear from in-memory schedule + schedule.clear(job_id) + + delay_seconds = run_at_timestamp - time.time() + if delay_seconds <= 0: + logger.info(f"Job {job_id} scheduled in the past. Skipping.") + return + # Schedule via schedule library + schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id) + + +def _schedule_recurring( + job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None +): + args = args or [] + kwargs = kwargs or {} + + def job_wrapper(): + if func_name not in FUNCTION_MAP: + logger.error(f"Function '{func_name}' is not registered.") + return + FUNCTION_MAP[func_name](*args, **kwargs) + + if unit == "seconds": + schedule.every(interval).seconds.do(job_wrapper).tag(job_id) + elif unit == "minutes": + schedule.every(interval).minutes.do(job_wrapper).tag(job_id) + elif unit == "hours": + schedule.every(interval).hours.do(job_wrapper).tag(job_id) + elif unit == "days": + schedule.every(interval).days.do(job_wrapper).tag(job_id) + else: + raise ValueError(f"Unsupported unit: {unit}") + + +# ------------------------------- +# Public APIs (with uniqueness + persistence) +# ------------------------------- + + +def run_once_job( + job_id: str, + func_name: str, + run_at_timestamp: float, + args=None, + kwargs=None, + *, + replace: bool = False, + persist: bool = True, +): + """ + Schedule a job to run once at a specific timestamp. + replace: if True, replace existing job with same id; otherwise print and skip. + persist: if False, do not write to JSON (used by reload_jobs()). + """ + if persist: + policy = "replace" if replace else "skip" + if not ensure_unique(job_id, on_conflict=policy): + return + elif job_in_scheduler(job_id): + schedule.clear(job_id) + _schedule_once(job_id, func_name, run_at_timestamp, args, kwargs) + if persist: + jobs = [j for j in load_jobs() if j["id"] != job_id] + jobs.append( + { + "id": job_id, + "type": "once", + "run_at": run_at_timestamp, + "function": func_name, + "args": args or [], + "kwargs": kwargs or {}, + } + ) + save_jobs(jobs) + + +def recurring_job( + job_id: str, + func_name: str, + interval: int, + unit: str, + args=None, + kwargs=None, + *, + replace: bool = False, + persist: bool = True, +): + """ + Schedule a recurring job. + replace: if True, replace existing job with same id; otherwise print and skip. + persist: if False, do not write to JSON (used by reload_jobs()). + """ + if persist: + policy = "replace" if replace else "skip" + if not ensure_unique(job_id, on_conflict=policy): + return + elif job_in_scheduler(job_id): + schedule.clear(job_id) + _schedule_recurring(job_id, func_name, interval, unit, args, kwargs) + if persist: + jobs = [j for j in load_jobs() if j["id"] != job_id] + jobs.append( + { + "id": job_id, + "type": "recurring", + "interval": interval, + "unit": unit, + "function": func_name, + "args": args or [], + "kwargs": kwargs or {}, + } + ) + save_jobs(jobs) + + +def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0): + """ + Find all jobs whose ID contains the given PID substring and reschedule them to run sooner. + """ + jobs = load_jobs() + matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")] + + if not matched_jobs: + logger.info(f"No jobs found containing PID substring '{pid_substring}'.") + return + + logger.info(f"Found {len(matched_jobs)} job(s) containing '{pid_substring}':") + + for job in matched_jobs: + job_id = job["id"] + logger.debug(f" - Prioritizing job: {job_id}") + + # Clear existing job from scheduler + schedule.clear(job_id) + + # Reschedule based on job type + if job["type"] == "once": + run_once_job( + job_id, + job["function"], + time.time() + new_delay_seconds, + job.get("args"), + job.get("kwargs"), + replace=True, + persist=True, + ) + elif job["type"] == "recurring": + recurring_job( + job_id, + job["function"], + job["interval"], + job["unit"], + job.get("args"), + job.get("kwargs"), + replace=True, + persist=True, + ) + else: + logger.warning(f"Unknown job type for job '{job_id}'") + + +# ------------------------------- +# Reload Saved Jobs +# ------------------------------- + + +def reload_jobs(): + """Reload jobs from JSON and reschedule them (no re-persist).""" + jobs = load_jobs() + for job in jobs: + if job["type"] == "once": + if job["run_at"] > time.time(): + run_once_job( + job["id"], + job["function"], + job["run_at"], + job.get("args"), + job.get("kwargs"), + persist=False, + ) + elif job["type"] == "recurring": + recurring_job( + job["id"], + job["function"], + job["interval"], + job["unit"], + job.get("args"), + job.get("kwargs"), + persist=False, + ) + + +# ------------------------------- +# Scheduler Loop +# ------------------------------- + + +def start_scheduler(): + """ + Start the scheduler loop (blocking). + Call this once in main to begin. + """ + try: + while True: + schedule.run_pending() + time.sleep(0.5) + except KeyboardInterrupt: + logger.critical("Scheduler stopped.") diff --git a/Development/WIP/Server/scheduler_async.py b/Development/WIP/Server/scheduler_async.py new file mode 100644 index 0000000..e20c123 --- /dev/null +++ b/Development/WIP/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/Development/WIP/Server/schema-v1.sql b/Development/WIP/Server/schema-v1.sql new file mode 100644 index 0000000..b061e34 --- /dev/null +++ b/Development/WIP/Server/schema-v1.sql @@ -0,0 +1,95 @@ + +-- Current state of agents +CREATE TABLE IF NOT EXISTS agents ( + agentid TEXT PRIMARY KEY, + clientversion TEXT, + domain TEXT, + freespace INTEGER, + groupid TEXT, + hostname TEXT, + ip TEXT, + localip TEXT, + lastcheckin TEXT, + os TEXT, + policyversion TEXT, + status INTEGER, + username TEXT, + groupname TEXT, + status_text TEXT, + firstseen TEXT +); + +-- History of clientversion changes +CREATE TABLE IF NOT EXISTS agent_clientversion_history ( + agentid TEXT, + clientversion TEXT, + timestamp TEXT, + PRIMARY KEY (agentid, clientversion, timestamp) +); + +-- History of IP changes +CREATE TABLE IF NOT EXISTS agent_ip_history ( + agentid TEXT, + ip TEXT, + timestamp TEXT, + PRIMARY KEY (agentid, ip, timestamp) +); + +-- History of local IP changes +CREATE TABLE IF NOT EXISTS agent_localip_history ( + agentid TEXT, + localip TEXT, + timestamp TEXT, + PRIMARY KEY (agentid, localip, timestamp) +); + +-- History of policyversion changes +CREATE TABLE IF NOT EXISTS agent_policyversion_history ( + agentid TEXT, + policyversion TEXT, + timestamp TEXT, + PRIMARY KEY (agentid, policyversion, timestamp) +); + + +-- Current state of policies +CREATE TABLE IF NOT EXISTS policies ( + groupid TEXT PRIMARY KEY, + hidden BOOLEAN, + name TEXT, + parent TEXT, + firstseen TEXT +); + +-- History of policy name changes +CREATE TABLE IF NOT EXISTS policy_name_history ( + groupid TEXT, + name TEXT, + timestamp TEXT, + PRIMARY KEY (groupid, name, timestamp) +); + + +-- Current state of allowlists +CREATE TABLE IF NOT EXISTS allowlists ( + applicationid TEXT PRIMARY KEY, + name TEXT, + version TEXT, + firstseen TEXT +); + +-- History of allowlist name changes +CREATE TABLE IF NOT EXISTS allowlist_name_history ( + applicationid TEXT, + name TEXT, + timestamp TEXT, + PRIMARY KEY (applicationid, name, timestamp) +); + +-- History of allowlist version changes +CREATE TABLE IF NOT EXISTS allowlist_version_history ( + applicationid TEXT, + version TEXT, + timestamp TEXT, + PRIMARY KEY (applicationid, version, timestamp) +); diff --git a/Development/WIP/Server/sqlhandler.py b/Development/WIP/Server/sqlhandler.py new file mode 100644 index 0000000..957e3f6 --- /dev/null +++ b/Development/WIP/Server/sqlhandler.py @@ -0,0 +1,232 @@ +import sqlite3 +import logging +from models.agent import Agent +from models.policy import Policy, Allowlist +from datetime import datetime +from typing import List +logger = logging.getLogger(__name__) + + +# ------------------ Database Initialization ------------------ + +def initialize_db(db_path: str, schema_path: str, schema_version: str = "1.0"): + try: + with open(schema_path, 'r') as f: + schema_sql = f.read() + + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.executescript(schema_sql) + + cursor.execute(""" + CREATE TABLE IF NOT EXISTS schema_version ( + version TEXT, + applied_on TEXT + ) + """) + cursor.execute(""" + INSERT INTO schema_version (version, applied_on) + VALUES (?, ?) + """, (schema_version, datetime.now().isoformat())) + + conn.commit() + conn.close() + logger.info(f"Tracking database initialized at {db_path} with schema version {schema_version}") + except Exception as e: + logger.error(f"Error initializing database: {e}") + +# ------------------ Change Tracking ------------------ + +def compare_and_track_changes( + conn, + current_agents: List[Agent], + current_policies: List[Policy], + current_allowlists: List[Allowlist] +): + try: + cursor = conn.cursor() + + # --- Agents --- + cursor.execute("SELECT * FROM agents") + agent_rows = cursor.fetchall() + agent_columns = [desc[0] for desc in cursor.description] + previous_agents = { + row[agent_columns.index("agentid")]: Agent(**dict(zip(agent_columns, row))) + for row in agent_rows + } + + for agent in current_agents: + if agent.agentid not in previous_agents: + insert_new_agent(conn, agent) + else: + log_agent_changes(conn, previous_agents[agent.agentid], agent) + + # --- Policies --- + cursor.execute("SELECT * FROM policies") + policy_rows = cursor.fetchall() + policy_columns = [desc[0] for desc in cursor.description] + previous_policies = { + row[policy_columns.index("groupid")]: Policy(**dict(zip(policy_columns, row))) + for row in policy_rows + } + + for policy in current_policies: + if policy.groupid not in previous_policies: + insert_new_policy(conn, policy) + else: + log_policy_changes(conn, previous_policies[policy.groupid], policy) + + # --- Allowlists --- + cursor.execute("SELECT * FROM allowlists") + allowlist_rows = cursor.fetchall() + allowlist_columns = [desc[0] for desc in cursor.description] + previous_allowlists = { + row[allowlist_columns.index("applicationid")]: Allowlist(**dict(zip(allowlist_columns, row))) + for row in allowlist_rows + } + + for allowlist in current_allowlists: + if allowlist.applicationid not in previous_allowlists: + insert_new_allowlist(conn, allowlist) + else: + log_allowlist_changes(conn, previous_allowlists[allowlist.applicationid], allowlist) + + conn.commit() + logger.info("Change tracking completed successfully.") + except Exception as e: + logger.error(f"Error during change tracking: {e}") + +# ------------------ Insert and Log Functions ------------------ + +def insert_new_agent(conn, agent: Agent): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + cursor.execute(""" + INSERT INTO agents ( + agentid, clientversion, domain, freespace, groupid, hostname, ip, localip, + lastcheckin, os, policyversion, status, username, groupname, status_text, firstseen + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(agentid) DO NOTHING + """, ( + agent.agentid, agent.clientversion, agent.domain, agent.freespace, agent.groupid, + agent.hostname, agent.ip, agent.localip, agent.lastcheckin, agent.os, + agent.policyversion, agent.status, agent.username, agent.groupname, agent.status_text, + timestamp + )) + + for field, value in [ + ("clientversion", agent.clientversion), + ("ip", agent.ip), + ("localip", agent.localip), + ("policyversion", agent.policyversion), + ("hostname", agent.hostname), + ("status", agent.status) + ]: + cursor.execute(f""" + INSERT INTO agent_{field}_history (agentid, {field}, timestamp) + VALUES (?, ?, ?) + """, (agent.agentid, value, timestamp)) + + conn.commit() + logger.info(f"Inserted new agent: {agent.agentid}") + except Exception as e: + logger.error(f"Error inserting new agent {agent.agentid}: {e}") + +def log_agent_changes(conn, previous: Agent, current: Agent): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + for field in ["clientversion", "ip", "localip", "policyversion", "hostname", "status"]: + if getattr(previous, field) != getattr(current, field): + cursor.execute(f""" + INSERT INTO agent_{field}_history (agentid, {field}, timestamp) + VALUES (?, ?, ?) + """, (current.agentid, getattr(current, field), timestamp)) + + cursor.execute(""" + UPDATE agents SET lastcheckin = ? WHERE agentid = ? + """, (current.lastcheckin, current.agentid)) + + conn.commit() + logger.info(f"Logged changes for agent: {current.agentid}") + except Exception as e: + logger.error(f"Error logger changes for agent {current.agentid}: {e}") + +def insert_new_policy(conn, policy: Policy): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + cursor.execute(""" + INSERT INTO policies (groupid, hidden, name, parent, firstseen) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(groupid) DO NOTHING + """, (policy.groupid, policy.hidden, policy.name, policy.parent, timestamp)) + + cursor.execute(""" + INSERT INTO policy_name_history (groupid, name, timestamp) + VALUES (?, ?, ?) + """, (policy.groupid, policy.name, timestamp)) + + conn.commit() + logger.info(f"Inserted new policy: {policy.groupid}") + except Exception as e: + logger.error(f"Error inserting new policy {policy.groupid}: {e}") + +def log_policy_changes(conn, previous: Policy, current: Policy): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + if previous.name != current.name: + cursor.execute(""" + INSERT INTO policy_name_history (groupid, name, timestamp) + VALUES (?, ?, ?) + """, (current.groupid, current.name, timestamp)) + + conn.commit() + logger.info(f"Logged changes for policy: {current.groupid}") + except Exception as e: + logger.error(f"Error logger changes for policy {current.groupid}: {e}") + +def insert_new_allowlist(conn, allowlist: Allowlist): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + cursor.execute(""" + INSERT INTO allowlists (applicationid, name, version, firstseen) + VALUES (?, ?, ?, ?) + ON CONFLICT(applicationid) DO NOTHING + """, (allowlist.applicationid, allowlist.name, allowlist.version, timestamp)) + + for field, value in [("name", allowlist.name), ("version", allowlist.version)]: + cursor.execute(f""" + INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp) + VALUES (?, ?, ?) + """, (allowlist.applicationid, value, timestamp)) + + conn.commit() + logger.info(f"Inserted new allowlist: {allowlist.applicationid}") + except Exception as e: + logger.error(f"Error inserting new allowlist {allowlist.applicationid}: {e}") + +def log_allowlist_changes(conn, previous: Allowlist, current: Allowlist): + try: + timestamp = datetime.now().isoformat() + cursor = conn.cursor() + + for field in ["name", "version"]: + if getattr(previous, field) != getattr(current, field): + cursor.execute(f""" + INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp) + VALUES (?, ?, ?) + """, (current.applicationid, getattr(current, field), timestamp)) + + conn.commit() + logger.info(f"Logged changes for allowlist: {current.applicationid}") + except Exception as e: + logger.error(f"Error logger changes for allowlist {current.applicationid}: {e}") \ No newline at end of file diff --git a/Development/WIP/Sync Code/API.py b/Development/WIP/Sync Code/API.py new file mode 100644 index 0000000..e6b8897 --- /dev/null +++ b/Development/WIP/Sync Code/API.py @@ -0,0 +1,307 @@ +# 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) + + def policy_set_script_custom( + self, + groupid: str, + script_custom: int, + scripts_audit: Optional[List[str]] = None, + scripts_disabled: Optional[List[str]] = None, + scripts_respect: Optional[List[str]] = None + ) -> dict: + """ + Set the script control behavior for a policy group. + + Parameters: + - groupid (str): Target Group ID. + - script_custom (int): 0 = Disabled, 1 = Enabled. + - scripts_audit (List[str], optional): Script types to audit. + - scripts_disabled (List[str], optional): Script types to disable. + - scripts_respect (List[str], optional): Script types to respect current policy. + + Returns: + - dict: JSON response from the API. + + Available script types are "batch","powershell","command","vbscript","javascript","windowsinstaller","htmlapplication","javaapplication","windowsscriptcomponent","compiledhtml","shellscript","dylib","python" + + """ + payload = { + "groupid": groupid, + "script_custom": script_custom, + "scripts_audit": scripts_audit or [], + "scripts_disabled": scripts_disabled or [], + "scripts_respect": scripts_respect or [] + } + return self._post("/v1/group/settings/script_custom", 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/Development/WIP/Sync Code/AirlockTools_Client.py b/Development/WIP/Sync Code/AirlockTools_Client.py new file mode 100644 index 0000000..829b262 --- /dev/null +++ b/Development/WIP/Sync Code/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 + + if username: + api = AirlockAPIWrapper( + base_url=str(os.getenv("URL")), + api_key = getAPI(username, "AirlockTools"), # pyright: ignore[reportArgumentType] + ) + + + menus.menu_main(api) + + +if __name__ == "__main__": + main() diff --git a/Development/WIP/Sync Code/agenthandler.py b/Development/WIP/Sync Code/agenthandler.py new file mode 100644 index 0000000..290baba --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/configmanager.py b/Development/WIP/Sync Code/configmanager.py new file mode 100644 index 0000000..56c8f3c --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/execution.py b/Development/WIP/Sync Code/execution.py new file mode 100644 index 0000000..a7466f6 --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/menus.py b/Development/WIP/Sync Code/menus.py new file mode 100644 index 0000000..01295a6 --- /dev/null +++ b/Development/WIP/Sync Code/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[0]) + 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/Development/WIP/Sync Code/otp.py b/Development/WIP/Sync Code/otp.py new file mode 100644 index 0000000..4f5d52b --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/policyhandler.py b/Development/WIP/Sync Code/policyhandler.py new file mode 100644 index 0000000..f76b969 --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/prepPolicy.py b/Development/WIP/Sync Code/prepPolicy.py new file mode 100644 index 0000000..249edc8 --- /dev/null +++ b/Development/WIP/Sync Code/prepPolicy.py @@ -0,0 +1,365 @@ +# 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, Optional + +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], + history_days: Optional[int] = None +): + + + if history_days is None: + 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 = [] + working_dir = load_env("WORKING_DIR") + # 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\\{selected_policies[0].name}_{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/Development/WIP/Sync Code/quietAgent.py b/Development/WIP/Sync Code/quietAgent.py new file mode 100644 index 0000000..1a8d79b --- /dev/null +++ b/Development/WIP/Sync Code/quietAgent.py @@ -0,0 +1,105 @@ +# 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") + selected_policy = selectPolicies(api, False) + if selected_policy: + agents = api.agents_find_by_group(selected_policy[0].groupid) + + 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โ€“150): ", + value_type=int, + valid_range=(1, 150), + ) + + 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 + + policy_exec_history.loc[:, "datetime"] = pd.to_datetime( + policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True + ) + + now = datetime.datetime.now(datetime.timezone.utc) + + policy_exec_history.loc[:, "days_ago"] = policy_exec_history["datetime"].apply( + lambda dt: (now - dt).days + ) + + hostname_counts = policy_exec_history["hostname"].value_counts() + + agents.loc[:, "execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int) + + most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates( + subset="hostname", keep="first" + ) + + agents.loc[:, "days_since"] = agents["hostname"].map( + most_recent_exec.set_index("hostname")["days_ago"] + ) + + agents.loc[:, "required_quiet"] = required_quiet + agents.loc[:, "enforce_ready"] = agents["days_since"].apply( + lambda x: True if pd.isna(x) or x > required_quiet else False + ) + + agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True]) + + 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) + + total_agents = len(agents) + ready_agents = agents["enforce_ready"].sum() + not_ready_agents = total_agents - ready_agents + ready_percentage = (ready_agents / total_agents) * 100 + + 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.info(message) + colorText(message, "green") diff --git a/Development/WIP/Sync Code/security.py b/Development/WIP/Sync Code/security.py new file mode 100644 index 0000000..841902e --- /dev/null +++ b/Development/WIP/Sync Code/security.py @@ -0,0 +1,159 @@ +# 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 +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 +from utils.utils import colorText +from sys import exit + +# 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() + + +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.") + print(colorText("โŒ Authentication failed. Exiting.","red")) + exit(1) # Exit cleanly without traceback + 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/Development/WIP/Sync Code/selector.py b/Development/WIP/Sync Code/selector.py new file mode 100644 index 0000000..9385eb8 --- /dev/null +++ b/Development/WIP/Sync Code/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 = 3, + 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/Development/WIP/Sync Code/setup.py b/Development/WIP/Sync Code/setup.py new file mode 100644 index 0000000..d2d70ef --- /dev/null +++ b/Development/WIP/Sync Code/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/Development/WIP/Sync Code/utils.py b/Development/WIP/Sync Code/utils.py new file mode 100644 index 0000000..6803411 --- /dev/null +++ b/Development/WIP/Sync Code/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}") + + diff --git a/Development/WIP/dynamicAsyncTaskQueue.py b/Development/WIP/dynamicAsyncTaskQueue.py new file mode 100644 index 0000000..4f847e9 --- /dev/null +++ b/Development/WIP/dynamicAsyncTaskQueue.py @@ -0,0 +1,135 @@ +import asyncio +import logging +from typing import Callable, Any +import psutil # For CPU and memory monitoring + +# Set up the logger +logger = logging.getLogger("AsyncTaskQueue") +logger.setLevel(logging.INFO) +handler = logging.StreamHandler() +formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') +handler.setFormatter(formatter) +logger.addHandler(handler) + +class AsyncTaskQueue: + def __init__( + self, + min_workers: int = 1, + max_workers: int = 10, + cpu_threshold: float = 90.0, + memory_threshold: float = 90.0, + ): + self.queue = asyncio.Queue() + self.min_workers = min_workers + self.max_workers = max_workers + self.cpu_threshold = cpu_threshold + self.memory_threshold = memory_threshold + self.workers = [] + self._stop_event = asyncio.Event() + self._scale_lock = asyncio.Lock() # Prevent race conditions + + async def _scale_up(self): + """Add a worker if under max limit and resources allow.""" + async with self._scale_lock: + if ( + len(self.workers) < self.max_workers + and psutil.cpu_percent() < self.cpu_threshold + and psutil.virtual_memory().percent < self.memory_threshold + ): + worker = asyncio.create_task(self.worker_loop(f"Worker-{len(self.workers) + 1}")) + self.workers.append(worker) + logger.info(f"Scaled up. Workers: {len(self.workers)}") + + async def _scale_down(self): + """Remove a worker if above min limit.""" + async with self._scale_lock: + if len(self.workers) > self.min_workers: + worker = self.workers.pop() + worker.cancel() + logger.info(f"Scaled down. Workers: {len(self.workers)}") + + async def _monitor_resources(self): + """Monitor CPU and memory usage, and adjust workers.""" + while not self._stop_event.is_set(): + cpu_usage = psutil.cpu_percent() + memory_usage = psutil.virtual_memory().percent + + if ( + cpu_usage > self.cpu_threshold + or memory_usage > self.memory_threshold + ): + await self._scale_down() + elif ( + len(self.workers) < self.max_workers + and self.queue.qsize() > 2 # Only scale up if there's work + ): + await self._scale_up() + + await asyncio.sleep(2) # Check every 2 seconds + + async def start_workers(self): + """Start initial workers and the resource monitor.""" + for i in range(self.min_workers): + worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}")) + self.workers.append(worker) + # Start the resource monitor + asyncio.create_task(self._monitor_resources()) + logger.info(f"Started {self.min_workers} workers and resource monitor.") + + async def stop_workers(self): + """Stop all workers and the resource monitor.""" + logger.info("Stopping workers...") + self._stop_event.set() + await self.queue.join() # Wait for all tasks to complete + for worker in self.workers: + worker.cancel() + await asyncio.gather(*self.workers, return_exceptions=True) + logger.info("All workers stopped.") + + async def worker_loop(self, name: str): + """Worker loop: Process tasks from the queue.""" + logger.info(f"{name} started.") + while not self._stop_event.is_set(): + try: + task = await self.queue.get() + logger.info(f"{name} processing: {task['name']}") + await task['func'](*task['args']) + except Exception as e: + logger.error(f"Error in {name}: {e}", exc_info=True) + finally: + self.queue.task_done() + logger.info(f"{name} exited.") + + async def enqueue(self, name: str, func: Callable, *args: Any): + """Add a task to the queue.""" + logger.info(f"Enqueuing task: {name}") + await self.queue.put({'name': name, 'func': func, 'args': args}) + +async def run_sync_task_in_thread(func: Callable, *args: Any): + """Run a synchronous function in a separate thread.""" + await asyncio.to_thread(func, *args) +"""import asyncio + +async def example_async_task(name: str, duration: int): + print(f"{name} started, will sleep for {duration} seconds") + await asyncio.sleep(duration) + print(f"{name} finished") + +async def main(): + queue = AsyncTaskQueue( + min_workers=2, + max_workers=10, + cpu_threshold=90.0, + memory_threshold=90.0, + ) + await queue.start_workers() + + # Enqueue tasks + for i in range(20): + await queue.enqueue(f"Task{i}", example_async_task, f"Task{i}", 1) + + await asyncio.sleep(10) # Let tasks run + await queue.stop_workers() + +asyncio.run(main()) +""" \ No newline at end of file diff --git a/Development/WIP/gu2i.py b/Development/WIP/gu2i.py new file mode 100644 index 0000000..758690b --- /dev/null +++ b/Development/WIP/gu2i.py @@ -0,0 +1,179 @@ +from textual.app import App, ComposeResult +from textual.screen import Screen +from textual.widgets import Header, Tabs, Tab, Static, Footer, DirectoryTree, Button +from textual.containers import Horizontal +import logging +import dotenv +import os + +from Development_Stubs.WIP.localApproval import moveToLocalApproval +from flows.prepPolicy import ( + buildPathsandPublishers, + buildPreflights, + selectAllowlists, + selectPolicies, + sortHashes, +) +from flows.quietAgent import findQuietAgents +from services.agenthandler import findAgents +from services.API import AirlockAPIWrapper +import services.policyhandler as policyh +from utils.utils import ( + areYouSure, + colorText, + displayIntro, + load_env, + open_directory, + printEnforceChecklist, +) +from screens.agent_results import AgentResultsScreen +logger = logging.getLogger(__name__) +dotenv.load_dotenv() + +ASCII_ART = displayIntro() + +class BaseScreen(Screen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("This is a dummy screen.", id="content") + yield Footer() + +class LandingScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static(ASCII_ART, id="ascii-art", markup=False) + yield Horizontal( + Button("Get Started", id="get-started-button"), + Button("Settings", id="settings-button"), + ) + yield Footer() + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "get-started-button": + # Go back to the previous screen (MainMenuScreen) + self.app.pop_screen() + +class LocalApprovalScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + + yield Footer() + +class OTPScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("OTP: Enter your one-time password here.", id="content") + yield Footer() + +class SearchScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("Search: Find items in the system.", id="content") + yield Horizontal( + Button("Search for Agents", id="searchAgent-button"), + ) + yield Footer() + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "searchAgent-button": + self.app.push_screen(AgentResultsScreen(self.app.api, self.app.working_dir)) + +class QuietHostsScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("Quiet Hosts: Manage hosts that are not responding.", id="content") + yield Footer() + +class PolicyPrepScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("Policy Prep: Prepare policies for deployment.", id="content") + yield Footer() + +class UpdatePoliciesScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("Update Policies: Update existing policies.", id="content") + yield Footer() + +class SettingsScreen(BaseScreen): + def compose(self) -> ComposeResult: + yield Header() + yield Static("Settings: Configure application settings.", id="content") + yield Footer() + +class DirectoryTreeScreen(Screen): + def compose(self) -> ComposeResult: + yield Header() + yield DirectoryTree("./") + yield Footer() + +class MainMenuScreen(Screen): + def compose(self) -> ComposeResult: + yield Header() + yield Tabs( + Tab("Home", id="home_screen"), + Tab("Local Approval", id="local_approval"), + Tab("OTP", id="otp"), + Tab("Search", id="search"), + Tab("Quiet Hosts", id="quiet"), + Tab("Policy Prep", id="policy"), + Tab("Update Policies", id="update"), + id="tabs" + ) + yield Static("Select a tab to begin.", id="content") + yield Footer() + + def on_tabs_tab_activated(self, event: Tabs.TabActivated) -> None: + match event.tab.id: + case "home_screen": + self.app.push_screen(LandingScreen()) + case "local_approval": + self.app.push_screen(LocalApprovalScreen()) + case "otp": + self.app.push_screen(OTPScreen()) + case "search": + self.app.push_screen(SearchScreen()) + case "quiet": + self.app.push_screen(QuietHostsScreen()) + case "policy": + self.app.push_screen(PolicyPrepScreen()) + case "update": + self.app.push_screen(UpdatePoliciesScreen()) + case "settings": + self.app.push_screen(SettingsScreen()) + +class AirlockTools(App): + BINDINGS = [ + ("q", "quit", "Quit"), + ("d", "open_dir", "Open Directory"), + ("b", "back", "Go Back"), + ] + + def __init__(self, api: AirlockAPIWrapper, working_dir: str): + super().__init__() + self.api = api + self.working_dir = working_dir + + def on_mount(self) -> None: + self.push_screen(MainMenuScreen()) + + def action_quit(self) -> None: + self.exit() + + def action_open_dir(self) -> None: + self.push_screen(DirectoryTreeScreen()) + + def action_back(self) -> None: + if len(self.screen_stack) > 2: + self.pop_screen() + else: + self.bell() + + +if __name__ == "__main__": + api = AirlockAPIWrapper(base_url=str(os.getenv("URL")),api_key = "b82ab7e2b12430fff0cbf0e798f79e16c1aea2c37ffee5dc884de7ba11d3a3a4") + working_dir = "./" # Or a dynamic path + AirlockTools(api, working_dir).run() + + diff --git a/services/security.py b/services/security.py index 093b58d..bbf2bf4 100644 --- a/services/security.py +++ b/services/security.py @@ -32,6 +32,8 @@ SALT_SIZE = 16 # 128-bit Salt NONCE_SIZE = 12 # AES-GCM KEY_SIZE = 32 # AES-256 +logger = logging.getLogger(__name__) + def _derive_key(password: bytes, salt: bytes) -> bytes: kdf = PBKDF2HMAC( @@ -66,6 +68,20 @@ def store_api_key(service: str, username: str, api_key: str, password: str): b64 = base64.b64encode(blob).decode() keyring.set_password(service, username, b64) + + logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.") + + print("\nโœ… API key stored securely.") + print("The program will now exit. Press Enter to continue...") + + try: + _ = input() + except Exception: + pass + + _ = None + sys.exit(0) + def retrieve_api_key(service: str, username: str, password: str) -> str: configure_keyring_backend() @@ -80,20 +96,6 @@ def retrieve_api_key(service: str, username: str, password: str) -> str: 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: