From 3620192bf48caba0ccab75386e9d8faabc43dd03 Mon Sep 17 00:00:00 2001 From: Amir Ben Nun <34831306+amirbenun@users.noreply.github.com> Date: Tue, 8 Aug 2023 17:16:18 +0300 Subject: [PATCH] Test CloudFormation agents tags (#1202) --- .../fleet_api/src/agents_enrolled.py | 60 ++++++++++++++++--- .../fleet_api/src/install_cnvm_integration.py | 3 +- .../fleet_api/src/install_cspm_integration.py | 2 +- .../src/install_kspm_eks_integration.py | 2 +- .../src/install_kspm_unmanaged_integration.py | 2 +- .../fleet_api/src/state_file_manager.py | 3 +- 6 files changed, 60 insertions(+), 12 deletions(-) diff --git a/deploy/test-environments/fleet_api/src/agents_enrolled.py b/deploy/test-environments/fleet_api/src/agents_enrolled.py index 4a697632e3..25ad198f28 100644 --- a/deploy/test-environments/fleet_api/src/agents_enrolled.py +++ b/deploy/test-environments/fleet_api/src/agents_enrolled.py @@ -4,6 +4,8 @@ """ import sys import time +import re +from dataclasses import dataclass from api.agent_policy_api import get_agents import configuration_fleet as cnfg from state_file_manager import state_manager @@ -12,6 +14,16 @@ TIMEOUT = 600 +@dataclass +class AgentExpected: + """ + Class to represent the details of an enrolled agent. + """ + + count: int + tags: list[str] + + def get_expected_agents() -> dict: """ Returns: @@ -20,7 +32,7 @@ def get_expected_agents() -> dict: logger.info("Loading agent policies state file") policies_dict = {} for policy in state_manager.get_policies(): - policies_dict[policy.agnt_policy_id] = policy.expected_agents + policies_dict[policy.agnt_policy_id] = AgentExpected(policy.expected_agents, policy.expected_tags) return policies_dict @@ -36,23 +48,57 @@ def get_actual_agents() -> dict: return policies_dict -def verify_agents_enrolled() -> bool: +def verify_agent_count(expected: dict, actual: dict) -> bool: """ Verify that the expected number of agents are enrolled """ - expected = get_expected_agents() - actual = get_actual_agents() result = True - for policy_id, expected_count in expected.items(): + for policy_id, expected_agents in expected.items(): if policy_id not in actual: result = False logger.info(f"Policy {policy_id} not found in the actual agents mapping") - elif actual[policy_id] != expected_count: + elif actual[policy_id] != expected_agents.count: result = False - logger.info(f"Policy {policy_id} expected {expected_count} agents, but got {actual[policy_id]}") + logger.info(f"Policy {policy_id} expected {expected_agents.count} agents, but got {actual[policy_id]}") + else: + logger.info(f"Policy {policy_id} has {actual[policy_id]} agents as expected") return result +def verify_agent_tags(agent, expected_agents) -> bool: + """ + Verify that the agent has the expected tags + """ + expected_tags = [] + if agent.policy_id in expected_agents: + expected_tags = expected_agents[agent.policy_id].tags + for pattern in expected_tags: + pattern_exist = False + for tag in agent.tags: + if re.match(pattern, tag): + pattern_exist = True + break + if not pattern_exist: + logger.warning(f"Agent {agent.id} does not have the expected tag {pattern}") + return False + return True + + +def verify_agents_enrolled() -> bool: + """ + Construct a dictionary of the expected agents and the actual agents + Returns: + bool: True if the expected agents are enrolled, False otherwise + """ + expected = get_expected_agents() + agents = get_agents(cfg=cnfg.elk_config) + actual = {} + for agent in agents: + if verify_agent_tags(agent, expected): + actual[agent.policy_id] = actual.get(agent.policy_id, 0) + 1 + return verify_agent_count(expected, actual) + + def wait_for_agents_enrolled(timeout) -> bool: """ Wait for agents to be enrolled diff --git a/deploy/test-environments/fleet_api/src/install_cnvm_integration.py b/deploy/test-environments/fleet_api/src/install_cnvm_integration.py index 9ec7ca3114..3b5cbc30e6 100644 --- a/deploy/test-environments/fleet_api/src/install_cnvm_integration.py +++ b/deploy/test-environments/fleet_api/src/install_cnvm_integration.py @@ -27,6 +27,7 @@ CNVM_PACKAGE_POLICY = "../../../cloud/data/package_policy_cnvm_aws.json" CNVM_EXPECTED_AGENTS = 1 CNVM_CLOUDFORMATION_CONFIG = "../../../cloudformation/config.json" +CNVM_AGENT_TAGS = ["cft_version:CFT_VERSION", "cft_arn:arn:aws:cloudformation:.*"] cnvm_agent_policy_data = Path(__file__).parent / CNVM_AGENT_POLICY cnvm_pkg_policy_data = Path(__file__).parent / CNVM_PACKAGE_POLICY @@ -60,7 +61,7 @@ def load_data() -> Tuple[Dict, Dict]: agent_policy_id=agent_policy_id, ) - state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CNVM_EXPECTED_AGENTS)) + state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CNVM_EXPECTED_AGENTS, CNVM_AGENT_TAGS)) cloudformation_params = Munch() cloudformation_params.ENROLLMENT_TOKEN = get_enrollment_token( diff --git a/deploy/test-environments/fleet_api/src/install_cspm_integration.py b/deploy/test-environments/fleet_api/src/install_cspm_integration.py index 04be0b322f..937a9250c0 100755 --- a/deploy/test-environments/fleet_api/src/install_cspm_integration.py +++ b/deploy/test-environments/fleet_api/src/install_cspm_integration.py @@ -76,7 +76,7 @@ def load_data() -> Tuple[Dict, Dict]: cspm_data=cspm_data, ) - state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CSPM_EXPECTED_AGENTS)) + state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CSPM_EXPECTED_AGENTS, [])) manifest_params = Munch() manifest_params.enrollment_token = get_enrollment_token( diff --git a/deploy/test-environments/fleet_api/src/install_kspm_eks_integration.py b/deploy/test-environments/fleet_api/src/install_kspm_eks_integration.py index 42408d6777..d253449d0d 100755 --- a/deploy/test-environments/fleet_api/src/install_kspm_eks_integration.py +++ b/deploy/test-environments/fleet_api/src/install_kspm_eks_integration.py @@ -71,7 +71,7 @@ def load_data() -> Tuple[Dict, Dict]: eks_data=eks_data, ) - state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_EKS_EXPECTED_AGENTS)) + state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_EKS_EXPECTED_AGENTS, [])) manifest_params = Munch() manifest_params.enrollment_token = get_enrollment_token( diff --git a/deploy/test-environments/fleet_api/src/install_kspm_unmanaged_integration.py b/deploy/test-environments/fleet_api/src/install_kspm_unmanaged_integration.py index a495ceb53f..af5c2c9726 100755 --- a/deploy/test-environments/fleet_api/src/install_kspm_unmanaged_integration.py +++ b/deploy/test-environments/fleet_api/src/install_kspm_unmanaged_integration.py @@ -66,7 +66,7 @@ def load_data() -> Tuple[Dict, Dict]: agent_policy_id=agent_policy_id, ) - state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_UNMANAGED_EXPECTED_AGENTS)) + state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_UNMANAGED_EXPECTED_AGENTS, [])) manifest_params = Munch() manifest_params.enrollment_token = get_enrollment_token( diff --git a/deploy/test-environments/fleet_api/src/state_file_manager.py b/deploy/test-environments/fleet_api/src/state_file_manager.py index 36ff34a8a3..84451c0819 100644 --- a/deploy/test-environments/fleet_api/src/state_file_manager.py +++ b/deploy/test-environments/fleet_api/src/state_file_manager.py @@ -27,7 +27,7 @@ class PolicyState: Class to represent a policy state. """ - def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int): + def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int, expected_tags: list[str]): """ Args: agnt_policy_id (str): ID of the agent policy. @@ -37,6 +37,7 @@ def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int self.agnt_policy_id = agnt_policy_id self.pkg_policy_id = pkg_policy_id self.expected_agents = expected_agents + self.expected_tags = expected_tags class StateFileManager: