Add job enqueueing for domain scans

.pre-commit-config.yaml

@@ -25,6 +25,7 @@ repos:
     - email-validator==2.1.0.post1
     - fastapi==0.104.1
     - Jinja2==3.1.2
+    - rq==1.15.1
     - sqlalchemy-stubs==0.4
     - types-redis==4.6.0.11
     - types-requests==2.31.0.10

app/requirements.txt

@@ -9,6 +9,7 @@ psycopg2-binary==2.9.9
 python-decouple==3.8
 python-multipart==0.0.6
 redis==5.0.1
+rq==1.15.1
 SQLAlchemy==2.0.25
 uvicorn==0.25.0
 validators==0.22.0

app/src/app.py

@@ -1,6 +1,5 @@
 import binascii
 import dataclasses
-import datetime
 import os
 import time
 import traceback
@@ -11,28 +10,26 @@
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from redis import Redis
+from rq import Queue
 from starlette.responses import Response
 
 from common.config import Config
 from common.language import Language
 from common.mail_receiver_utils import get_key_from_username
 
-from .app_utils import (
-    get_from_and_dkim_domain,
-    recipient_username_to_address,
-    scan_and_log,
-)
-from .check_results import load_check_results, save_check_results
+from .app_utils import recipient_username_to_address, scan_and_log
+from .check_results import load_check_results
 from .db import ScanLogEntrySource, ServerErrorLogEntry, Session
 from .logging import build_logger
 from .resolver import setup_resolver
 from .scan import DomainValidationException, ScanningException, ScanResult
 from .templates import setup_templates
-from .translate import translate
+from .worker import check_domain_scan_job
 
 app = FastAPI()
 LOGGER = build_logger(__name__)
 REDIS = Redis.from_url(Config.Data.REDIS_URL)
+job_queue = Queue(connection=REDIS)
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -118,42 +115,24 @@ async def check_email_results(request: Request, recipient_username: str) -> Resp
             },
         )
 
-    message_timestamp = datetime.datetime.fromisoformat(message_timestamp_raw.decode("ascii"))
-
     _, mail_from = parseaddr(mail_from_raw.decode("ascii"))
     _, envelope_domain = tuple(mail_from.split("@", 1))
 
-    from_domain, dkim_domain = get_from_and_dkim_domain(message_data)
-    if not from_domain:
-        result = None
-        error = translate("Invalid or no e-mail domain in the message From header", Language(Config.UI.LANGUAGE))
-    else:
-        try:
-            result = scan_and_log(
-                request=request,
-                source=ScanLogEntrySource.GUI,
-                envelope_domain=envelope_domain,
-                from_domain=from_domain,
-                dkim_domain=dkim_domain,
-                message=message_data,
-                message_timestamp=message_timestamp,
-                nameservers=Config.Network.NAMESERVERS,
-                language=Language(Config.UI.LANGUAGE),
-            )
-            error = None
-        except (DomainValidationException, ScanningException) as e:
-            result = None
-            error = translate(e.message, Language(Config.UI.LANGUAGE))
-
-    token = save_check_results(
-        envelope_domain=envelope_domain,
-        from_domain=from_domain or envelope_domain,
-        dkim_domain=dkim_domain,
-        result=result,
-        error=error,
-        rescan_url="/check-email/",
-        message_recipient_username=recipient_username,
+    client_ip = request.client.host if request.client else None
+    client_user_agent = request.headers.get("user-agent", None)
+
+    token = binascii.hexlify(os.urandom(32)).decode("ascii")
+    job_queue.enqueue(
+        check_domain_scan_job,
+        client_ip,
+        client_user_agent,
+        envelope_domain,
+        token,
+        key,
+        recipient_username,
+        job_id=token,
     )
+
     return RedirectResponse(f"/check-results/{token}", status_code=302)
 
 
@@ -169,37 +148,24 @@ async def check_domain_scan_get(request: Request) -> Response:
 
 @app.post("/check-domain/scan", response_class=HTMLResponse, include_in_schema=False)
 async def check_domain_scan_post(request: Request, domain: str = Form()) -> Response:
-    try:
-        result = scan_and_log(
-            request=request,
-            source=ScanLogEntrySource.GUI,
-            envelope_domain=domain,
-            from_domain=domain,
-            dkim_domain=None,
-            message=None,
-            message_timestamp=None,
-            nameservers=Config.Network.NAMESERVERS,
-            language=Language(Config.UI.LANGUAGE),
-        )
-        error = None
-    except (DomainValidationException, ScanningException) as e:
-        result = None
-        error = translate(e.message, Language(Config.UI.LANGUAGE))
-
-    token = save_check_results(
-        envelope_domain=domain,
-        from_domain=domain,
-        dkim_domain=None,
-        result=result,
-        error=error,
-        rescan_url="/check-domain/",
-        message_recipient_username=None,
-    )
+    client_ip = request.client.host if request.client else None
+    client_user_agent = request.headers.get("user-agent", None)
+
+    token = binascii.hexlify(os.urandom(32)).decode("ascii")
+    job_queue.enqueue(check_domain_scan_job, client_ip, client_user_agent, domain, token, job_id=token)
+
     return RedirectResponse(f"/check-results/{token}", status_code=302)
 
 
 @app.get("/check-results/{token}", response_class=HTMLResponse, include_in_schema=False)
 async def check_results(request: Request, token: str) -> Response:
+    if job := job_queue.fetch_job(token):
+        if job.get_status(refresh=False) not in ["finished", "canceled", "failed"]:
+            return templates.TemplateResponse(
+                "check_running.html",
+                {"request": request},
+            )
+
     check_results = load_check_results(token)
 
     if not check_results:
@@ -228,8 +194,10 @@ async def check_domain_api(request: Request, domain: str) -> ScanAPICallResult:
     object will be empty, as DKIM can't be checked when given only a domain.
     """
     try:
+        client_ip = request.client.host if request.client else None
+        client_user_agent = request.headers.get("user-agent", None)
+
         result = scan_and_log(
-            request=request,
             source=ScanLogEntrySource.API,
             envelope_domain=domain,
             from_domain=domain,
@@ -238,6 +206,8 @@ async def check_domain_api(request: Request, domain: str) -> ScanAPICallResult:
             message_timestamp=None,
             nameservers=Config.Network.NAMESERVERS,
             language=Language(Config.UI.LANGUAGE),
+            client_ip=client_ip,
+            client_user_agent=client_user_agent,
         )
         return ScanAPICallResult(result=result)
     except (DomainValidationException, ScanningException):

app/src/app_utils.py

@@ -9,7 +9,6 @@
 
 import dkim.util
 from email_validator import EmailNotValidError, validate_email
-from fastapi import Request
 
 from common.config import Config
 from common.language import Language
@@ -70,7 +69,6 @@ def dkim_implementation_mismatch_callback(message: bytes, dkimpy_valid: bool, op
 
 
 def scan_and_log(
-    request: Request,
     source: ScanLogEntrySource,
     envelope_domain: str,
     from_domain: str,
@@ -79,15 +77,17 @@ def scan_and_log(
     message_timestamp: Optional[datetime.datetime],
     nameservers: List[str],
     language: Language,
+    client_ip: Optional[str],
+    client_user_agent: Optional[str],
 ) -> ScanResult:
     scan_log_entry = ScanLogEntry(
         envelope_domain=envelope_domain,
         from_domain=from_domain,
         dkim_domain=dkim_domain,
         message=message,
         source=source.value,
-        client_ip=request.client.host if request.client else None,
-        client_user_agent=request.headers.get("user-agent", None),
+        client_ip=client_ip,
+        client_user_agent=client_user_agent,
         check_options={
             "nameservers": nameservers,
         },

app/src/check_results.py

@@ -1,8 +1,6 @@
-import binascii
 import dataclasses
 import datetime
 import json
-import os
 from typing import Any, Dict, Optional
 
 import dacite
@@ -35,8 +33,8 @@ def save_check_results(
     error: Optional[str],
     rescan_url: str,
     message_recipient_username: Optional[str],
-) -> str:
-    token = binascii.hexlify(os.urandom(32)).decode("ascii")
+    token: str,
+) -> None:
     # We don't use HSET or HMSET, as result is a recursive dict, and values that can be stored
     # using HSET/HMSET are bytes, string, int or float, so we still wouldn't avoid serialization.
     REDIS.set(
@@ -56,7 +54,6 @@ def save_check_results(
             cls=JSONEncoderAdditionalTypes,
         ),
     )
-    return token
 
 
 def load_check_results(token: str) -> Optional[Dict[str, Any]]:

app/src/worker.py

@@ -0,0 +1,106 @@
+import binascii
+import datetime
+import os
+from typing import Optional
+
+from redis import Redis
+
+from common.config import Config
+from common.language import Language
+
+from .app_utils import get_from_and_dkim_domain, scan_and_log
+from .check_results import save_check_results
+from .db import ScanLogEntrySource
+from .logging import build_logger
+from .resolver import setup_resolver
+from .scan import DomainValidationException, ScanningException
+from .translate import translate
+
+LOGGER = build_logger(__name__)
+REDIS = Redis.from_url(Config.Data.REDIS_URL)
+
+setup_resolver()
+
+
+def check_domain_scan_job(
+    client_ip: Optional[str],
+    client_user_agent: Optional[str],
+    envelope_domain: str,
+    token: str,
+    message_key: Optional[bytes] = None,
+    recipient_username: Optional[str] = None,
+) -> None:
+    if message_key:
+        message_data = REDIS.get(message_key)
+        message_timestamp_raw = REDIS.get(message_key + b"-timestamp")
+
+        if not message_data or not message_timestamp_raw:
+            raise RuntimeError("Worker coudn't access message data")
+
+        message_timestamp = datetime.datetime.fromisoformat(message_timestamp_raw.decode("ascii"))
+
+        from_domain, dkim_domain = get_from_and_dkim_domain(message_data)
+        if not from_domain:
+            result = None
+            error = translate("Invalid or no e-mail domain in the message From header", Language(Config.UI.LANGUAGE))
+        else:
+            try:
+                result = scan_and_log(
+                    source=ScanLogEntrySource.GUI,
+                    envelope_domain=envelope_domain,
+                    from_domain=from_domain,
+                    dkim_domain=dkim_domain,
+                    message=message_data,
+                    message_timestamp=message_timestamp,
+                    nameservers=Config.Network.NAMESERVERS,
+                    language=Language(Config.UI.LANGUAGE),
+                    client_ip=client_ip,
+                    client_user_agent=client_user_agent,
+                )
+                error = None
+            except (DomainValidationException, ScanningException) as e:
+                result = None
+                error = translate(e.message, Language(Config.UI.LANGUAGE))
+
+        token = binascii.hexlify(os.urandom(32)).decode("ascii")
+
+        save_check_results(
+            envelope_domain=envelope_domain,
+            from_domain=from_domain or envelope_domain,
+            dkim_domain=dkim_domain,
+            result=result,
+            error=error,
+            rescan_url="/check-email/",
+            message_recipient_username=recipient_username,
+            token=token,
+        )
+
+    else:
+        try:
+            result = scan_and_log(
+                source=ScanLogEntrySource.GUI,
+                envelope_domain=envelope_domain,
+                from_domain=envelope_domain,
+                dkim_domain=None,
+                message=None,
+                message_timestamp=None,
+                nameservers=Config.Network.NAMESERVERS,
+                language=Language(Config.UI.LANGUAGE),
+                client_ip=client_ip,
+                client_user_agent=client_user_agent,
+            )
+            error = None
+        except (DomainValidationException, ScanningException) as e:
+            result = None
+            error = translate(e.message, Language(Config.UI.LANGUAGE))
+
+        save_check_results(
+            envelope_domain=envelope_domain,
+            from_domain=envelope_domain,
+            dkim_domain=None,
+            result=result,
+            error=error,
+            rescan_url="/check-domain/",
+            message_recipient_username=None,
+            token=token,
+        )

app/templates/check_running.html

@@ -0,0 +1,19 @@
+{% extends "custom_layout.html" %}
+
+{% block header_additional %}
+    <meta http-equiv="refresh" content="5">
+{% endblock %}
+
+{% block body %}
+    <div class="container">
+        <div class="card">
+            <div class="card-body">
+                <h5 class="card-title">{% trans %}Domain analysis is running{% endtrans %}</h5>
+                <div class="card-text">
+                    <img class="spinner waiting" src="/static/images/spinner.svg" />
+                    <h5 class="waiting">{% trans %}Waiting for the domain analysis to finish{% endtrans %}</h5>
+                </div>
+            </div>
+        </div>
+   </div>
+{% endblock %}

app/translations/en_US/LC_MESSAGES/messages.po

@@ -240,6 +240,14 @@ msgid ""
 "detected only if earlier checks complete successfully."
 msgstr ""
 
+#: app/templates/check_running.html:11
+msgid "Domain analysis is running"
+msgstr ""
+
+#: app/templates/check_running.html:14
+msgid "Waiting for the domain analysis to finish"
+msgstr ""
+
 #: app/templates/root.html:9
 msgid ""
 "Verify your <b>DKIM</b>, <b>DMARC</b>, and <b>SPF</b> settings by sending"
@@ -262,3 +270,6 @@ msgid ""
 "<b>SPF</b> and <b>DMARC</b> mechanisms will be checked. To check DKIM, "
 "you need to send a test e-mail."
 msgstr ""
+
+#~ msgid "Waiting for the domain analysis to finfish"
+#~ msgstr ""