First Round Async. Much work left to do, dont trust results of hash categorization presently.

services/API.py

@@ -1,274 +1,175 @@
-#    Copyright (C) 2025 James Brotosky, Brandon Wickline
-#
-#    This program is free software: you can redistribute it and/or modify
-#    it under the terms of the GNU Affero General Public License as published
-#    by the Free Software Foundation, either version 3 of the License, or
-#    (at your option) any later version.
-#
-#    This program is distributed in the hope that it will be useful,
-#    but WITHOUT ANY WARRANTY; without even the implied warranty of
-#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-#    GNU Affero General Public License for more details.
-#
-#    You should have received a copy of the GNU Affero General Public License
-#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-
 import json
 import logging
 from typing import Dict, List, Optional
 
+import httpx
 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.
-        """
+    async def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
         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
+        timeout = httpx.Timeout(300.0)
+        async with httpx.AsyncClient(verify=False, timeout=timeout) as client:
+            try:
+                logger.debug(f"POST Request to {url} with payload: {payload}")
+                response = await client.post(url, headers=self.headers, data=data) # pyright: ignore[reportArgumentType]
+                response.raise_for_status()
+                logger.debug(f"Response received from {url}")
+                return response.json()
+            except httpx.RequestError 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", {})
+    async def allowlist_find_all(self) -> pd.DataFrame:
+        result = await 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", {})
+    async def agent_find_all(self) -> pd.DataFrame:
+        result = await 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."""
+    async def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
         payload = {"hostname": hostname}
-        result = self._post("/v1/agent/find", payload)
+        result = await 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."""
+    async def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
         payload = {"agentid": agentid}
-        result = self._post("/v1/agent/find", payload)
+        result = await 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)."""
+    async def agent_find_by_status(self, status: int) -> pd.DataFrame:
         payload = {"status": status}
-        result = self._post("/v1/agent/find", payload)
+        result = await 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."""
+    async def agent_find_by_username(self, username: str) -> pd.DataFrame:
         payload = {"username": username}
-        result = self._post("/v1/agent/find", payload)
+        result = await 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."""
+    async def agent_move(self, agentid: str, groupid: str) -> dict:
         payload = {"agentid": agentid, "groupid": groupid}
-        return self._post("/v1/agent/move", payload)
+        return await self._post("/v1/agent/move", payload)
 
-    def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
-        """Find agents by group ID."""
+    async def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
         payload = {"groupid": groupid}
-        result = self._post("/v1/agent/find", payload)
+        result = await 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."""
+    async def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
         payload = {"applicationid": applicationid, "hashes": hashes}
-        return self._post("/v1/hash/application/add", payload)
+        return await self._post("/v1/hash/application/add", payload)
 
-    def hash_query(self, hashes: List[str]) -> pd.DataFrame:
-        """Query information about specific hashes."""
+    async def hash_query(self, hashes: List[str]) -> pd.DataFrame:
         payload = {"hashes": hashes}
-        result = self._post("/v1/hash/query", payload)
+        result = await self._post("/v1/hash/query", payload)
         return pd.DataFrame(result["response"]["results"])
 
     # OTP Management
-    def otp_find_active(self) -> pd.DataFrame:
-        """Find active OTPs."""
+    async def otp_find_active(self) -> pd.DataFrame:
         payload = {"status": "1"}
-        result = self._post("/v1/otp/usage", payload)
+        result = await 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."""
+    async def otp_find_awaiting(self) -> pd.DataFrame:
         payload = {"status": "0"}
-        result = self._post("/v1/otp/usage", payload)
+        result = await 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."""
+
+    async def otp_find_by_agent(self, agentid) -> pd.DataFrame:
         payload = {"agentid": agentid}
-        result = self._post("/v1/otp/usage", payload)
+        result = await 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."""
+
+    async def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
         payload = {
             "duration": str(duration),
             "agentid": str(agentid),
             "purpose": purpose,
         }
-        result = self._post("/v1/otp/retrieve", payload)
+        result = await 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."""
+    async def otp_get_activities(self, otpid: str) -> pd.DataFrame:
         payload = {"otpid": otpid}
-        result = self._post("/v1/otp/activities", payload)
+        result = await 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)
 
+    async def otp_revoke(self, otpid: str) -> dict:
+        payload = {"otpid": otpid}
+        return await self._post("/v1/otp/revoke", payload)
+
+    async def otp_validate(self, otpcode: str) -> dict:
+        payload = {"otpcode": otpcode}
+        return await 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."""
+    async def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
         payload = {"groupid": groupid, "path": paths}
-        return self._post("/v1/group/path/add", payload)
+        return await self._post("/v1/group/path/add", payload)
 
-    def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
-        """Add publishers to a policy group."""
+    async def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
         payload = {"groupid": groupid, "publisher": publishers}
-        return self._post("/v1/group/publisher/add", payload)
+        return await 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."""
+    async def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
         payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
-        return self._post("/v1/group/assign", payload)
+        return await self._post("/v1/group/assign", payload)
 
-    def policy_find_all(self) -> pd.DataFrame:
-        """Retrieve all policy groups."""
-        result = self._post("/v1/group")
+    async def policy_find_all(self) -> pd.DataFrame:
+        result = await 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."""
+    async def policy_list_agents(self, groupid: str) -> pd.DataFrame:
         payload = {"groupid": groupid}
-        result = self._post("/v1/group/agents", payload)
+        result = await 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."""
+
+    async def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
         payload = {"groupid": groupid}
-        result = self._post("/v1/group/policies", payload)
+        result = await 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"""
+    async def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
         payload = {"groupid": groupid, "auditmode": auditmode}
-        return self._post("/v1/group/settings/auditmode", payload)
+        return await self._post("/v1/group/settings/auditmode", payload)
+
+    async 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:
+        payload = {
+            "groupid": groupid,
+            "script_custom": script_custom,
+            "scripts_audit": scripts_audit or [],
+            "scripts_disabled": scripts_disabled or [],
+            "scripts_respect": scripts_respect or []
+        }
+        return await 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."""
+    async def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> List[Dict]:
         payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
-        result = self._post("/v1/logging/exechistories", payload)
+        result = await 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
-
-        """
-
+    async def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]:
         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)
-"""
+        result = await self._post("/v1/getexechistory", payload)
+        return result["response"]["exechistory"]

services/TaskQueue.py

@@ -0,0 +1,108 @@
+import asyncio
+import logging
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncTaskQueue:
+    def __init__(self, worker_count: int = 3):
+        self.queue = asyncio.Queue()
+        self.worker_count = worker_count
+        self.workers = []
+        self._stop_event = asyncio.Event()
+
+    async def start_workers(self):
+        """Start the worker pool."""
+        logger.debug(f"Starting {self.worker_count} workers...")
+        for i in range(self.worker_count):
+            worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}"))
+            self.workers.append(worker)
+        logger.debug("All workers started.")
+
+    async def stop_workers(self):
+        """Stop the worker pool and wait for all tasks to complete."""
+        logger.debug("Stopping workers...")
+        self._stop_event.set()  # Signal workers to stop
+        await self.queue.join()  # Wait for all tasks to be processed
+        for worker in self.workers:
+            worker.cancel()
+        await asyncio.gather(*self.workers, return_exceptions=True)
+        logger.debug("All workers stopped.")
+
+
+    async def worker_loop(self, name: str):
+        """Worker loop: Process tasks from the queue."""
+        logger.debug(f"{name} started.")
+        while not self._stop_event.is_set():
+            task = None
+            try:
+                task = await self.queue.get()
+                logger.info(f"{name} processing: {task['name']}")
+                await task['func'](*task['args'], **task['kwargs'])            
+            except asyncio.CancelledError:
+                logger.debug(f"{name} received cancellation.")
+                break
+            except Exception as e:
+                logger.error(f"Error in {name}: {e}", exc_info=True)
+            finally:
+                if task is not None:
+                    self.queue.task_done()
+        logger.debug(f"{name} exited.")
+
+
+
+    async def enqueue(self, name: str, func: Callable, *args: Any, **kwargs: Any):
+        """Add a task to the queue."""
+        logger.debug(f"Enqueuing task: {name}")
+        await self.queue.put({'name': name, 'func': func, 'args': args, 'kwargs': kwargs})
+
+
+
+async def run_sync_task_in_thread(func: Callable, *args: Any, **kwargs: Any):
+    """Run a synchronous function in a separate thread.
+
+    Args:
+        func: The synchronous function to run.
+        *args: Arguments to pass to the function.
+    """
+    await asyncio.to_thread(func, *args, **kwargs)
+    
+    """
+from asyncTaskQueue import AsyncTaskQueue, run_sync_task_in_thread
+import asyncio
+
+# Example async task
+async def your_async_function(name: str, duration: int):
+    print(f"{name} started, will sleep for {duration} seconds")
+    await asyncio.sleep(duration)
+    print(f"{name} finished")
+
+# Example sync task
+def your_sync_function(name: str, duration: int):
+    print(f"{name} started, will sleep for {duration} seconds")
+    import time
+    time.sleep(duration)
+    print(f"{name} finished")
+
+async def main():
+    queue = AsyncTaskQueue(worker_count=2)
+    await queue.start_workers()
+
+    # Enqueue async tasks directly
+    await queue.enqueue("AsyncTask1", your_async_function, "AsyncTask1", 2)
+    await queue.enqueue("AsyncTask2", your_async_function, "AsyncTask2", 1)
+
+    # Enqueue sync tasks using the wrapper
+    await queue.enqueue("SyncTask1", run_sync_task_in_thread, your_sync_function, "SyncTask1", 2)
+
+    # Let tasks run for a while
+    await asyncio.sleep(5)
+
+    # Stop workers
+    await queue.stop_workers()
+
+asyncio.run(main())
+
+
+"""

services/agenthandler.py

@@ -1,19 +1,4 @@
-#    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 <https://www.gnu.org/licenses/>.
-
-
+import asyncio
 import json
 import logging
 import os
@@ -27,16 +12,16 @@ import pandas as pd
 from models.agent import Agent
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
+from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread
 from utils.configmanager import get_protected_json, load_env
-from utils.selector import Selector
+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(
+async def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
+    agents = await selectAgents(api)
+    history_days = await Selector.select_value(
         prompt="Enter how many days of history to pull (1–150): ",
         value_type=int,
         valid_range=(1, 150),
@@ -48,63 +33,68 @@ def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
 
     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:
+    async def fetch_history(agent):
         try:
-            exechistory = api.history_execution(today, historical_date, agent.hostname)
+            exechistory = await api.history_execution(today, historical_date, agent.hostname)
+            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"))
         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"))
+    await asyncio.gather(*(fetch_history(agent) for agent in agents))
 
     if outputjson:
         print(json.dumps(all_history, indent=2))
 
+async def findAllAgents(api: AirlockAPIWrapper):
+    policies_df = await api.policy_find_all()
+    agents_df = await api.agent_find_all()
 
-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()]
+    policies = [Policy(**row["data"]) for _, row in policies_df.iterrows()]
+    agents = [Agent(**row["data"]) for _, row in agents_df.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:
+    queue = AsyncTaskQueue()
+    await queue.start_workers()
+
+    async def enrich_agent(agent):
         agent.enrich(groupid_to_name)
 
+    for agent in agents:
+        await queue.enqueue(f"enrich_{agent.hostname}", enrich_agent, agent)
+
+    await asyncio.sleep(1)
+    await queue.stop_workers()
+
     return agents
 
-def findAgents(api, return_dataframe):
-    agents = selectAgents(api)
-    working_dir = load_env("WORKING_DIR")
+async def findAgents(api: AirlockAPIWrapper, return_dataframe: bool):
+    agents = await selectAgents(api)
+    working_dir = await 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)
 
@@ -112,30 +102,24 @@ def findAgents(api, 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()
+    user_input = await get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ")
+    user_input = user_input.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}")
+        await run_sync_task_in_thread(agent_df.to_csv, file_path, index=False)
 
-        print(
-            colorText(
-                f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
-                "green",
-            )
-        )
+        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]:
+async 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"))
@@ -144,50 +128,44 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
     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()]
+
+    policies_df = await api.policy_find_all()
+    policies = [Policy(**row.to_dict()) for _, row in policies_df.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("")
+        line = await 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
+            continue
         else:
             empty_line_count = 0
+            if valid_line_pattern.match(stripped_line):
+                device_input_lines.append(stripped_line)
+            else:
+                print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow"))
 
-        # 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))
+    pattern = "\n".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()]
+    agents_df = await api.agent_find_all()
+    agents = [Agent(**row.to_dict()) for _, row in agents_df.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)}")
@@ -197,30 +175,17 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
         logger.debug("❌ No matching devices found.")
         print(colorText("❌ No matching devices found.", "red"))
     else:
-        logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")        
+        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", "{}")
 
+async def moveAgentToRelatedPolicy(api: AirlockAPIWrapper, agent: Agent, mode: str = "audit"):
+    policy_relationship_map = await get_protected_json("POLICY_MAP_ENF_AUD", "{}")
+    #TODO - Have this return the policy name it was moved to instead of the groupid
     if mode == "audit":
         if agent.groupid in policy_relationship_map:
             target_policy = policy_relationship_map[agent.groupid]
@@ -231,7 +196,6 @@ def moveAgentToRelatedPolicy(
         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:
@@ -242,9 +206,11 @@ def moveAgentToRelatedPolicy(
         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)
+    result = await api.agent_move(agent.agentid, target_policy)
+    if result == {'error': 'Success'}: logger.info(f"{agent.hostname} has been moved to {target_policy}")
+    else: logger.debug(result)
+        

services/policyhandler.py

@@ -1,29 +1,15 @@
-#    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 <https://www.gnu.org/licenses/>.
-
-
+import asyncio
 import datetime
 import gc
 import json
 import logging
 import os
-import sys
+import uuid
 
+import aiofiles
 import pandas as pd
-import tqdm
 from bson import ObjectId
+from tqdm.asyncio import tqdm_asyncio
 
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
@@ -35,181 +21,101 @@ 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.")
-
+async def pullPolicyExechistories(api: AirlockAPIWrapper, policy: Policy, type, days, outputjson):
+    file_path = f"{get_base_directory()}\\cache\\chunkinator_{policy.name}_{uuid.uuid4().hex}.json"
     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]
-                )
+    if not os.path.exists(file_path):
+        async with aiofiles.open(file_path, "w") as file:
+            await file.write(json.dumps(json_output))
 
-                # 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 = tqdm_asyncio(total=10000, desc=f"Checkpoint Progress: {checkpoint}", colour="blue")
+    pbar = tqdm_asyncio(total=100, desc=f"Total of {policy.name} Complete: ")
 
-                filebar.total = len(histories)
+    while True:
+        histories = await api.history_logging(type=type, checkpoint=checkpoint, policy=[policy.name])
+        if not histories:
+            break
 
-                if not histories:
-                    break
+        for index, history_item in enumerate(histories):
+            if "checkpoint" not in history_item or "datetime" not in history_item:
+                continue
 
-                for index, history_item in enumerate(histories):
-                    if (
-                        "checkpoint" not in history_item
-                        or "datetime" not in history_item
-                    ):
-                        continue  # Skip malformed entries
+            if index == len(histories) - 1:
+                checkpoint = history_item["checkpoint"]
+                filebar.set_description(f"Checkpoint Progress: {checkpoint}")
+                break
 
-                    # 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", ""),
+                    "%Y-%m-%dT%H:%M:%SZ"
+                ).date()
+            except ValueError:
+                continue
 
-                    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)
 
-                    if (
-                        datetime.date.today() - datetime.timedelta(days=days)
-                    ) <= history_date:
-                        json_output["response"]["exechistories"].append(history_item)
+            filebar.update(1)
+            await asyncio.sleep(0)
 
-                    filebar.update(1)
-                    filebar.refresh()
+        # Deduplication
+        seen = {}
+        if os.path.exists(file_path):
+            async with aiofiles.open(file_path, "r") as file:
+                content = await file.read()
+                existing_data = json.loads(content)
+            combined = existing_data["response"]["exechistories"] + json_output["response"]["exechistories"]
+        else:
+            combined = json_output["response"]["exechistories"]
 
-                # 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
 
-                for entry in combined:
-                    key = (
-                        entry.get("sha256"),
-                        entry.get("filename"),
-                        entry.get("hostname"),
-                    )
-                    seen[key] = entry
+        deduplicated = list(seen.values())
+        async with aiofiles.open(file_path, "w") as file:
+            await file.write(json.dumps({"error": "Success", "response": {"exechistories": deduplicated}}))
 
-                deduplicated = list(seen.values())
-                with open(file_path, "w") as file:
-                    json.dump(
-                        {
-                            "error": "Success",
-                            "response": {"exechistories": deduplicated},
-                        },
-                        file,
-                    )
+        json_output["response"]["exechistories"].clear()
 
-                json_output["response"]["exechistories"].clear()
+        try:
+            last_date = datetime.datetime.strptime(
+                history_item["datetime"].replace(" +0000 UTC", ""), # pyright: ignore[reportPossiblyUnboundVariable]
+                "%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(f"Total of {policy.name} Complete: ")
+        except Exception:
+            pass
 
-                # 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
 
-                filebar.n = 1
-
-    # Final output
-    with open(file_path, "r") as file:
-        final_output = json.load(file)
+    async with aiofiles.open(file_path, "r") as file:
+        final_output = await file.read()
     os.remove(file_path)
 
-    return json.dumps(final_output) if outputjson else None
+    return final_output if outputjson else None
 
 
-def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
+async def getPolicyInfo(api, policy, type, days):
     executionhist_policy = pd.DataFrame()
-    exehist = pullPolicyExechistories(api, policy, type, days, True)
+    exehist = await 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",
-                ]
+                ["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",
-            )
-        )
+            executionhist_policy["policy"] = policy.name
+            executionhist_policy = executionhist_policy.drop_duplicates(subset=["sha256", "filename", "hostname"])
+            executionhist_policy = executionhist_policy.sort_values(by=["sha256", "filename"])
+        print(colorText(f"Staging of Execution history for policy: {policy.name} is complete", "green"))
         del data
         del exehist
         gc.collect()
@@ -230,8 +136,8 @@ def skipback(days):
     return ObjectId(objectid_hex)
 
 
-def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
-    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
+async def updateAuditPoliciesFromEnforcementPolices(api):
+    policy_relationship_map = await 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")
+        await api.policy_clone(enforcement_policy, audit_policy)
+        await api.policy_set_auditmode(audit_policy, "1")

services/security.py

@@ -1,18 +1,3 @@
-#    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 <https://www.gnu.org/licenses/>.
-
 import base64
 import logging
 import os
@@ -25,10 +10,12 @@ 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
+
 # Constants
 KDF_ITERATIONS = 200_000
-SALT_SIZE = 16 # 128-bit Salt
-NONCE_SIZE = 12 # AES-GCM 
+SALT_SIZE = 16  # 128-bit Salt
+NONCE_SIZE = 12  # AES-GCM
 KEY_SIZE = 32  # AES-256
 
 
@@ -54,7 +41,7 @@ def configure_keyring_backend():
         raise EnvironmentError(f"Unsupported OS: {system}")
 
 
-def store_api_key(service: str, username: str, api_key: str, password: str):
+async 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)
@@ -66,7 +53,7 @@ def store_api_key(service: str, username: str, api_key: str, password: str):
     keyring.set_password(service, username, b64)
 
 
-def retrieve_api_key(service: str, username: str, password: str) -> str:
+async def retrieve_api_key(service: str, username: str, password: str) -> str:
     configure_keyring_backend()
     b64 = keyring.get_password(service, username)
     if b64 is None:
@@ -81,41 +68,38 @@ def retrieve_api_key(service: str, username: str, password: str) -> str:
     return pt.decode()
 
 
-def api_key_exists(service: str, username: str) -> bool:
+async 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
+    return (
+        len(password) >= 12
+        and bool(re.search(r"[A-Z]", password))
+        and bool(re.search(r"[a-z]", password))
+        and bool(re.search(r"[0-9]", password))
+        and bool(re.search(r"[^A-Za-z0-9]", password))
+    )
 
 
-def getAPI(USERNAME, SERVICE_NAME):
+async 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):
+    if await 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)
+                apikey = await retrieve_api_key(SERVICE_NAME, USERNAME, password)
                 logging.debug("API key successfully retrieved.")
                 return apikey
             except Exception as e:
                 logging.warning(f"Attempt {attempt} failed: {str(e)}")
         logging.error("Failed to retrieve API key after 3 incorrect attempts.")
-        raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
+        print(colorText("❌ Authentication failed. Exiting.", "red"))
+        exit(1)
     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()
@@ -131,7 +115,7 @@ def getAPI(USERNAME, SERVICE_NAME):
 
             if check_password_complexity(password):
                 try:
-                    store_api_key(SERVICE_NAME, USERNAME, api_key, password)
+                    await store_api_key(SERVICE_NAME, USERNAME, api_key, password)
                     logging.info("API key stored securely.")
                     break
                 except Exception as e:
@@ -145,8 +129,8 @@ class APIKeyManager:
     _api_key = None
 
     @classmethod
-    def load(cls, service: str, username: str, password: str):
-        cls._api_key = retrieve_api_key(service, username, password)
+    async def load(cls, service: str, username: str, password: str):
+        cls._api_key = await retrieve_api_key(service, username, password)
 
     @classmethod
     def get(cls) -> str: