AirlockToolClient.py

@@ -21,9 +21,6 @@
 #TODO Fix Requirements.txt
 #TODO Create Generic system_config.json for gitea
 
-
-import argparse
-import json
 import logging
 import os
 
@@ -33,7 +30,7 @@ import urllib3
 import utils.menus as menus
 from services.API import AirlockAPIWrapper
 from services.security import getAPI
-from services.setup import setup
+from utils.setup import setup
 
 urllib3.disable_warnings(
     urllib3.exceptions.InsecureRequestWarning
@@ -70,54 +67,8 @@ def main():
         api_key = getAPI(username, "AirlockTools"),
     )    
 
-    parser = argparse.ArgumentParser(description="Program for Managing Airlock via  API & CMD")
-    parser.add_argument("--monitor", action="store_true", help="Run in non-interactive mode")
-    # Add other arguments as needed
-
-    args = parser.parse_args()
-
-    if args.monitor:
-        # Non-interactive logic
-        logger.info("Running non-interactively to start monitoring Airlock Changes")
-        """
-            os.makedirs("scheduling", exist_ok=True)
-            os.makedirs("OTP/HTML", exist_ok=True)
-            os.makedirs("OTP/PARQ", exist_ok=True)
-            os.makedirs("Local_Approval/HTML", exist_ok=True)
-            os.makedirs("Local_Approval/PARQ", exist_ok=True)
-                                
-            register_function("monitorOTP", utils.otpfunctions.monitorOTP)
-            register_function("monitorLA", la.scheduleAddingLAHashes)
-            register_function("updateAuditPolicies", utils.policyfunctions.updateAuditPoliciesFromEnforcementPolices)
-
-            
-            if not os.path.exists("scheduling\\jobs.json"): 
-                recurring_job("monitorOTP", "monitorOTP", interval=60, unit="seconds", args=[url, pups])
-                recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant])
-                recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[url, policy_relationship_map])
-            else:
-                reload_jobs()
-            start_scheduler()
-            """
-    else:
-        # Interactive logic
-    
-        
-        
-        raw = os.getenv("POLICY_MAP_ENF_AUD", "{}")
-
-        try:
-            # Escape backslashes before parsing
-            escaped = raw.encode('unicode_escape').decode('utf-8')
-            badpathparts = json.loads(escaped)
-        except Exception as e:
-            logging.error(f"Failed to parse BAD_PATH_PARTS: {e}")
-            badpathparts = []
-        print(badpathparts)
-
-
-
-        menus.menu_main(api)
+   
+    menus.menu_main(api)
 
 
 if __name__ == "__main__":

AirlockToolServer.py

@@ -0,0 +1,89 @@
+#    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/>.
+
+
+#TODO Continue implementing logger
+#TODO Add input sanitation and CSV injection prevention
+#TODO Continue OTP and Local approval rewrites
+#TODO Explore pywin32
+#TODO Fix Requirements.txt
+#TODO Create Generic system_config.json for gitea
+
+
+import logging
+import os
+
+import dotenv
+import urllib3
+
+from Server.scheduler_async import start_scheduler, register_function, recurring_job, reload_jobs
+from services.API import AirlockAPIWrapper
+from services.security import getAPI
+from utils.setup import setup
+import flows.localApproval as la
+from services.policyhandler import updateAuditPoliciesFromEnforcementPolices
+urllib3.disable_warnings(
+    urllib3.exceptions.InsecureRequestWarning
+)
+
+def main():
+
+    #Determine  working directory, setup directory, configure logging, sent env, get API and URL if not already stored 
+  
+    working_dir = setup()
+    
+    logger = logging.getLogger(__name__)
+
+    dotenv.load_dotenv(dotenv_path=working_dir / ".env")
+
+    try:
+            url = os.getenv("URL")
+            username = os.getenv("USERNAME")
+
+            if not url:
+                raise ValueError("Missing URL in environment variables.")
+            if not username:
+                raise ValueError("Missing USERNAME in environment variables.")
+
+            logger.debug(f"Retrieved URL: {url}")
+            logger.debug(f"Retrieved Username: {username}")
+
+    except ValueError as e:
+            logger.error(f"Configuration error: {e}", exc_info=True)
+            raise
+
+
+    api = AirlockAPIWrapper(
+        base_url=str(os.getenv("URL")),
+        api_key = getAPI(username, "AirlockTools"),
+    )    
+
+    logger.info("Running non-interactively to start monitoring Airlock Changes")
+
+
+    register_function("monitorLA", la.scheduleAddingLAHashes)
+    register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices)
+
+    
+    if not os.path.exists("scheduling\\jobs.json"): 
+        recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api])
+        recurring_job("updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api])
+    else:
+        reload_jobs()
+    
+    start_scheduler()
+        
+if __name__ == "__main__":
+    main()

Server/scheduler_async.py

@@ -0,0 +1,177 @@
+import asyncio
+import json
+import logging
+import os
+from typing import Callable, Dict, Any, 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.")

flows/localApproval.py

@@ -23,11 +23,11 @@ import time
 import dotenv
 import numpy as np
 import pandas as pd
-
+from utils.setup import get_base_directory
 from models.agent import Agent
 from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
 from services.API import AirlockAPIWrapper
-from services.scheduler import (
+from Server.scheduler import (
     register_function,
     run_once_job,
 )

flows/otp.py

@@ -16,205 +16,43 @@
 
 
 import logging
+from services.agenthandler import selectAgents
+from utils.utils import colorText
+from utils.selector import Selector
 
 logger = logging.getLogger(__name__)
 
 
-"""
-def getActiveOTP(url):
-
-    endpoint = url + f'/v1/otp/usage'
-    payload = {
-            "status" : "1"
-    }
-
-    headers = {"X-APIKey": load_env('APIKEY')}
-    payload = json.dumps(payload)
-
-    response = requests.post(endpoint, headers=headers, data=payload, verify=False)
-    result = json.loads(response.text)
-    otp = pd.DataFrame(result["response"]["otpusage"])
-    if os.path.exists("OTP\\PARQ\\newest_active_OTP.parquet"):
-        previous_run = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet")
-        previous_run.to_parquet("OTP\\PARQ\\old_active_OTP.parquet", index=False)
-        os.remove("OTP\\PARQ\\newest_active_OTP.parquet")
-    otp.to_parquet("OTP\\PARQ\\newest_active_OTP.parquet", index=False)
-    if not otp.empty:
-        formatHTML(otp, f"OTP\\HTML\\newest_active_OTP.html")
-
-def getOTPActivities(url, otpid):
-    endpoint = url + f'/v1/otp/activities'
-    payload = {"otpid": f"{otpid}"}
-    headers = {"X-APIKey": load_env('APIKEY')}
-    payload = json.dumps(payload)
-
-    response = requests.post(endpoint, headers=headers, data=payload, verify=False)
-    result = json.loads(response.text)
-    new_data = pd.DataFrame(result["response"]["otpactivities"])
-
-    # Define file path
-    parquet_path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
-
-    # Check if file exists and read it
-    if os.path.exists(parquet_path):
-        existing_data = pd.read_parquet(parquet_path)
-        combined_data = pd.concat([existing_data, new_data], ignore_index=True)
-        combined_data.drop_duplicates(inplace=True)
-    else:
-        combined_data = new_data
-
-    # Save combined data
-    combined_data.to_parquet(parquet_path, index=False)
-
-    # Optional: generate styled HTML if there's data
-    if not combined_data.empty:
-        formatHTML(combined_data, f"OTP/HTML/OTP_activities_{otpid}.html")
-
-def monitorOTP(url, pups):
-
-    getActiveOTP(url)
+from services.API import AirlockAPIWrapper
 
 
-    old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet"
-    new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet"
-
-    if os.path.exists(old_otp_path):
-        old_active_OTP = pd.read_parquet(old_otp_path)
-    else:
-        old_active_OTP = pd.DataFrame(columns=['otpid'])  # Ensure expected column exists
-
-    current_active_OTP = pd.read_parquet(new_otp_path)
-
-
-    if 'otpid' not in current_active_OTP.columns: current_active_OTP = pd.DataFrame(columns=['otpid'])
-    if 'otpid' not in old_active_OTP.columns: old_active_OTP = pd.DataFrame(columns=['otpid'])
-
-    newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])]
-    still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
-    no_longer_OTP = old_active_OTP[~old_active_OTP['otpid'].isin(current_active_OTP['otpid'])]
-
-
-    register_function("addhash", addOTPHashes)
-
-    for _, row in newly_added.iterrows():
-        clientid = row['clientid']
-        duration = (int(row['duration']) * 60)
-        hostname = row['hostname']
-        purpose = row ['purpose']
-        pid = row['otpid']
-        early = math.floor(duration * .95)
-        #If newly added to the list - schedule adding the majority of the executions prior to the expiration of OTP period.
-        run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}", "addhash", time.time() + early, [url, clientid, pid, pups], None)
-        print(f"Processing: {pid} with other data: {row}")
-
-
-    for _, row in still_in_OTP.iterrows():
-        pid = row['otpid']
-        #While still in OTP, continue to update activities list
-        getOTPActivities(url,pid)
-
-    for _, row in no_longer_OTP.iterrows():
-        clientid = row['clientid']
-        hostname = row['hostname']
-        purpose = row ['purpose']
-        pid = row['otpid']
-        allowlist = clientf.getDestAllowlistFromClientID(url,clientid)
-        policy, policyid = clientf.getPolicyFromClientID(url,clientid)
-
-"""
-
-# Devices can come out of OTP either by timeout, or by early move out of OTP.  If they are manually moved out prior to the job to add hashes can run, we want to accelerate the job.
-# But first, we want to update the OTP activities one final time for the pid, then move up any jobs if they exist, then add the hashes to the local approval allowlist
-"""
-        
-        getOTPActivities(url,pid)
-        find_and_prioritize_jobs_by_pid(pid, 1)        
-        addOTPHashes(url, clientid,pid, pups)
-        
-        finalhashesadded = pd.read_parquet(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
-        finalhashesadded['policy'] = policy
-        finalhashesadded['allowlist'] = allowlist
-        finalhashesadded['added_at'] = time.localtime()
-        
-        if not os.path.exists(f"OTP\\PARQ\\localapprovalhistory.parquet"): 
-            df = pd.DataFrame()
-            df.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
-        
-        history = pd.read_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
-        history = pd.concat([history, finalhashesadded], ignore_index=True)
-        formatHTML(history, f"localapproval_history.html")
-
-        history.to_parquet(f"OTP\\PARQ\\localapprovalhistory.parquet")
-        
-        os.remove(f"OTP\\PARQ\\otp_activities_{pid}.parquet")
-        del finalhashesadded
-
-def addOTPHashes(url, clientid, otpid, pups):
-    path = f"OTP\\PARQ\\otp_activities_{otpid}.parquet"
-    activities = pd.read_parquet(path)
-
-    pattern = regulator(pups)
-    allowlist = clientf.getDestAllowlistFromClientID(url, clientid)
-
-    # Initialize or preserve 'hash_added' column
-    if "hash_added" not in activities.columns: activities["hash_added"] = None
-
-    # Identify rows that should be added (not matching pattern and not already added)
-    approve_by_hash = activities[
-        ~activities["filename"].str.contains(pattern, na=False) & (activities["hash_added"] != "added")
-    ]
-
-    hashes_to_add = approve_by_hash["sha256"].tolist()
-
-    # Add hashes to policy
-    if hashes_to_add:
-        #TODO add the api call 
-        pass
-
-    # Update 'hash_added' column
-    activities["hash_added"] = activities.apply(
-        lambda row: "do not add" if pd.notna(row["filename"]) and pattern in row["filename"]
-        else ("added" if row["sha256"] in hashes_to_add else row["hash_added"]),
-        axis=1
-    )
-
-    # Save the updated DataFrame
-    activities.to_parquet(path)
-
-def generateOTP(url, agentid):
+def generate(api: AirlockAPIWrapper):
+    agents = selectAgents(api)
     purpose = input(colorText(" Please enter the purpose for the OTP: ", "white"))
-
     possible_durations = [15, 60, 360, 1440, 10080]
-    duration_selected = " "
 
     print(colorText("Please select a duration:", "white"))
-    for i, option in enumerate(possible_durations, start=1):
-        print(f"{i}. {option}")
+    duration_selected = Selector.select_int(possible_durations)
+    otp_dict = {}
+    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
 
-    try:
-        choice = int(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"))
-        else:
-            print(colorText("Invalid choice.", "red"))
-    except ValueError:
-        print(colorText("Invalid input. Please enter a number.", "red"))
+    return otp_dict
 
-    endpoint = url + '/v1/otp/retrieve'
-    payload = {
-            "duration" : f"{duration_selected}",
-            "agentid"  : f"{agentid}",
-            "purpose" : f"{purpose}"
-    }
+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
 
-    headers = {"X-APIKey": load_env('APIKEY')}
-    payload = json.dumps(payload)
+    return otp_dict
 
-
-    response = requests.post(endpoint, headers=headers, data=payload, verify=False)
-    result = json.loads(response.text)
-    otpcode = result["response"]["otpcode"]
-    print(colorText(f"The OPT code is: {otpcode}", "yellow"))
-
-"""
+def revoke(api: AirlockAPIWrapper):
+    otp_dict = otp_activities_by_agent(api)
+    list_to_revoke = [entry["otpid"] for entry in otp_dict]
+    for revokee in list_to_revoke:
+        api.otp_revoke(revokee)

flows/prepPolicy.py

@@ -24,7 +24,7 @@ import pandas as pd
 from models.execution import ExecutionHistoryRecord, Hash
 from models.policy import Allowlist, Policy
 from services.API import AirlockAPIWrapper
-from services.selector import Selector
+from utils.selector import Selector
 from utils.utils import (
     colorText,
     formatHTML,
@@ -42,10 +42,10 @@ dotenv.load_dotenv()
 
 def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
 
-    allowlists = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
+    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(allowlists, allow_multiple, prompt_each=True)
+    selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)
 
     if selected is None:
         return []

flows/quietAgent.py

@@ -21,7 +21,7 @@ import pandas as pd
 from flows.prepPolicy import selectPolicies
 from services.API import AirlockAPIWrapper
 from services.policyhandler import getPolicyInfo
-from services.selector import Selector
+from utils.selector import Selector
 from utils.utils import colorText, load_env
 
 logger = logging.getLogger(__name__)

models/agent.py

@@ -15,6 +15,8 @@
 
 from dataclasses import dataclass, field
 from typing import ClassVar, Optional
+from models.policy import Policy
+from typing import List
 
 
 @dataclass
@@ -23,7 +25,7 @@ class Agent:
     clientversion: str
     domain: str
     freespace: int
-    groupid: int
+    groupid: str  # Changed to str to match UUID-style IDs
     hostname: str
     ip: str
     localip: str
@@ -36,14 +38,23 @@ class Agent:
     status_text: Optional[str] = field(default=None)
 
     # Class-level status map
-    status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"}
-
-    def enrich(self, groupid_to_name: dict):
-        """Enrich the agent with groupname and human-readable status."""
-        self.groupname = groupid_to_name.get(self.groupid, None)
-        self.status_text = self.status_map.get(self.status, "Unknown")
+    status_map: ClassVar[dict] = {
+        0: "Offline",
+        1: "Online",
+        2: "Hidden",
+        3: "Safemode"
+    }
 
 
+    def enrich_with_policies(self, policies: List[Policy]):
+            """Enrich the agent with groupname and human-readable status."""
+            self.status_text = self.status_map.get(self.status, "Unknown")
+            for policy in policies:
+                if policy.groupid == self.groupid:
+                    self.groupname = policy.name
+                    break
+            if not self.groupname:
+                self.groupname = "Unknown"
 """
 
 from models.agent import Agent

services/API.py

@@ -141,7 +141,13 @@ class AirlockAPIWrapper:
         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 = {
@@ -157,6 +163,29 @@ class AirlockAPIWrapper:
         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:

services/agenthandler.py

@@ -27,7 +27,7 @@ import pandas as pd
 from models.agent import Agent
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
-from services.selector import Selector
+from utils.selector import Selector
 from utils.utils import colorText, load_env, load_env_json
 
 logger = logging.getLogger(__name__)
@@ -94,43 +94,44 @@ def findAllAgents(api):
 
     return agents
 
-
-
 def findAgents(api, return_dataframe):
     agents = selectAgents(api)
     working_dir = load_env("WORKING_DIR")
-    if agents:
-        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
-        else:
-            print(agent_df)
-            logging.debug("Displayed DataFrame to console.")
-
-            # Ask user if they want to export
-            user_input = 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.")
-    else:
-        logging.warning("No agents found.")
+    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 = 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"))
@@ -143,6 +144,7 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
     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
@@ -185,6 +187,10 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
         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(

services/policyhandler.py

@@ -27,7 +27,7 @@ from bson import ObjectId
 
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
-from services.setup import get_base_directory
+from utils.setup import get_base_directory
 from utils.utils import colorText, load_env_json
 
 logger = logging.getLogger(__name__)

utils/menus.py

@@ -30,6 +30,13 @@ from flows.prepPolicy import (
     sortHashes,
 )
 from flows.quietAgent import findQuietAgents
+
+from flows.otp import (
+    generate,
+    otp_activities_by_agent,
+    revoke
+)
+
 from services.agenthandler import findAgents
 from services.API import AirlockAPIWrapper
 from utils.utils import (
@@ -41,6 +48,8 @@ from utils.utils import (
     printEnforceChecklist,
 )
 
+
+
 logger = logging.getLogger(__name__)
 
 dotenv.load_dotenv()
@@ -237,24 +246,33 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
 
 
 def menu_otp(api: AirlockAPIWrapper):
+    working_dir = load_env("WORKING_DIR")
     while True:
+
         print(colorText("\n--- 🎫 OTP Submenu 🎫 ---", "cyan"))
-        print(colorText("1. Generate OTP", "cyan"))
-        # print(colorText("2. Sub-option B","cyan"))
-        print(colorText("Q. Return to Main Menu", "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("Q. 🔚  - Quit", "yellow"))
+
         choice = input("Enter your choice: ")
 
         if choice == "1":
-            # TODO generateOTP(api,findAgents()
-            break
-
+            otp_list = generate(api)
+            print(colorText(otp_list,"green"))
         elif choice == "2":
-            print("You selected Sub-option B")
+            otp_activities_by_agent(api)
+        elif choice == "3":
+            revoke(api)
+        elif choice == "F":
+            open_directory(working_dir)
+        elif choice == "S":
+            menu_settings()
         elif choice == "Q":
-            print("Returning to Main Menu...")
             break
-        else:
-            print("Invalid choice. Please try again.")
+   
 
 
 def menu_settings():

utils/selector.py

@@ -20,51 +20,62 @@ from typing import Any, List, Optional, Union
 logger = logging.getLogger(__name__)
 
 
+from typing import List, Optional, Union, Any, Callable
+import logging
+
+logger = logging.getLogger(__name__)
+
 class Selector:
     @staticmethod
-    def select_objects(
-        objects: List[Any],
-        allow_multiple: bool = False,
-        prompt_each: bool = False
-    ) -> Union[Optional[Any], List[Any]]:
-        if not objects:
-            logger.warning("No objects available for selection.")
-            return None
-
-        # Sort objects alphabetically by their 'name' attribute
-        sorted_objects = sorted(objects, key=lambda obj: getattr(obj, "name", str(obj)).lower())
-
-        # Display in 4 columns with extra spacing
-        num_columns = 4
-        rows = (len(sorted_objects) + num_columns - 1) // num_columns
-        print("\nAvailable Choices:")
+    def _display_choices(
+        items: List[Any],
+        label_func: Callable[[Any], str],
+        num_columns: int = 4,
+        header: str = "Available Choices:"
+    ) -> None:
+        sorted_items = sorted(items, key=lambda item: label_func(item).lower())
+        rows = (len(sorted_items) + num_columns - 1) // num_columns
+        print(f"\n{header}")
         for row in range(rows):
             line = ""
             for col in range(num_columns):
                 idx = row + col * rows
-                if idx < len(sorted_objects):
-                    obj = sorted_objects[idx]
-                    name = getattr(obj, "name", str(obj))
-                    line += f"{idx + 1}: {name:<30}"
+                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 = input("Select an object by number (or Q to finish): ").strip().lower()
+                choice = 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_objects):
-                        obj = sorted_objects[index - 1]
-                        if obj not in selected:
-                            selected.append(obj)
+                    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: {getattr(obj, 'name', str(obj))}")
+                                logger.info(f"Selected: {label_func(item)}")
                         else:
-                            logger.warning("Object already selected.")
+                            logger.warning("Item already selected.")
                     else:
                         logger.warning("Selection out of range. Try again.")
                 except ValueError:
@@ -72,17 +83,59 @@ class Selector:
             return selected if selected else None
         else:
             try:
-                choice = int(input("Select one object by number: "))
-                if 1 <= choice <= len(sorted_objects):
-                    selected_obj = sorted_objects[choice - 1]
-                    logger.info(f"Selected: {getattr(selected_obj, 'name', str(selected_obj))}")
-                    return selected_obj
+                choice = int(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,

utils/setup.py

@@ -147,8 +147,7 @@ def setup() -> Path:
         "Approved": [],
         "Needs_Review": ["Review_First", "Review_Second", "HTML"],
         "Preflight": ["HTML"],
-        "Archived": ["HTML"],
-        "Scheduling": []
+        "Archived": []
     }
 
     for folder_name, subfolders in folders_structure.items():