Add machine tags from jamf groups

tests/inventory/test_machine_snapshot.py

@@ -2,6 +2,7 @@
 from datetime import datetime, timedelta
 from dateutil import parser
 from django.test import TestCase
+from django.utils.crypto import get_random_string
 from django.utils.timezone import is_aware, make_naive
 from zentral.contrib.inventory.conf import DESKTOP, MACOS, SERVER, update_ms_tree_type, VM
 from zentral.contrib.inventory.models import (BusinessUnit,
@@ -12,7 +13,7 @@
                                               MetaBusinessUnitTag,
                                               MetaMachine,
                                               Source,
-                                              Tag)
+                                              Tag, Taxonomy)
 from zentral.utils.mt_models import MTOError
 
 
@@ -267,6 +268,42 @@ def test_meta_machine(self):
         self.assertEqual(MachineSnapshotCommit.objects.count(), 3)
         self.assertEqual(CurrentMachineSnapshot.objects.count(), 0)
 
+    def test_meta_machine_update_taxonomy_tags(self):
+        # one machine
+        serial_number = get_random_string(13)
+        # two tags from taxonomy1
+        taxonomy1 = Taxonomy.objects.create(name=get_random_string(34))
+        tag11 = Tag.objects.create(taxonomy=taxonomy1, name=get_random_string(17))
+        MachineTag.objects.get_or_create(serial_number=serial_number, tag=tag11)
+        tag12 = Tag.objects.create(taxonomy=taxonomy1, name=get_random_string(18))
+        MachineTag.objects.get_or_create(serial_number=serial_number, tag=tag12)
+        # one tag from taxonomy2
+        taxonomy2 = Taxonomy.objects.create(name=get_random_string(27))
+        tag21 = Tag.objects.create(taxonomy=taxonomy2, name=get_random_string(20))
+        MachineTag.objects.get_or_create(serial_number=serial_number, tag=tag21)
+        # one detached tag
+        tag31 = Tag.objects.create(name=get_random_string(21))
+        MachineTag.objects.get_or_create(serial_number=serial_number, tag=tag31)
+        # update the taxonomy1  tags. keep one, add two new ones, one collision, remove one.
+        new_tag_names = [get_random_string(22), get_random_string(33)]
+        updated_tag_names = [
+            tag11.name,  # existing,
+            # tag12.name  # removed
+            tag31.name,  # collision, because we will try to add a tag with the same name, but within the taxonomy1
+        ] + new_tag_names  # new ones
+        mm = MetaMachine(serial_number)
+        mm.update_taxonomy_tags(taxonomy1, updated_tag_names)
+        # verify
+        # two new tags
+        new_tags = list(Tag.objects.filter(name__in=new_tag_names))
+        self.assertEqual(len(new_tags), 2)
+        # in the taxonomy1
+        self.assertTrue(all(t.taxonomy == taxonomy1 for t in new_tags))
+        # expected tags for the machine
+        expected_tags = [("machine", t)
+                         for t in [tag11, tag21, tag31] + new_tags]
+        self.assertEqual(set(expected_tags), set(mm.tags_with_types))
+
     def test_machine_name(self):
         tree = {"source": {"module": "godzilla",
                            "name": "test"},

zentral/contrib/inventory/models.py

@@ -866,6 +866,28 @@ def available_tags(self):
         tags.sort(key=lambda t: (t.meta_business_unit is None, str(t).upper()))
         return tags
 
+    def update_taxonomy_tags(self, taxonomy, tag_names):
+        tag_names = set(tag_names)
+        with transaction.atomic():
+            existing_machine_tags = (MachineTag.objects.select_for_update()
+                                                       .select_related("tag")
+                                                       .filter(serial_number=self.serial_number,
+                                                               tag__taxonomy=taxonomy))
+            existing_tag_names = set(mt.tag.name for mt in existing_machine_tags)
+            # delete old tags
+            tag_names_to_delete = existing_tag_names - tag_names
+            if tag_names_to_delete:
+                existing_machine_tags.filter(tag__name__in=tag_names_to_delete).delete()
+            # add missing tags
+            for tag_name in tag_names - existing_tag_names:
+                try:
+                    with transaction.atomic():
+                        tag, _ = Tag.objects.get_or_create(taxonomy=taxonomy, name=tag_name)
+                except IntegrityError:
+                    logger.error("Tag collision, taxonomy '%s', name '%s'", taxonomy.pk, tag_name)
+                else:
+                    MachineTag.objects.get_or_create(serial_number=self.serial_number, tag=tag)
+
     def max_incident_severity(self):
         return (MachineIncident.objects.select_related("incident")
                                        .filter(serial_number=self.serial_number, status__in=OPEN_STATUSES)

zentral/contrib/jamf/api_client.py

@@ -1,5 +1,6 @@
 import logging
 from dateutil import parser
+import re
 from urllib.parse import urlparse
 from xml.etree import ElementTree as ET
 from xml.sax.saxutils import escape as xml_escape
@@ -43,6 +44,13 @@ def __init__(self, host, port, path, user, password, secret, business_unit=None,
                            requests.adapters.HTTPAdapter(max_retries=max_retries))
         self.mobile_device_groups = {}
         self.reverse_computer_groups = {}
+        self.group_tag_regex = None
+        # tags from groups
+        self.tag_configs = []
+        for tag_config in kwargs.get("tag_configs", []):
+            tag_config = tag_config.copy()
+            tag_config["regex"] = re.compile(tag_config.pop("regex"))
+            self.tag_configs.append(tag_config)
 
     def get_source_d(self):
         return {"module": "zentral.contrib.jamf",
@@ -169,6 +177,28 @@ def get_machine_d(self, device_type, jamf_id):
         else:
             raise APIClientError("Unknown device type %s", device_type)
 
+    def get_machine_d_and_tags(self, device_type, jamf_id):
+        machine_d = self.get_machine_d(device_type, jamf_id)
+        tags = {}
+        groups = None
+        for tag_config in self.tag_configs:
+            tag_names = tags.setdefault(tag_config["taxonomy_id"], [])
+            if tag_config["source"] == "GROUP":
+                if groups is None:
+                    groups = machine_d.get("groups", [])
+                for group in groups:
+                    regex = tag_config["regex"]
+                    group_name = group["name"]
+                    if regex.match(group_name):
+                        tag_name = regex.sub(tag_config["replacement"], group_name)
+                        if tag_name:
+                            tag_names.append(tag_name)
+                        else:
+                            logger.error("Empty group tag name %s %s %s", device_type, jamf_id, regex)
+            else:
+                logger.error("Unknown tag config source: %s", tag_config["source"])
+        return machine_d, tags
+
     def get_computer_machine_d(self, jamf_id):
         computer = self._computer(jamf_id)
         # serial number, reference

zentral/contrib/jamf/forms.py

@@ -1,12 +1,17 @@
+import re
 from django import forms
 from zentral.contrib.inventory.models import BusinessUnit
-from .models import JamfInstance
+from .models import JamfInstance, TagConfig
 
 
 class JamfInstanceForm(forms.ModelForm):
     class Meta:
         model = JamfInstance
-        fields = ("business_unit", "host", "port", "path", "user", "password")
+        fields = (
+            "business_unit",
+            "host", "port", "path",
+            "user", "password",
+        )
         widgets = {
             'password': forms.PasswordInput(render_value=True)
         }
@@ -17,3 +22,18 @@ def __init__(self, *args, **kwargs):
             BusinessUnit.objects.filter(source__module="zentral.contrib.inventory")
                                 .order_by('meta_business_unit__name')
         )
+
+
+class TagConfigForm(forms.ModelForm):
+    class Meta:
+        model = TagConfig
+        fields = ("source", "taxonomy", "regex", "replacement")
+
+    def clean_regex(self):
+        regex = self.cleaned_data["regex"]
+        if regex:
+            try:
+                re.compile(regex)
+            except re.error:
+                raise forms.ValidationError("Not a valid regex")
+        return regex

zentral/contrib/jamf/migrations/0002_tagconfig.py

@@ -0,0 +1,34 @@
+# Generated by Django 2.2.17 on 2020-11-19 14:34
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('inventory', '0049_machine_snapshot_certificates_on_delete_cascade'),
+        ('jamf', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='TagConfig',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('source', models.CharField(choices=[('GROUP', 'Group')], default='GROUP', max_length=16)),
+                ('regex', models.CharField(
+                     help_text='matching names will be used to automatically generate tags',
+                     max_length=256
+                 )),
+                ('replacement', models.CharField(
+                     help_text='replacement pattern used to generate a tag name from a tag regex match',
+                     max_length=32
+                 )),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('updated_at', models.DateTimeField(auto_now=True)),
+                ('instance', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='jamf.JamfInstance')),
+                ('taxonomy', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='inventory.Taxonomy')),
+            ],
+        ),
+    ]

zentral/contrib/jamf/models.py

@@ -3,6 +3,7 @@
 from django.core.validators import MinValueValidator, MaxValueValidator
 from django.db import models
 from django.db.models import F
+from django.urls import reverse
 from django.utils.crypto import get_random_string
 
 
@@ -39,6 +40,9 @@ class Meta:
     def __str__(self):
         return self.host
 
+    def get_absolute_url(self):
+        return reverse("jamf:jamf_instance", args=(self.pk,))
+
     def save(self, *args, **kwargs):
         if not self.id:
             self.version = 0
@@ -65,6 +69,7 @@ def serialize(self):
             "user": self.user,
             "password": self.password,
             "secret": self.secret,
+            "tag_configs": [tm.serialize() for tm in self.tagconfig_set.select_related("taxonomy").all()],
         }
         if self.business_unit:
             d["business_unit"] = self.business_unit.serialize()
@@ -76,3 +81,36 @@ def observer_dict(self):
                 "type": "Jamf Pro",
                 "content_type": "jamf.jamfinstance",
                 "pk": self.pk}
+
+
+class TagConfig(models.Model):
+    GROUP_SOURCE = "GROUP"
+    SOURCE_CHOICES = (
+        (GROUP_SOURCE, "Group"),
+    )
+    instance = models.ForeignKey(JamfInstance, on_delete=models.CASCADE)
+    source = models.CharField(max_length=16, choices=SOURCE_CHOICES, default=GROUP_SOURCE)
+    taxonomy = models.ForeignKey("inventory.Taxonomy", on_delete=models.CASCADE)
+    regex = models.CharField(
+        max_length=256,
+        help_text="matching names will be used to automatically generate tags"
+    )
+    replacement = models.CharField(
+        max_length=32,
+        help_text="replacement pattern used to generate a tag name from a tag regex match"
+    )
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
+
+    def get_absolute_url(self):
+        return "{}#tag-config-{}".format(self.instance.get_absolute_url(), self.pk)
+
+    def save(self, *args, **kwargs):
+        super().save(*args, **kwargs)
+        self.instance.save()
+
+    def serialize(self):
+        return {"source": self.source,
+                "taxonomy_id": self.taxonomy.id,
+                "regex": self.regex,
+                "replacement": self.replacement}

zentral/contrib/jamf/preprocessors/webhook.py

@@ -1,5 +1,7 @@
 import logging
-from zentral.contrib.inventory.models import MachineGroup, MachineSnapshot, MachineSnapshotCommit
+from django.db import transaction
+from zentral.contrib.inventory.models import (MachineGroup, MachineSnapshot, MachineSnapshotCommit,
+                                              MetaMachine, Taxonomy)
 from zentral.contrib.inventory.utils import inventory_events_from_machine_snapshot_commit
 from zentral.contrib.jamf.api_client import APIClient
 from zentral.core.events import event_cls_from_type
@@ -14,6 +16,7 @@ class WebhookEventPreprocessor(object):
 
     def __init__(self):
         self.clients = {}
+        self.taxonomies = {}
 
     def get_client(self, jamf_instance_d):
         key = (jamf_instance_d["pk"], jamf_instance_d["version"])
@@ -23,6 +26,14 @@ def get_client(self, jamf_instance_d):
             self.clients[key] = client
         return client
 
+    def get_taxonomy(self, taxonomy_id):
+        if taxonomy_id not in self.taxonomies:
+            try:
+                self.taxonomies[taxonomy_id] = Taxonomy.objects.get(pk=taxonomy_id)
+            except Taxonomy.DoesNotExist:
+                logger.error("Could not get taxonomy %s", taxonomy_id)
+        return self.taxonomies.get(taxonomy_id)
+
     def is_known_machine(self, client, serial_number):
         kwargs = {"serial_number": serial_number}
         for k, v in client.get_source_d().items():
@@ -32,13 +43,14 @@ def is_known_machine(self, client, serial_number):
     def update_machine(self, client, device_type, jamf_id):
         logger.info("Update machine %s %s %s", client.get_source_d(), device_type, jamf_id)
         try:
-            machine_d = client.get_machine_d(device_type, jamf_id)
+            machine_d, tags = client.get_machine_d_and_tags(device_type, jamf_id)
         except Exception:
-            logger.exception("Could not get machine_d. %s %s %s",
+            logger.exception("Could not get machine_d and tags. %s %s %s",
                              client.get_source_d(), device_type, jamf_id)
         else:
             try:
-                msc, ms = MachineSnapshotCommit.objects.commit_machine_snapshot_tree(machine_d)
+                with transaction.atomic():
+                    msc, ms = MachineSnapshotCommit.objects.commit_machine_snapshot_tree(machine_d)
             except Exception:
                 logger.exception("Could not commit machine snapshot")
             else:
@@ -53,6 +65,12 @@ def update_machine(self, client, device_type, jamf_id):
                                                  tags=event_cls.tags)
                         event = event_cls(metadata, payload)
                         yield event
+            if tags:
+                machine = MetaMachine(machine_d["serial_number"])
+                for taxonomy_id, tag_names in tags.items():
+                    taxonomy = self.get_taxonomy(taxonomy_id)
+                    if taxonomy:
+                        machine.update_taxonomy_tags(taxonomy, tag_names)
 
     def get_inventory_groups(self, client, device_type, jamf_id, is_smart):
         kwargs = {"reference": client.group_reference(device_type, jamf_id, is_smart)}