From 8d785882a6579a70fef4581139b25abe9836c36f Mon Sep 17 00:00:00 2001 From: mimi89999 Date: Thu, 18 Jan 2024 18:24:29 +0100 Subject: [PATCH] Add job enqueueing for domain scans --- app/requirements.txt | 1 + app/src/app.py | 45 ++++++++++++++++--- app/src/app_utils.py | 7 +-- app/src/check_results.py | 5 +-- app/templates/check_running.html | 19 ++++++++ .../en_US/LC_MESSAGES/messages.po | 8 ++++ app/translations/messages.pot | 8 ++++ .../pl_PL/LC_MESSAGES/messages.po | 8 ++++ docker-compose.yml | 11 +++++ test/base.py | 9 +++- 10 files changed, 106 insertions(+), 15 deletions(-) create mode 100644 app/templates/check_running.html diff --git a/app/requirements.txt b/app/requirements.txt index 7cbc5a0..59e9f90 100644 --- a/app/requirements.txt +++ b/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 diff --git a/app/src/app.py b/app/src/app.py index 92a6663..826567d 100644 --- a/app/src/app.py +++ b/app/src/app.py @@ -11,6 +11,7 @@ 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 @@ -33,6 +34,7 @@ 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") @@ -129,8 +131,10 @@ async def check_email_results(request: Request, recipient_username: str) -> Resp error = translate("Invalid or no e-mail domain in the message From header", Language(Config.UI.LANGUAGE)) else: 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.GUI, envelope_domain=envelope_domain, from_domain=from_domain, @@ -139,13 +143,17 @@ async def check_email_results(request: Request, recipient_username: str) -> Resp 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 = save_check_results( + 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, @@ -153,6 +161,7 @@ async def check_email_results(request: Request, recipient_username: str) -> Resp error=error, rescan_url="/check-email/", message_recipient_username=recipient_username, + token=token, ) return RedirectResponse(f"/check-results/{token}", status_code=302) @@ -167,11 +176,9 @@ async def check_domain_scan_get(request: Request) -> Response: return RedirectResponse("/check-domain") -@app.post("/check-domain/scan", response_class=HTMLResponse, include_in_schema=False) -async def check_domain_scan_post(request: Request, domain: str = Form()) -> Response: +def check_domain_scan_job(client_ip: Optional[str], client_user_agent: Optional[str], domain: str, token: str): try: result = scan_and_log( - request=request, source=ScanLogEntrySource.GUI, envelope_domain=domain, from_domain=domain, @@ -180,13 +187,15 @@ async def check_domain_scan_post(request: Request, domain: str = Form()) -> Resp 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)) - token = save_check_results( + save_check_results( envelope_domain=domain, from_domain=domain, dkim_domain=None, @@ -194,12 +203,30 @@ async def check_domain_scan_post(request: Request, domain: str = Form()) -> Resp error=error, rescan_url="/check-domain/", message_recipient_username=None, + token=token, ) + + +@app.post("/check-domain/scan", response_class=HTMLResponse, include_in_schema=False) +async def check_domain_scan_post(request: Request, domain: str = Form()) -> Response: + 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 +255,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 +267,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): diff --git a/app/src/app_utils.py b/app/src/app_utils.py index b15eae8..b7d18ff 100644 --- a/app/src/app_utils.py +++ b/app/src/app_utils.py @@ -70,7 +70,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,6 +78,8 @@ 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, @@ -86,8 +87,8 @@ def scan_and_log( 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, }, diff --git a/app/src/check_results.py b/app/src/check_results.py index 3cfceea..2b65728 100644 --- a/app/src/check_results.py +++ b/app/src/check_results.py @@ -35,8 +35,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 +56,6 @@ def save_check_results( cls=JSONEncoderAdditionalTypes, ), ) - return token def load_check_results(token: str) -> Optional[Dict[str, Any]]: diff --git a/app/templates/check_running.html b/app/templates/check_running.html new file mode 100644 index 0000000..ed51809 --- /dev/null +++ b/app/templates/check_running.html @@ -0,0 +1,19 @@ +{% extends "custom_layout.html" %} + +{% block header_additional %} + +{% endblock %} + +{% block body %} +
+
+
+
{% trans %}Domain analysis is running{% endtrans %}
+
+ +
{% trans %}Waiting for the domain analysis to finfish{% endtrans %}
+
+
+
+
+{% endblock %} diff --git a/app/translations/en_US/LC_MESSAGES/messages.po b/app/translations/en_US/LC_MESSAGES/messages.po index 10cb03d..f5e09eb 100644 --- a/app/translations/en_US/LC_MESSAGES/messages.po +++ b/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:7 +msgid "Domain analysis is running" +msgstr "" + +#: app/templates/check_running.html:10 +msgid "Waiting for the domain analysis to finfish" +msgstr "" + #: app/templates/root.html:9 msgid "" "Verify your DKIM, DMARC, and SPF settings by sending" diff --git a/app/translations/messages.pot b/app/translations/messages.pot index 10cb03d..f5e09eb 100644 --- a/app/translations/messages.pot +++ b/app/translations/messages.pot @@ -240,6 +240,14 @@ msgid "" "detected only if earlier checks complete successfully." msgstr "" +#: app/templates/check_running.html:7 +msgid "Domain analysis is running" +msgstr "" + +#: app/templates/check_running.html:10 +msgid "Waiting for the domain analysis to finfish" +msgstr "" + #: app/templates/root.html:9 msgid "" "Verify your DKIM, DMARC, and SPF settings by sending" diff --git a/app/translations/pl_PL/LC_MESSAGES/messages.po b/app/translations/pl_PL/LC_MESSAGES/messages.po index dce71a2..4e27586 100644 --- a/app/translations/pl_PL/LC_MESSAGES/messages.po +++ b/app/translations/pl_PL/LC_MESSAGES/messages.po @@ -267,6 +267,14 @@ msgstr "" "Po poprawie błędów prosimy ponowić skanowanie - niektóre błędy mogą " "zostać znalezione dopiero po udanym wykonaniu wcześniejszych testów." +#: app/templates/check_running.html:7 +msgid "Domain analysis is running" +msgstr "Trwa analiza domeny" + +#: app/templates/check_running.html:10 +msgid "Waiting for the domain analysis to finfish" +msgstr "Oczekiwanie na zakończenie analizy" + #: app/templates/root.html:9 msgid "" "Verify your DKIM, DMARC, and SPF settings by sending" diff --git a/docker-compose.yml b/docker-compose.yml index e66349b..8d85134 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -22,6 +22,17 @@ services: ports: - 8000:8000 <<: *common-configuration + worker: + build: + context: . + dockerfile: app/docker/Dockerfile + environment: + DB_URL: postgresql+psycopg2://postgres:postgres@db:5432/mailgoose + REDIS_URL: redis://redis:6379/0 + env_file: + - .env + command: bash -c "/wait-for-it.sh db:5432 -- rq worker" + <<: *common-configuration mail_receiver: build: context: . diff --git a/test/base.py b/test/base.py index 31b9c3d..67b93f8 100644 --- a/test/base.py +++ b/test/base.py @@ -1,3 +1,4 @@ +from time import sleep from typing import Any, Dict from unittest import TestCase @@ -7,8 +8,12 @@ class BaseTestCase(TestCase): def check_domain(self, domain: str) -> str: - response = requests.post(APP_URL + "/check-domain/scan", {"domain": domain}) - return response.text.replace("\n", " ") + submission_response = requests.post(APP_URL + "/check-domain/scan", {"domain": domain}, allow_redirects=False) + submission_response.raise_for_status() + results_url = submission_response.next.url + sleep(10) + results_response = requests.get(results_url) + return results_response.text.replace("\n", " ") def check_domain_api_v1(self, domain: str) -> Dict[str, Any]: response = requests.post(APP_URL + "/api/v1/check-domain?domain=" + domain)