Parcourir la source

stop the email handler process when PGP error

Son NK il y a 5 ans
Parent
commit
16df2acb29
2 fichiers modifiés avec 22 ajouts et 9 suppressions
  1. 1 1
      app/pgp_utils.py
  2. 21 8
      email_handler.py

+ 1 - 1
app/pgp_utils.py

@@ -39,7 +39,7 @@ def encrypt_file(data: BytesIO, fingerprint: str) -> str:
             full_path = f"/tmp/{random_file_name}"
             with open(full_path, "wb") as f:
                 f.write(data.getbuffer())
-            LOG.error("Log to %s", full_path)
+            LOG.error("PGP fail - log to %s", full_path)
             raise PGPException("Cannot encrypt")
 
     return str(r)

+ 21 - 8
email_handler.py

@@ -30,13 +30,10 @@ It should contain the following info:
 
 
 """
-import arrow
 import email
-import spf
+import os
 import time
 import uuid
-from aiosmtpd.controller import Controller
-from aiosmtpd.smtp import Envelope
 from email import encoders
 from email.message import Message
 from email.mime.application import MIMEApplication
@@ -46,6 +43,11 @@ from io import BytesIO
 from smtplib import SMTP
 from typing import List, Tuple
 
+import arrow
+import spf
+from aiosmtpd.controller import Controller
+from aiosmtpd.smtp import Envelope
+
 from app import pgp_utils, s3
 from app.alias_utils import try_auto_create
 from app.config import (
@@ -90,6 +92,7 @@ from app.models import (
     RefusedEmail,
     Mailbox,
 )
+from app.pgp_utils import PGPException
 from app.utils import random_string
 from init_app import load_pgp_public_keys
 from server import create_app
@@ -341,10 +344,14 @@ def prepare_pgp_message(orig_msg: Message, pgp_fingerprint: str):
 
     second = MIMEApplication("octet-stream", _encoder=encoders.encode_7or8bit)
     second.add_header("Content-Disposition", "inline")
-    # encrypt original message
-    encrypted_data = pgp_utils.encrypt_file(
-        BytesIO(orig_msg.as_bytes()), pgp_fingerprint
-    )
+    try:
+        # encrypt original message
+        encrypted_data = pgp_utils.encrypt_file(
+            BytesIO(orig_msg.as_bytes()), pgp_fingerprint
+        )
+    except PGPException:
+        LOG.error("Exit due to PGP fail")
+        exit()
     second.set_payload(encrypted_data)
     msg.attach(second)
 
@@ -1053,6 +1060,12 @@ class MailHandler:
             return handle(envelope, smtp)
 
 
+def exit():
+    pid = os.getpid()
+    LOG.warning("kill pid %s", pid)
+    os.kill(pid, 9)
+
+
 if __name__ == "__main__":
     controller = Controller(MailHandler(), hostname="0.0.0.0", port=20381)