Add optional functionality to create extra siteadmin host groups (#49)

config.sample.toml

@@ -20,6 +20,7 @@ managed_inventory = ["location"]
 #hostgroup_disabled = "All-auto-disabled-hosts"
 #hostgroup_source_prefix = "Source-"
 #hostgroup_importance_prefix = "Importance-"
+#extra_siteadmin_hostgroup_prefixes = []
 
 [source_collectors.mysource]
 module_name = "mysource"

tests/conftest.py

@@ -1,4 +1,6 @@
 import os
+from pathlib import Path
+from typing import Iterable
 import pytest
 
 
@@ -95,3 +97,27 @@ def sample_config():
         os.path.dirname(os.path.dirname(__file__)) + "/config.sample.toml"
     ) as config:
         yield config.read()
+
+
+@pytest.fixture
+def hostgroup_map_file(tmp_path: Path) -> Iterable[Path]:
+    contents = hostgroup_map = """
+# This file defines assosiation between siteadm fetched from Nivlheim and hostsgroups in Zabbix.
+# A siteadm can be assosiated only with one hostgroup or usergroup.
+# Example: <siteadm>:<host/user groupname>
+#
+#****************************************************************************************
+# ATT: First letter will be capitilazed, leading and trailing spaces will be removed and 
+#      spaces within the hostgroupname will be replaced with "-" by the script automatically 
+#****************************************************************************************
+#
+[email protected]:Hostgroup-user1-primary
+#
+[email protected]:Hostgroup-user2-primary
+[email protected]:Hostgroup-user2-secondary
+#
+[email protected]:Hostgroup-user3-primary
+"""
+    map_file_path = tmp_path / "siteadmin_hostgroup_map.txt"
+    map_file_path.write_text(contents)
+    yield map_file_path

tests/test_utils.py

@@ -1,11 +1,14 @@
-from ipaddress import IPv4Address, IPv6Address
 import logging
+from ipaddress import IPv4Address, IPv6Address
 from pathlib import Path
 from typing import Dict, List, Set, Tuple, Union
 
 import pytest
+from pytest import LogCaptureFixture
+from hypothesis import HealthCheck, given, settings
+from hypothesis import strategies as st
+
 from zabbix_auto_config import utils
-from hypothesis import given, settings, strategies as st, HealthCheck, assume
 
 
 @pytest.mark.parametrize(
@@ -155,3 +158,75 @@ def test_zac_tags2zabbix_tags(
     zabbix_tags = utils.zac_tags2zabbix_tags(tags)
     for tag in expected:
         assert tag in zabbix_tags
+
+
+# Test with the two prefixes we use + no prefix
+@pytest.mark.parametrize(
+    "prefix",
+    ["Templates-", "Siteadmin-"],
+)
+def test_mapping_values_with_prefix(hostgroup_map_file: Path, prefix: str):
+    m = utils.read_map_file(hostgroup_map_file)
+
+    # Make sure we read the map file correctly
+    assert len(m) == 3
+
+    old_prefix = "Hostgroup-"
+    new_map = utils.mapping_values_with_prefix(
+        m,
+        prefix=prefix,
+    )
+
+    # Compare new dict to old dict
+    assert new_map is not m  # we should get a new dict
+    assert len(new_map) == len(m)
+    assert sum(len(v) for v in new_map.values()) == sum(len(v) for v in m.values())
+
+    # Check values in new map
+    assert new_map["[email protected]"] == [f"{prefix}user1-primary"]
+    assert new_map["[email protected]"] == [
+        f"{prefix}user2-primary",
+        f"{prefix}user2-secondary",
+    ]
+    assert new_map["[email protected]"] == [f"{prefix}user3-primary"]
+
+    # Check values in old map (they should be untouched)
+    assert m["[email protected]"] == [f"{old_prefix}user1-primary"]
+    assert m["[email protected]"] == [
+        f"{old_prefix}user2-primary",
+        f"{old_prefix}user2-secondary",
+    ]
+    assert m["[email protected]"] == [f"{old_prefix}user3-primary"]
+
+
+def test_mapping_values_with_prefix_no_prefix_arg(caplog: LogCaptureFixture) -> None:
+    """Passing an empty string as the prefix should be ignored and logged."""
+    res = utils.mapping_values_with_prefix(
+        {"[email protected]": ["Hostgroup-user1-primary"]},
+        prefix="",
+    )
+    assert res == {"[email protected]": []}
+    assert caplog.text.count("WARNING") == 1
+
+
+def test_mapping_values_with_prefix_no_group_prefix(caplog: LogCaptureFixture) -> None:
+    """Passing a group name with no prefix separated by the separator
+    should be ignored and logged."""
+    res = utils.mapping_values_with_prefix(
+        {"[email protected]": ["Mygroup"]},
+        prefix="Foo-",
+    )
+    assert res == {"[email protected]": []}
+    assert caplog.text.count("WARNING") == 1
+
+
+def test_mapping_values_with_prefix_no_prefix_separator(
+    caplog: LogCaptureFixture,
+) -> None:
+    """Passing a prefix with no separator emits a warning (but is otherwise legal)."""
+    res = utils.mapping_values_with_prefix(
+        {"[email protected]": ["Hostgroup-user1-primary", "Hostgroup-user1-secondary"]},
+        prefix="Foo",
+    )
+    assert res == {"[email protected]": ["Foouser1-primary", "Foouser1-secondary"]}
+    assert caplog.text.count("WARNING") == 2

zabbix_auto_config/models.py

@@ -14,6 +14,7 @@
     BaseModel,
     BaseSettings,
     conint,
+    root_validator,
     validator,
     Extra,
 )
@@ -42,6 +43,14 @@ class ZabbixSettings(BaseSettings):
     hostgroup_source_prefix: str = "Source-"
     hostgroup_importance_prefix: str = "Importance-"
 
+    # Prefixes for extra host groups to create based on the host groups
+    # in the siteadmin mapping. 
+    # e.g. Siteadmin-foo -> Templates-foo if list is ["Templates-"]
+    # The groups must have prefixes separated by a hyphen (-) in order 
+    # to replace them with any of these prefixes.
+    # These groups are not managed by ZAC beyond creating them.
+    extra_siteadmin_hostgroup_prefixes: Set[str] = set()
+
 class ZacSettings(BaseSettings):
     source_collector_dir: str
     host_modifier_dir: str
@@ -83,7 +92,7 @@ class Host(BaseModel):
     enabled: bool
     hostname: str
 
-    importance: Optional[conint(ge=0)]  # type: ignore # mypy blows up: https://github.com/samuelcolvin/pydantic/issues/156#issuecomment-614748288
+    importance: Optional[conint(ge=0)]  # type: ignore # mypy blows up: https://github.com/pydantic/pydantic/issues/239 & https://github.com/pydantic/pydantic/issues/156
     interfaces: List[Interface] = []
     inventory: Dict[str, str] = {}
     macros: Optional[None] = None  # TODO: What should macros look like?

zabbix_auto_config/processing.py

@@ -11,6 +11,7 @@
 import signal
 import itertools
 import queue
+from typing import Dict, List
 
 import psycopg2
 import pyzabbix
@@ -332,7 +333,7 @@ def merge_sources(self):
 
 
 class ZabbixUpdater(BaseProcess):
-    def __init__(self, name, state, db_uri, zabbix_config):
+    def __init__(self, name, state, db_uri, zabbix_config: models.ZabbixSettings):
         super().__init__(name, state)
 
         self.db_uri = db_uri
@@ -362,9 +363,15 @@ def __init__(self, name, state, db_uri, zabbix_config):
             logging.error("Unable to login to Zabbix API: %s", str(e))
             raise exceptions.ZACException(*e.args)
 
-        self.property_template_map = utils.read_map_file(os.path.join(self.config.map_dir, "property_template_map.txt"))
-        self.property_hostgroup_map = utils.read_map_file(os.path.join(self.config.map_dir, "property_hostgroup_map.txt"))
-        self.siteadmin_hostgroup_map = utils.read_map_file(os.path.join(self.config.map_dir, "siteadmin_hostgroup_map.txt"))
+        self.property_template_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "property_template_map.txt")
+        )
+        self.property_hostgroup_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "property_hostgroup_map.txt")
+        )
+        self.siteadmin_hostgroup_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "siteadmin_hostgroup_map.txt")
+        )
 
     def work(self):
         start_time = time.time()
@@ -670,6 +677,8 @@ def clear_templates(self, templates, host):
                 self.api.host.update(hostid=host["hostid"], templates_clear=templates)
             except pyzabbix.ZabbixAPIException as e:
                 logging.error("Error when clearing templates on host '%s': %s", host["host"], e.args)
+        else:
+            logging.debug("DRYRUN: Clearing templates on host: '%s'", host["host"])
 
     def set_templates(self, templates, host):
         logging.debug("Setting templates on host: '%s'", host["host"])
@@ -679,6 +688,8 @@ def set_templates(self, templates, host):
                 self.api.host.update(hostid=host["hostid"], templates=templates)
             except pyzabbix.ZabbixAPIException as e:
                 logging.error("Error when setting templates on host '%s': %s", host["host"], e.args)
+        else:
+            logging.debug("DRYRUN: Setting templates on host: '%s'", host["host"])
 
     def do_update(self):
         managed_template_names = set(itertools.chain.from_iterable(self.property_template_map.values()))
@@ -741,36 +752,70 @@ def do_update(self):
 class ZabbixHostgroupUpdater(ZabbixUpdater):
 
     def set_hostgroups(self, hostgroups, host):
-        logging.debug("Setting hostgroups on host: '%s'", host["host"])
         if not self.config.dryrun:
+            logging.debug("Setting hostgroups on host: '%s'", host["host"])
             try:
                 groups = [{"groupid": hostgroup_id} for _, hostgroup_id in hostgroups.items()]
                 self.api.host.update(hostid=host["hostid"], groups=groups)
             except pyzabbix.ZabbixAPIException as e:
                 logging.error("Error when setting hostgroups on host '%s': %s", host["host"], e.args)
+        else:
+            logging.debug("DRYRUN: Setting hostgroups on host: '%s'", host["host"])
 
     def create_hostgroup(self, hostgroup_name):
         if not self.config.dryrun:
+            logging.debug("Creating hostgroup: '%s'", hostgroup_name)
             try:
                 result = self.api.hostgroup.create(name=hostgroup_name)
                 return result["groupids"][0]
             except pyzabbix.ZabbixAPIException as e:
                 logging.error("Error when creating hostgroups '%s': %s", hostgroup_name, e.args)
         else:
+            logging.debug("DRYRUN: Creating hostgroup: '%s'", hostgroup_name)
             return "-1"
 
+    def create_extra_hostgroups(
+        self, existing_hostgroups: List[Dict[str, str]]
+    ) -> None:
+        """Creates additonal host groups based on the prefixes specified
+        in the config file. These host groups are not assigned hosts by ZAC."""
+        hostgroup_names = [h["name"] for h in existing_hostgroups]
+
+        for prefix in self.config.extra_siteadmin_hostgroup_prefixes:
+            mapping = utils.mapping_values_with_prefix(
+                self.siteadmin_hostgroup_map,  # this is copied in the function
+                prefix=prefix,
+            )
+            for hostgroups in mapping.values():
+                for hostgroup in hostgroups:
+                    if hostgroup in hostgroup_names:
+                        continue
+                    self.create_hostgroup(hostgroup)
+
     def do_update(self):
-        managed_hostgroup_names = set(itertools.chain.from_iterable(self.property_hostgroup_map.values()))
-        managed_hostgroup_names.update(itertools.chain.from_iterable(self.siteadmin_hostgroup_map.values()))
+        managed_hostgroup_names = set(
+            itertools.chain.from_iterable(self.property_hostgroup_map.values())
+        )
+        managed_hostgroup_names.update(
+            itertools.chain.from_iterable(self.siteadmin_hostgroup_map.values())
+        )
+
+        existing_hostgroups = self.api.hostgroup.get(output=["name", "groupid"])
+
+        # Create extra host groups if necessary
+        if self.config.extra_siteadmin_hostgroup_prefixes:
+            self.create_extra_hostgroups(existing_hostgroups)
+
         zabbix_hostgroups = {}
-        for zabbix_hostgroup in self.api.hostgroup.get(output=["name", "groupid"]):
+        for zabbix_hostgroup in existing_hostgroups:
             zabbix_hostgroups[zabbix_hostgroup["name"]] = zabbix_hostgroup["groupid"]
             if zabbix_hostgroup["name"].startswith(self.config.hostgroup_source_prefix):
                 managed_hostgroup_names.add(zabbix_hostgroup["name"])
             if zabbix_hostgroup["name"].startswith(self.config.hostgroup_importance_prefix):
                 managed_hostgroup_names.add(zabbix_hostgroup["name"])
         managed_hostgroup_names.update([self.config.hostgroup_all])
 
+
         with self.db_connection, self.db_connection.cursor() as db_cursor:
             db_cursor.execute(f"SELECT data FROM {self.db_hosts_table} WHERE data->>'enabled' = 'true'")
             db_hosts = {t[0]["hostname"]: models.Host(**t[0]) for t in db_cursor.fetchall()}

zabbix_auto_config/utils.py

@@ -1,8 +1,9 @@
+import copy
 import ipaddress
 import logging
 from pathlib import Path
 import re
-from typing import Dict, Iterable, List, Sequence, Set, Tuple, Union
+from typing import Dict, Iterable, List, Mapping, MutableMapping, Set, Tuple, Union
 
 
 def is_valid_regexp(pattern: str):
@@ -84,3 +85,68 @@ def read_map_file(path: Union[str, Path]) -> Dict[str, List[str]]:
             )
         _map[key] = values_dedup
     return _map
+
+
+def with_prefix(
+    text: str,
+    prefix: str,
+    separator: str = "-",
+) -> str:
+    """Replaces the prefix of `text` with `prefix`. Assumes the separator
+    between the prefix and the text is `separator` (default: "-").
+
+    Parameters
+    ----
+    text: str
+        The text to format.
+    prefix: str
+        The prefix to add to `text`.
+    separator: str
+        The separator between the prefix and the text.
+
+    Returns
+    -------
+    str
+        The formatted string.
+    """
+    if not all(s for s in (text, prefix, separator)):
+        raise ValueError("Text, prefix, and separator cannot be empty")
+
+    _, _, suffix = text.partition(separator)
+
+    # Unable to split text, nothing to do
+    if not suffix:
+        raise ValueError(
+            f"Could not find prefix in {text!r} with separator {separator!r}"
+        )
+
+    groupname = f"{prefix}{suffix}"
+    if not prefix.endswith(separator) and not suffix.startswith(separator):
+        logging.warning(
+            "Prefix '%s' for group name '%s' does not contain separator '%s'",
+            prefix,
+            groupname,
+            separator,
+        )
+    return groupname
+
+def mapping_values_with_prefix(
+    m: MutableMapping[str, List[str]],
+    prefix: str,
+    separator: str = "-",
+) -> MutableMapping[str, List[str]]:
+    """Calls `with_prefix` on all items in the values (list) in the mapping `m`."""
+    m = copy.copy(m) # don't modify the original mapping
+    for key, value in m.items():
+        new_values = []
+        for v in value:
+            try:
+                new_value = with_prefix(text=v, prefix=prefix, separator=separator)
+            except ValueError:
+                logging.warning(
+                    f"Unable to replace prefix in '%s' with '%s'", v, prefix
+                )
+                continue
+            new_values.append(new_value)
+        m[key] = new_values
+    return m