UI Improvements

services/agenthandler.py

@@ -27,9 +27,9 @@ import pandas as pd
 from models.agent import Agent
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
-from utils.configmanager import load_env, load_env_json, get_protected_json
+from utils.configmanager import get_protected_json, load_env
 from utils.selector import Selector
-from utils.utils import colorText
+from utils.utils import colorText, get_sanitized_input
 
 logger = logging.getLogger(__name__)
 
@@ -116,7 +116,7 @@ def findAgents(api, return_dataframe):
     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()
+    user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
     if user_input == 'y':
         timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
         filename = f"agentsearch_{timestamp}.csv"
@@ -134,6 +134,7 @@ def findAgents(api, return_dataframe):
     else:
         logging.debug("User declined to export the DataFrame.")
 
+
 def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
     print(colorText("🔍 Device Search", "cyan"))
     print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
@@ -150,23 +151,34 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
     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 = input()
-        if line.strip() == "":
+        line = get_sanitized_input("")
+        stripped_line = line.strip()
+
+        if stripped_line == "":
             empty_line_count += 1
             if empty_line_count == 2:
                 break
+            continue  # Don't validate empty lines
         else:
             empty_line_count = 0
-            device_input_lines.append(line.strip())
 
+        # 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
+    # Build regex pattern to match hostnames
     pattern = "|".join(map(re.escape, device_names))
     regex = re.compile(pattern, re.IGNORECASE)
 
@@ -188,12 +200,11 @@ 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
+    # 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,