Cleaning up menu, added Move to Audit/Enforcement

services/agenthandler.py

@@ -27,8 +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.selector import Selector
-from utils.utils import colorText, load_env, load_env_json
+from utils.utils import colorText
 
 logger = logging.getLogger(__name__)
 
@@ -207,7 +208,7 @@ def moveAgentToRelatedPolicy(
         policy_relationship_map: Dict mapping enforcement → audit.
         mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
     """
-    policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
+    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
 
     if mode == "audit":
         if agent.groupid in policy_relationship_map:

services/policyhandler.py

@@ -27,8 +27,9 @@ from bson import ObjectId
 
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
+from utils.configmanager import get_protected_json
 from utils.setup import get_base_directory
-from utils.utils import colorText, load_env_json
+from utils.utils import colorText
 
 logger = logging.getLogger(__name__)
 
@@ -230,7 +231,7 @@ def skipback(days):
 
 
 def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
-    policy_relationship_map = load_env_json("POLICY_MAP_ENF_AUD", "{}")
+    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")
     for enforcement_policy, audit_policy in policy_relationship_map.items():
         api.policy_clone(enforcement_policy, audit_policy)
         api.policy_set_auditmode(audit_policy, "1")