Have custom domains set up multiple dkim records to be able to rotate keys (#1334)

* Have custom domains set up multiple dkim records to be able to rotate keys

* Apply suggestions from code review

* Some PR comments

* Keep dkim enabled if it is already

* Format

* PR updates

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>

app/custom_domain_validation.py

@@ -0,0 +1,37 @@
+from app.db import Session
+from app.dns_utils import get_cname_record
+from app.models import CustomDomain
+
+
+class CustomDomainValidation:
+    def __init__(self, dkim_domain: str):
+        self.dkim_domain = dkim_domain
+        self._dkim_records = {
+            (f"{key}._domainkey", f"{key}._domainkey.{self.dkim_domain}")
+            for key in ("dkim", "dkim02", "dkim03")
+        }
+
+    def get_dkim_records(self) -> {str: str}:
+        """
+        Get a list of dkim records to set up. It will be
+
+        """
+        return self._dkim_records
+
+    def validate_dkim_records(self, custom_domain: CustomDomain) -> dict[str, str]:
+        """
+        Check if dkim records are properly set for this custom domain.
+        Returns empty list if all records are ok. Other-wise return the records that aren't properly configured
+        """
+        invalid_records = {}
+        for prefix, expected_record in self.get_dkim_records():
+            custom_record = f"{prefix}.{custom_domain.domain}"
+            dkim_record = get_cname_record(custom_record)
+            if dkim_record != expected_record:
+                invalid_records[custom_record] = dkim_record or "empty"
+        # HACK: If dkim is enabled, don't disable it to give users time to update their CNAMES
+        if custom_domain.dkim_verified:
+            return invalid_records
+        custom_domain.dkim_verified = len(invalid_records) == 0
+        Session.commit()
+        return invalid_records

app/dashboard/views/domain_detail.py

@@ -7,13 +7,13 @@ from flask_wtf import FlaskForm
 from wtforms import StringField, validators, IntegerField
 
 from app.config import EMAIL_SERVERS_WITH_PRIORITY, EMAIL_DOMAIN, JOB_DELETE_DOMAIN
+from app.custom_domain_validation import CustomDomainValidation
 from app.dashboard.base import dashboard_bp
 from app.db import Session
 from app.dns_utils import (
     get_mx_domains,
     get_spf_domain,
     get_txt_record,
-    get_cname_record,
     is_mx_equivalent,
 )
 from app.log import LOG
@@ -46,8 +46,7 @@ def domain_detail_dns(custom_domain_id):
 
     spf_record = f"v=spf1 include:{EMAIL_DOMAIN} ~all"
 
-    # hardcode the DKIM selector here
-    dkim_cname = f"dkim._domainkey.{EMAIL_DOMAIN}"
+    domain_validator = CustomDomainValidation(EMAIL_DOMAIN)
 
     dmarc_record = "v=DMARC1; p=quarantine; pct=100; adkim=s; aspf=s"
 
@@ -122,23 +121,17 @@ def domain_detail_dns(custom_domain_id):
                 spf_errors = get_txt_record(custom_domain.domain)
 
         elif request.form.get("form-name") == "check-dkim":
-            dkim_record = get_cname_record("dkim._domainkey." + custom_domain.domain)
-            if dkim_record == dkim_cname:
+            dkim_errors = domain_validator.validate_dkim_records(custom_domain)
+            if len(dkim_errors) == 0:
                 flash("DKIM is setup correctly.", "success")
-                custom_domain.dkim_verified = True
-                Session.commit()
-
                 return redirect(
                     url_for(
                         "dashboard.domain_detail_dns", custom_domain_id=custom_domain.id
                     )
                 )
             else:
-                custom_domain.dkim_verified = False
-                Session.commit()
-                flash("DKIM: the CNAME record is not correctly set", "warning")
                 dkim_ok = False
-                dkim_errors = [dkim_record or "[Empty]"]
+                flash("DKIM: the CNAME record is not correctly set", "warning")
 
         elif request.form.get("form-name") == "check-dmarc":
             txt_records = get_txt_record("_dmarc." + custom_domain.domain)
@@ -164,6 +157,7 @@ def domain_detail_dns(custom_domain_id):
     return render_template(
         "dashboard/domain_detail/dns.html",
         EMAIL_SERVERS_WITH_PRIORITY=EMAIL_SERVERS_WITH_PRIORITY,
+        dkim_records=domain_validator.get_dkim_records(),
         **locals(),
     )
 

templates/dashboard/domain_detail/dns.html

@@ -236,25 +236,28 @@
           folder.
         </div>
         <div class="mb-2">
-          Add the following CNAME DNS record to your domain.
+          Add the following CNAME DNS records to your domain.
         </div>
-        <div class="mb-2 p-3 dns-record">
-          Record: CNAME
-          <br />
-          Domain: <em data-toggle="tooltip"
+        {% for dkim_prefix, dkim_cname_value in dkim_records %}
+
+          <div class="mb-2 p-3 dns-record">
+            Record: CNAME
+            <br />
+            Domain: <em data-toggle="tooltip"
      title="Click to copy"
      class="clipboard"
-     data-clipboard-text="dkim._domainkey">dkim._domainkey</em>
-          <br />
-          Value:
-          <em data-toggle="tooltip"
-              title="Click to copy"
-              class="clipboard"
-              data-clipboard-text="{{ dkim_cname + '.' }}"
-              style="overflow-wrap: break-word">
-            {{ dkim_cname }}.
-          </em>
-        </div>
+     data-clipboard-text="dkim._domainkey">{{ dkim_prefix }}</em>
+            <br />
+            Value:
+            <em data-toggle="tooltip"
+                title="Click to copy"
+                class="clipboard"
+                data-clipboard-text="{{ dkim_cname_value }}."
+                style="overflow-wrap: break-word">
+              {{ dkim_cname_value }}.
+            </em>
+          </div>
+        {% endfor %}
         <div class="alert alert-info">
           Some DNS registrar might require a full record path, in this case please use
           <i>dkim._domainkey.{{ custom_domain.domain }}</i> as domain value instead.
@@ -285,24 +288,28 @@
         {% if not dkim_ok %}
 
           <div class="text-danger mt-4">
-            Your DNS is not correctly set.
-            {% if dkim_errors %}
+            <p>
+              Your DNS is not correctly set.
+            </p>
+            <ul>
+              {% for custom_record, retrieved_cname in dkim_errors.items() %}
 
-              The CNAME record we obtain for
-              <em>dkim._domainkey.{{ custom_domain.domain }}</em> is:
-              <div class="mb-3 p-3 dns-record">
-                {% for r in dkim_errors %}
-
-                  {{ r }}
-                  <br />
-                {% endfor %}
-              </div>
-            {% endif %}
+                <li>
+                  The CNAME record we obtain for <em>{{ custom_record }}</em> is {{ retrieved_cname }}
+                </li>
+              {% endfor %}
+            </ul>
             {% if custom_domain.dkim_verified %}
 
               Without DKIM setup, emails you sent from your alias might end up in Spam/Junk folder.
             {% endif %}
           </div>
+          {% if custom_domain.dkim_verified %}
+
+            <div class="text-danger mt-4">
+              DKIM is still enabled. Please update your DKIM settings with all CNAME records
+            </div>
+          {%  endif %}
         {% endif %}
       </div>
       <hr />