From fa370f39b3a73ce409ea4d0f59854d29c744a03c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9ophile=20Diot?= Date: Fri, 27 Dec 2024 11:02:52 +0100 Subject: [PATCH] Refactor Docker and Swarm controllers to improve container and service retrieval with enhanced error handling and namespace filtering --- src/autoconf/DockerController.py | 80 +++++++++++++++++++++++--------- src/autoconf/SwarmController.py | 71 +++++++++++++++++++++------- 2 files changed, 111 insertions(+), 40 deletions(-) diff --git a/src/autoconf/DockerController.py b/src/autoconf/DockerController.py index d3f49772a..c41535fd3 100644 --- a/src/autoconf/DockerController.py +++ b/src/autoconf/DockerController.py @@ -7,6 +7,7 @@ from re import compile as re_compile from traceback import format_exc from docker.models.containers import Container +from docker.errors import DockerException from Controller import Controller @@ -16,33 +17,66 @@ class DockerController(Controller): self.__client = DockerClient(base_url=docker_host) self.__custom_confs_rx = re_compile(r"^bunkerweb.CUSTOM_CONF_(SERVER_STREAM|SERVER_HTTP|MODSEC_CRS|MODSEC|CRS_PLUGINS_BEFORE|CRS_PLUGINS_AFTER)_(.+)$") - def _get_controller_instances(self) -> List[Container]: - containers: List[Container] = self.__client.containers.list(filters={"label": "bunkerweb.INSTANCE"}) + def _get_controller_containers(self, label_key: str) -> List[Container]: + """ + Fetch containers based on a specific label and filter them by namespace. + + Args: + label_key (str): The key of the label to filter containers by (e.g., "bunkerweb.INSTANCE"). + + Returns: + List[Container]: A list of containers matching the label and namespace criteria. + """ + try: + # Retrieve containers with the specific label + containers: List[Container] = self.__client.containers.list(filters={"label": label_key}) + except DockerException as e: + self._logger.error(f"Failed to retrieve containers with label '{label_key}': {e}") + return [] + if not self._namespaces: return containers - return [ - container - for container in containers - if any( - ({label: "" for label in container.labels} if isinstance(container.labels, list) else container.labels).get("bunkerweb.NAMESPACE", "") - == namespace - for namespace in self._namespaces - ) - ] + + namespace_set = set(self._namespaces) + valid_containers = [] + + for container in containers: + try: + # Safely retrieve and validate labels + labels = getattr(container, "labels", {}) + if not isinstance(labels, dict): + if isinstance(labels, list): + labels = {label: "" for label in labels} + else: + self._logger.warning(f"Unexpected label format for container {container.id}: {labels}") + continue + + # Check if the namespace label matches any in the set + namespace = labels.get("bunkerweb.NAMESPACE", "") + if namespace in namespace_set: + self._logger.debug(f"Container {container.id} matches namespace '{namespace}'.") + valid_containers.append(container) + else: + self._logger.debug(f"Container {container.id} does not match any namespace.") + + except AttributeError as e: + self._logger.warning(f"Container {container.id} missing expected attributes: {e}") + except Exception as e: + self._logger.error(f"Unexpected error while processing container {container.id}: {e}") + + return valid_containers + + def _get_controller_instances(self) -> List[Container]: + """ + Fetch containers labeled as 'bunkerweb.INSTANCE'. + """ + return self._get_controller_containers(label_key="bunkerweb.INSTANCE") def _get_controller_services(self) -> List[Container]: - containers: List[Container] = self.__client.containers.list(filters={"label": "bunkerweb.SERVER_NAME"}) - if not self._namespaces: - return containers - return [ - container - for container in containers - if any( - ({label: "" for label in container.labels} if isinstance(container.labels, list) else container.labels).get("bunkerweb.NAMESPACE", "") - == namespace - for namespace in self._namespaces - ) - ] + """ + Fetch containers labeled as 'bunkerweb.SERVER_NAME'. + """ + return self._get_controller_containers(label_key="bunkerweb.SERVER_NAME") def _to_instances(self, controller_instance) -> List[dict]: instance = { diff --git a/src/autoconf/SwarmController.py b/src/autoconf/SwarmController.py index 6b3e7b940..6fa941ccd 100644 --- a/src/autoconf/SwarmController.py +++ b/src/autoconf/SwarmController.py @@ -9,6 +9,7 @@ from docker import DockerClient from base64 import b64decode from docker.models.services import Service +from docker.errors import DockerException from Controller import Controller @@ -22,27 +23,63 @@ class SwarmController(Controller): self.__swarm_configs = [] self._logger.warning("Swarm integration is deprecated and will be removed in a future release") - def _get_controller_instances(self) -> List[Service]: - self.__swarm_instances = [] - services = self.__client.services.list(filters={"label": "bunkerweb.INSTANCE"}) + def _get_controller_swarm_services(self, label_key: str) -> List[Service]: + """ + Fetch Swarm services based on a specific label and filter them by namespace. + + Args: + label_key (str): The key of the label to filter services by (e.g., "bunkerweb.INSTANCE"). + + Returns: + List[Service]: A list of services matching the label and namespace criteria. + """ + try: + # Retrieve services with the specific label + services: List[Service] = self.__client.services.list(filters={"label": label_key}) + except DockerException as e: + self._logger.error(f"Failed to retrieve services with label '{label_key}': {e}") + return [] + if not self._namespaces: return services - return [ - service - for service in services - if any(service.attrs["Spec"]["Labels"].get("bunkerweb.NAMESPACE", "") == namespace for namespace in self._namespaces) - ] + + namespace_set = set(self._namespaces) + valid_services = [] + + for service in services: + try: + # Safely retrieve and validate labels + labels = service.attrs.get("Spec", {}).get("Labels", {}) + if not isinstance(labels, dict): + self._logger.warning(f"Unexpected label format for service {service.id}: {labels}") + continue + + # Check if the namespace label matches any in the set + namespace = labels.get("bunkerweb.NAMESPACE", "") + if namespace in namespace_set: + self._logger.debug(f"Service {service.id} matches namespace '{namespace}'.") + valid_services.append(service) + else: + self._logger.debug(f"Service {service.id} does not match any namespace.") + + except AttributeError as e: + self._logger.warning(f"Service {service.id} missing expected attributes: {e}") + except Exception as e: + self._logger.error(f"Unexpected error while processing service {service.id}: {e}") + + return valid_services + + def _get_controller_instances(self) -> List[Service]: + """ + Fetch Swarm services labeled as 'bunkerweb.INSTANCE'. + """ + return self._get_controller_swarm_services(label_key="bunkerweb.INSTANCE") def _get_controller_services(self) -> List[Service]: - self.__swarm_services = [] - services = self.__client.services.list(filters={"label": "bunkerweb.SERVER_NAME"}) - if not self._namespaces: - return services - return [ - service - for service in services - if any(service.attrs["Spec"]["Labels"].get("bunkerweb.NAMESPACE", "") == namespace for namespace in self._namespaces) - ] + """ + Fetch Swarm services labeled as 'bunkerweb.SERVER_NAME'. + """ + return self._get_controller_swarm_services(label_key="bunkerweb.SERVER_NAME") def _to_instances(self, controller_instance) -> List[dict]: self.__swarm_instances.append(controller_instance.id)