Merge pull request #3927 from modoboa/fix/rspamd_maps_generation

Improved dkim key creation to avoid duplicate runs.

modoboa/admin/api/v2/serializers.py

@@ -111,6 +111,8 @@ class DomainSerializer(v1_serializers.DomainSerializer):
         domain_admin = validated_data.pop("domain_admin", None)
         transport_def = validated_data.pop("transport", None)
         domain = super().create(validated_data)
+        if domain.enable_dkim:
+            domain.generate_dkim_key()
         if transport_def:
             transport = transport_models.Transport(**transport_def)
             transport.pattern = domain.name
@@ -182,6 +184,13 @@ class DomainSerializer(v1_serializers.DomainSerializer):
         transport_def = validated_data.pop("transport", None)
         resources = validated_data.pop("domainobjectlimit_set", None)
         instance = super().update(instance, validated_data)
+        if (
+            instance._loaded_values.get("enable_dkim") != instance.enable_dkim
+            and instance.enable_dkim
+        ) or (
+            instance._loaded_values.get("dkim_key_length") != instance.dkim_key_length
+        ):
+            instance.generate_dkim_key()
         if transport_def and instance.type == "relaydomain":
             transport = getattr(instance, "transport", None)
             created = False

modoboa/admin/handlers.py

@@ -1,13 +1,10 @@
 """Django signal handlers for admin."""
 
 from django.contrib.contenttypes.models import ContentType
-from django.core.management import call_command
 from django.db.models import signals
 from django.dispatch import receiver
 from django.utils.translation import gettext as _
 
-import django_rq
-
 from modoboa.core import models as core_models, signals as core_signals
 from modoboa.lib import exceptions, permissions, signals as lib_signals
 from modoboa.lib.email_utils import split_mailbox
@@ -42,14 +39,6 @@ def update_domain_mxs_and_mailboxes(sender, instance, **kwargs):
         alias.save(update_fields=["address"])
 
 
-@receiver(signals.post_save, sender=models.Domain)
-def create_dkim_key(sender, instance, **kwargs):
-    if not instance.enable_dkim:
-        return
-    queue = django_rq.get_queue("dkim")
-    queue.enqueue(call_command, "modo", "manage_dkim_keys", f"--domain={instance.name}")
-
-
 @receiver(signals.post_save, sender=models.DomainAlias)
 def create_alias_for_domainalias(sender, instance, **kwargs):
     """Create a dedicated alias for domain alias."""

modoboa/admin/management/commands/subcommands/_manage_dkim_keys.py

@@ -6,6 +6,7 @@ from django.core.management.base import BaseCommand
 from django.utils.encoding import smart_str
 from django.utils.translation import gettext as _
 
+from modoboa.admin import signals
 from modoboa.lib import sysutils
 from modoboa.parameters import tools as param_tools
 
@@ -86,14 +87,19 @@ class ManageDKIMKeys(BaseCommand):
             "dkim_default_key_length"
         )
 
+        domains = []
         if options["domain"] != "":
             domain = models.Domain.objects.filter(
                 name=options["domain"], enable_dkim=True, dkim_private_key_path=""
+            ).first()
+            if domain:
+                self.create_new_dkim_key(domain)
+                domains.append(domain)
+        else:
+            qset = models.Domain.objects.filter(
+                enable_dkim=True, dkim_private_key_path=""
             )
-            if domain.exists():
-                self.create_new_dkim_key(domain[0])
-            return
-
-        qset = models.Domain.objects.filter(enable_dkim=True, dkim_private_key_path="")
-        for domain in qset:
-            self.create_new_dkim_key(domain)
+            for domain in qset:
+                self.create_new_dkim_key(domain)
+                domains.append(domain)
+        signals.dkim_keys_created.send(self, domains=domains)

modoboa/admin/models/domain.py

@@ -4,12 +4,15 @@ import datetime
 
 from reversion import revisions as reversion
 
+from django.core.management import call_command
 from django.db import models
 from django.utils import timezone
 from django.utils.encoding import force_str, smart_str
 from django.utils.functional import cached_property
 from django.utils.translation import gettext as _, gettext_lazy
 
+import django_rq
+
 from modoboa.core import signals as core_signals
 from modoboa.core.models import User
 from modoboa.lib.exceptions import BadRequest, Conflict
@@ -305,6 +308,11 @@ class Domain(mixins.MessageLimitMixin, AdminObject):
     def __str__(self):
         return smart_str(self.name)
 
+    def generate_dkim_key(self):
+        """Launch domain key creation."""
+        queue = django_rq.get_queue("dkim")
+        queue.enqueue(call_command, "modo", "manage_dkim_keys", f"--domain={self.name}")
+
     def from_csv(self, user, row):
         """Create a new domain from a CSV entry.
 

modoboa/admin/signals.py

@@ -7,3 +7,4 @@ extra_domain_qset_filters = django.dispatch.Signal()  # Provides domfilter, extr
 import_object = django.dispatch.Signal()  # Provides objtype
 use_external_recipients = django.dispatch.Signal()  # Provides recipients
 extra_account_identities_actions = django.dispatch.Signal()  # Provides account
+dkim_keys_created = django.dispatch.Signal()  # Provides domains

modoboa/admin/tests/test_domain.py

@@ -4,11 +4,14 @@ import os
 import shutil
 import tempfile
 
+from rq import SimpleWorker
 from testfixtures import compare
 
 from django.core.management import call_command
 from django.urls import reverse
 
+import django_rq
+
 from modoboa.core.models import User
 from modoboa.lib.tests import ModoAPITestCase, SETTINGS_SAMPLE
 
@@ -194,7 +197,7 @@ class DKIMTestCase(ModoAPITestCase):
         self.assertEqual(domain.dkim_private_key_path, "")
 
     def test_dkim_key_length_modification(self):
-        """ """
+        """Check that key length change triggers a new key creation."""
         self.set_global_parameter("dkim_keys_storage_dir", self.workdir)
         values = {
             "name": "pouet.com",
@@ -207,9 +210,12 @@ class DKIMTestCase(ModoAPITestCase):
         url = reverse("v2:domain-list")
         response = self.client.post(url, values, format="json")
         self.assertEqual(response.status_code, 201)
-        call_command("modo", "manage_dkim_keys")
+        queue = django_rq.get_queue("dkim")
+        worker = SimpleWorker([queue], connection=queue.connection)
+        worker.work(burst=True)
         key_path = os.path.join(self.workdir, "{}.pem".format(values["name"]))
         self.assertTrue(os.path.exists(key_path))
+
         dom = Domain.objects.get(name="pouet.com")
         values["dkim_key_length"] = 4096
         response = self.client.put(
@@ -219,5 +225,5 @@ class DKIMTestCase(ModoAPITestCase):
         dom.refresh_from_db()
         self.assertEqual(dom.dkim_private_key_path, "")
         os.unlink(key_path)
-        call_command("modo", "manage_dkim_keys")
+        worker.work(burst=True)
         self.assertTrue(os.path.exists(key_path))

modoboa/core/commands/templates/cron_config.py.tpl

@@ -2,7 +2,6 @@ from rq.cron import register
 
 from django.conf import settings
 
-from modoboa.amavis import jobs as amavis_jobs
 from modoboa.calendars import jobs as calendars_jobs
 from modoboa.core import jobs as core_jobs
 from modoboa.admin import jobs as admin_jobs
@@ -29,5 +28,6 @@ register(calendars_jobs.generate_rights, queue_name="privileged", cron="*/2 * *
 register(webmail_jobs.send_scheduled_messages, queue_name="modoboa", cron="* * * * *")
 
 if "modoboa.amavis" in settings.MODOBOA_APPS:
+    from modoboa.amavis import jobs as amavis_jobs
     register(amavis_jobs.qcleanup, queue_name="modoboa", cron="0 0 * * *")
     register(amavis_jobs.amnotify, queue_name="modoboa", cron="0 12 * * *")

modoboa/rspamd/handlers.py

@@ -1,27 +1,15 @@
 """Modoboa rspamd handlers."""
 
-from django.core.management import call_command
 from django.dispatch import receiver
-from django.db.models.signals import post_save
 
 import django_rq
 
-from modoboa.admin import models as admin_models
+from modoboa.admin import signals
+from modoboa.rspamd import jobs
 
 
-@receiver(post_save, sender=admin_models.Domain)
-def update_rspamd_dkim_maps(sender, instance, created, **kwargs):
+@receiver(signals.dkim_keys_created)
+def update_rspamd_dkim_maps(sender, domains, **kwargs):
     """Update config maps."""
-    # Modify or create
-    condition = (
-        instance._loaded_values.get("dkim_private_key_path")
-        != instance.dkim_private_key_path
-        or instance._loaded_values.get("dkim_key_selector")
-        != instance.dkim_key_selector
-        or instance._loaded_values.get("enable_dkim") != instance.enable_dkim
-    )
-    if condition:
-        # TODO : check how to handle it
-        # Either with modoboa user or _rspamd...
-        queue = django_rq.get_queue("modoboa")
-        queue.enqueue(call_command, "manage_rspamd_maps", f"--domain={instance.name}")
+    queue = django_rq.get_queue("modoboa")
+    queue.enqueue(jobs.update_rspamd_maps, [domain.id for domain in domains])

modoboa/rspamd/jobs.py

@@ -1,23 +1,30 @@
-"""Management command to create DKIM keys."""
+"""Async jobs."""
 
-from django.core.management.base import BaseCommand, CommandError
+import logging
 
-from modoboa.parameters import tools as param_tools
 from modoboa.admin import models
+from modoboa.parameters import tools as param_tools
+
+logger = logging.getLogger("modoboa.jobs")
 
 
-class Command(BaseCommand):
-    """Command class."""
+class MapsUpdater:
+
+    def __init__(self) -> None:
+        self.modified_selector_file = False
+        self.modified_key_path_file = False
+        self.config = dict(param_tools.get_global_parameters("rspamd"))
+        self.dkim_path_map: dict[str, str] = {}
+        self.selector_map: dict[str, str] = {}
 
     def load_files(self):
-        self.config = dict(param_tools.get_global_parameters("rspamd"))
         if not self.config["key_map_path"] or not self.config["selector_map_path"]:
-            raise CommandError(
+            logger.error(
                 "path map path and/or selector map path "
                 "not set in modoboa rspamd settings."
             )
+            return False
 
-        self.dkim_path_map = {}
         try:
             with open(self.config["key_map_path"]) as f:
                 for line in f:
@@ -25,7 +32,6 @@ class Command(BaseCommand):
                     self.dkim_path_map[domain_name] = path.replace("\n", "")
         except FileNotFoundError:
             pass
-        self.selector_map = {}
         try:
             with open(self.config["selector_map_path"]) as f:
                 for line in f:
@@ -33,6 +39,7 @@ class Command(BaseCommand):
                     self.selector_map[domain_name] = selector
         except FileNotFoundError:
             pass
+        return True
 
     def manage_domain(self, domain_instance):
         domain_name = domain_instance.name
@@ -65,30 +72,13 @@ class Command(BaseCommand):
             self.dkim_path_map[domain_name] = domain_instance.dkim_private_key_path
             self.modified_key_path_file = True
 
-    def add_arguments(self, parser):
-        """Add arguments to command."""
-        parser.add_argument(
-            "--domain",
-            type=str,
-            dest="domain",
-            default="",
-            help="Domain target for keys generation.",
-        )
-
-    def handle(self, *args, **options):
+    def run(self, domains):
         """Entry point."""
-        self.modified_selector_file = False
-        self.modified_key_path_file = False
-        self.load_files()
+        if not self.load_files():
+            return
 
-        if options["domain"] != "":
-            domain = models.Domain.objects.filter(name=options["domain"])
-            if domain.exists():
-                self.manage_domain(domain[0])
-        else:
-            qset = models.Domain.objects.all()
-            for domain in qset:
-                self.manage_domain(domain)
+        for domain in domains:
+            self.manage_domain(domain)
 
         if self.modified_selector_file:
             with open(self.config["selector_map_path"], "w") as f:
@@ -98,3 +88,9 @@ class Command(BaseCommand):
             with open(self.config["key_map_path"], "w") as f:
                 for domain_name, key_path in self.dkim_path_map.items():
                     f.write(f"{domain_name} {key_path}\n")
+
+
+def update_rspamd_maps(domain_ids: list[int]) -> None:
+    """Launch map files update."""
+    domains = models.Domain.objects.filter(id__in=domain_ids)
+    MapsUpdater().run(domains)

modoboa/rspamd/tests.py

@@ -1,12 +1,17 @@
 import os
 import shutil
 
-from django.core import management
+from rq import SimpleWorker
+from testfixtures import LogCapture
+
 from django.test import modify_settings
 from django.urls import reverse
 
+import django_rq
+
 from modoboa.admin import factories as admin_factories, models as admin_models
 from modoboa.lib.tests import ModoAPITestCase
+from modoboa.rspamd import jobs
 
 
 @modify_settings(
@@ -35,13 +40,25 @@ class ManagementCommandTestCase(ModoAPITestCase):
         self.set_global_parameter("key_map_path", "")
         self.set_global_parameter("selector_map_path", "")
 
-    def test_command(self):
+    def test_job(self):
         self.unconfigure()
-        with self.assertRaises(management.CommandError):
-            management.call_command("manage_rspamd_maps")
+        with LogCapture("modoboa.jobs") as log:
+            jobs.update_rspamd_maps(
+                list(admin_models.Domain.objects.all().values_list("id", flat=True))
+            )
+        log.check(
+            (
+                "modoboa.jobs",
+                "ERROR",
+                "path map path and/or selector map path "
+                "not set in modoboa rspamd settings.",
+            )
+        )
 
         self.configure()
-        management.call_command("manage_rspamd_maps")
+        jobs.update_rspamd_maps(
+            list(admin_models.Domain.objects.all().values_list("id", flat=True))
+        )
         self.assertFalse(os.path.exists(self.key_map_path))
         self.assertFalse(os.path.exists(self.selector_map_path))
 
@@ -49,24 +66,44 @@ class ManagementCommandTestCase(ModoAPITestCase):
         domain.enable_dkim = True
         domain.dkim_private_key_path = "xxx"
         domain.save()
-        management.call_command("manage_rspamd_maps")
+        jobs.update_rspamd_maps(
+            list(admin_models.Domain.objects.all().values_list("id", flat=True))
+        )
         self.assertTrue(os.path.exists(self.key_map_path))
         self.assertTrue(os.path.exists(self.selector_map_path))
 
         # Now, empty map files
         domain.enable_dkim = False
         domain.save()
-        management.call_command("manage_rspamd_maps")
+        jobs.update_rspamd_maps(
+            list(admin_models.Domain.objects.all().values_list("id", flat=True))
+        )
         with open(self.key_map_path) as fp:
             content = fp.read()
         self.assertNotIn(domain.name, content)
 
-    def test_command_with_arg(self):
+    def test_signal_handler(self):
+        self.set_global_parameter("dkim_keys_storage_dir", self.workdir, app="admin")
         self.configure()
-        domain = admin_models.Domain.objects.get(name="test.com")
-        domain.enable_dkim = True
-        domain.save()
-        management.call_command("manage_rspamd_maps", "--domain", "test.com")
+        values = {
+            "name": "pouet.com",
+            "quota": 1000,
+            "default_mailbox_quota": 100,
+            "type": "domain",
+            "enable_dkim": True,
+        }
+        url = reverse("v2:domain-list")
+        response = self.client.post(url, values, format="json")
+        self.assertEqual(response.status_code, 201)
+
+        queue = django_rq.get_queue("dkim")
+        worker = SimpleWorker([queue], connection=queue.connection)
+        worker.work(burst=True)
+
+        queue = django_rq.get_queue("modoboa")
+        worker = SimpleWorker([queue], connection=queue.connection)
+        worker.work(burst=True)
+
         self.assertTrue(os.path.exists(self.key_map_path))
         self.assertTrue(os.path.exists(self.selector_map_path))