diff --git a/Development/AdminTools/policycloner.py b/Development/AdminTools/policycloner.py deleted file mode 100644 index a4d642a..0000000 --- a/Development/AdminTools/policycloner.py +++ /dev/null @@ -1,46 +0,0 @@ -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 deleted file mode 100644 index beb5f6e..0000000 --- a/Development/AdminTools/policylister.py +++ /dev/null @@ -1,44 +0,0 @@ -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 deleted file mode 100644 index 1f32549..0000000 --- a/Development/AdminTools/scriptPolicySetter.py +++ /dev/null @@ -1,75 +0,0 @@ -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 deleted file mode 100644 index 9af12c9..0000000 --- a/Development/AdminTools/testsetup.py +++ /dev/null @@ -1,39 +0,0 @@ -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 deleted file mode 100644 index df9721f..0000000 --- a/Development/WIP/Server/datastore.py +++ /dev/null @@ -1,170 +0,0 @@ -""" -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 deleted file mode 100644 index 789b917..0000000 --- a/Development/WIP/Server/localApproval.py +++ /dev/null @@ -1,291 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import datetime -import 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 deleted file mode 100644 index a653c3e..0000000 --- a/Development/WIP/Server/scheduler.py +++ /dev/null @@ -1,352 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import 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 deleted file mode 100644 index e20c123..0000000 --- a/Development/WIP/Server/scheduler_async.py +++ /dev/null @@ -1,177 +0,0 @@ -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 deleted file mode 100644 index b061e34..0000000 --- a/Development/WIP/Server/schema-v1.sql +++ /dev/null @@ -1,95 +0,0 @@ - --- 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 deleted file mode 100644 index 957e3f6..0000000 --- a/Development/WIP/Server/sqlhandler.py +++ /dev/null @@ -1,232 +0,0 @@ -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 deleted file mode 100644 index e6b8897..0000000 --- a/Development/WIP/Sync Code/API.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 829b262..0000000 --- a/Development/WIP/Sync Code/AirlockTools_Client.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -#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 deleted file mode 100644 index 290baba..0000000 --- a/Development/WIP/Sync Code/agenthandler.py +++ /dev/null @@ -1,250 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import 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 deleted file mode 100644 index 56c8f3c..0000000 --- a/Development/WIP/Sync Code/configmanager.py +++ /dev/null @@ -1,115 +0,0 @@ -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 deleted file mode 100644 index a7466f6..0000000 --- a/Development/WIP/Sync Code/execution.py +++ /dev/null @@ -1,398 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 01295a6..0000000 --- a/Development/WIP/Sync Code/menus.py +++ /dev/null @@ -1,316 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 4f5d52b..0000000 --- a/Development/WIP/Sync Code/otp.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - - -import 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 deleted file mode 100644 index f76b969..0000000 --- a/Development/WIP/Sync Code/policyhandler.py +++ /dev/null @@ -1,237 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import datetime -import 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 deleted file mode 100644 index 249edc8..0000000 --- a/Development/WIP/Sync Code/prepPolicy.py +++ /dev/null @@ -1,365 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 1a8d79b..0000000 --- a/Development/WIP/Sync Code/quietAgent.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import datetime -import 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 deleted file mode 100644 index 841902e..0000000 --- a/Development/WIP/Sync Code/security.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 9385eb8..0000000 --- a/Development/WIP/Sync Code/selector.py +++ /dev/null @@ -1,171 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import 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 deleted file mode 100644 index d2d70ef..0000000 --- a/Development/WIP/Sync Code/setup.py +++ /dev/null @@ -1,175 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import 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 deleted file mode 100644 index 6803411..0000000 --- a/Development/WIP/Sync Code/utils.py +++ /dev/null @@ -1,708 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -import 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 deleted file mode 100644 index 4f847e9..0000000 --- a/Development/WIP/dynamicAsyncTaskQueue.py +++ /dev/null @@ -1,135 +0,0 @@ -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 deleted file mode 100644 index 758690b..0000000 --- a/Development/WIP/gu2i.py +++ /dev/null @@ -1,179 +0,0 @@ -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() - -