diff --git a/flows/prepPolicy.py b/flows/prepPolicy.py index 910e629..b80a07f 100644 --- a/flows/prepPolicy.py +++ b/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) diff --git a/models/execution.py b/models/execution.py index f980c91..139777e 100644 --- a/models/execution.py +++ b/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", "[]")) diff --git a/services/API.py b/services/API.py index a51eaf6..2b047bd 100644 --- a/services/API.py +++ b/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""" diff --git a/services/policyhandler.py b/services/policyhandler.py index c1698d1..101fe32 100644 --- a/services/policyhandler.py +++ b/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__) diff --git a/services/security.py b/services/security.py index 2eb7de1..fa66249 100644 --- a/services/security.py +++ b/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 diff --git a/services/setup.py b/services/setup.py index 0b67e47..7ed20aa 100644 --- a/services/setup.py +++ b/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 = [ diff --git a/utils/menus.py b/utils/menus.py index 6edfa76..d59e13e 100644 --- a/utils/menus.py +++ b/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(