flows/prepPolicy.py

@@ -55,9 +55,9 @@ def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
     return selected if isinstance(selected, list) else [selected]
 
 
-def selectAllowlists(api: AirlockAPIWrapper, allow_multiple=True) -> List[Allowlist]:
-
-    allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
+def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]:
+    if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
+    else:  allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()]
     logger.debug("Prompting for Allowlist(s)")
     print(colorText("Please select allowlist(s)", "white"))
     selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)

models/execution.py

@@ -126,8 +126,6 @@ class Hash:
 
     @classmethod
     def categorize_hashes(cls, hashes):
-        import re
-        from utils.utils import load_env, load_env_json, regulator
 
         threat_tolerance = load_env("VT_THREAT_TOLERANCE", cast_type=int)
         bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))

services/API.py

@@ -184,6 +184,12 @@ class AirlockAPIWrapper:
         payload = {"groupid": groupid}
         result = self._post("/v1/group/agents", payload)
         return pd.DataFrame(result["response"]["agents"])
+        
+    def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
+        """List allowlists assigned to a specific policy group."""
+        payload = {"groupid": groupid}
+        result = self._post("/v1/group/policies", payload)
+        return pd.DataFrame(result["response"]["applications"])
 
     def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
         """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""

services/policyhandler.py

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

services/security.py

@@ -118,10 +118,16 @@ def getAPI(USERNAME, SERVICE_NAME):
         raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
     else:
         logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.")
-        api_key = input(f"🔑 No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
+        api_key = getpass(f"🔑 No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
 
         while True:
             password = getpass("🔐 Create a password to encrypt your API key: ")
+            confirm_password = getpass("🔐 Confirm your password: ")
+
+            if password != confirm_password:
+                logging.warning("❌ Passwords do not match. Try again.")
+                continue
+
             if check_password_complexity(password):
                 try:
                     store_api_key(SERVICE_NAME, USERNAME, api_key, password)
@@ -131,8 +137,8 @@ def getAPI(USERNAME, SERVICE_NAME):
                     logging.error(f"Failed to store API key: {e}")
                     break
             else:
-                print("❌ Password does not meet complexity requirements. Try again.")
-        return api_key
+                logging.warning("❌ Password does not meet complexity requirements. Try again.")
+
 
 class APIKeyManager:
     _api_key = None

services/setup.py

@@ -22,6 +22,7 @@ import os
 import platform
 import sys
 from pathlib import Path
+
 from dotenv import load_dotenv, set_key
 
 PROTECTED_KEYS = [

utils/menus.py

@@ -113,7 +113,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
 
             print(colorText("Please choose Allowlist for Hashes", "white"))
 
-            destination_allowlist = selectAllowlists(api, False)
+            destination_allowlist = selectAllowlists(api, destination_policy, False)
 
         elif choice == "3":
             sortHashes(