UI Improvements

flows/localApproval.py

@@ -23,13 +23,13 @@ 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 utils.utils import colorText
 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__)
 
@@ -38,14 +38,15 @@ dotenv.load_dotenv()
 
 def getLocalApprovals(api: AirlockAPIWrapper):
     working_dir = load_env("WORKING_DIR")
+    base_dir = get_base_directory
     result = api.otp_find_awaiting()
     local_approval = pd.DataFrame(result["response"]["otpusage"])
-    if os.path.exists(f"{working_dir}\\Scheduling\\newest_local_approval.parquet"):
-        previous_run = pd.read_parquet(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
+    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"{working_dir}\\Scheduling\\last_local_approval.parquet", index=False
+            f"{base_dir}\\cache\\last_local_approval.parquet", index=False
         )
-        os.remove(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
+        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[
@@ -58,7 +59,7 @@ def getLocalApprovals(api: AirlockAPIWrapper):
 
     if not local_approval.empty:
         local_approval.to_parquet(
-            f"{working_dir}\\Scheduling\\newest_local_approval.parquet", index=False
+            f"{base_dir}\\cache\\newest_local_approval.parquet", index=False
         )
 
     return local_approval
@@ -181,7 +182,8 @@ def moveToLocalApproval(api: AirlockAPIWrapper):
         print(f"{i}. {option}")
 
     try:
-        choice = int(input("Enter the number of your choice: "))
+
+        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"))

flows/otp.py

@@ -16,9 +16,10 @@
 
 
 import logging
+
 from services.agenthandler import selectAgents
-from utils.utils import colorText
 from utils.selector import Selector
+from utils.utils import colorText, get_sanitized_input
 
 logger = logging.getLogger(__name__)
 
@@ -27,18 +28,27 @@ from services.API import AirlockAPIWrapper
 
 
 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]
-
-    print(colorText("Please select a duration:", "white"))
-    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
+    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
 
@@ -54,5 +64,6 @@ def otp_activities_by_agent(api: AirlockAPIWrapper):
 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)
+    if otp_dict and list_to_revoke:
+        for revokee in list_to_revoke:
+            api.otp_revoke(revokee)

flows/quietAgent.py

@@ -46,7 +46,7 @@ def findQuietAgents(api: AirlockAPIWrapper):
         required_quiet = Selector.select_value(
             prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
             value_type=int,
-            valid_range=(1, 365),
+            valid_range=(1, 150),
         )
 
         # Get execution history as a DataFrame