Merge pull request #223 from simple-login/pgp-fix

Fix intermitten PGP errors

app/pgp_utils.py

@@ -1,6 +1,7 @@
 from io import BytesIO
 
 import gnupg
+from memory_profiler import memory_usage
 
 from app.config import GNUPGHOME
 from app.log import LOG
@@ -24,6 +25,10 @@ def load_public_key(public_key: str) -> str:
 
 
 def encrypt_file(data: BytesIO, fingerprint: str) -> str:
+    LOG.d("encrypt for %s", fingerprint)
+    mem_usage = memory_usage(-1, interval=1, timeout=1)
+    LOG.d("mem_usage %s", mem_usage)
+
     r = gpg.encrypt_file(data, fingerprint, always_trust=True)
     if not r.ok:
         LOG.error("Try encrypt again %s", fingerprint)
@@ -34,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)

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)
 
@@ -1063,7 +1076,7 @@ if __name__ == "__main__":
         LOG.warning("LOAD PGP keys")
         app = create_app()
         with app.app_context():
-            load_pgp_public_keys(app)
+            load_pgp_public_keys()
 
     while True:
         time.sleep(2)

init_app.py

@@ -6,21 +6,18 @@ from app.pgp_utils import load_public_key
 from server import create_app
 
 
-def load_pgp_public_keys(app):
+def load_pgp_public_keys():
     """Load PGP public key to keyring"""
-    with app.app_context():
-        for mailbox in Mailbox.query.filter(Mailbox.pgp_public_key != None).all():
-            LOG.d("Load PGP key for mailbox %s", mailbox)
-            fingerprint = load_public_key(mailbox.pgp_public_key)
+    for mailbox in Mailbox.query.filter(Mailbox.pgp_public_key != None).all():
+        LOG.d("Load PGP key for mailbox %s", mailbox)
+        fingerprint = load_public_key(mailbox.pgp_public_key)
 
-            # sanity check
-            if fingerprint != mailbox.pgp_finger_print:
-                LOG.error(
-                    "fingerprint %s different for mailbox %s", fingerprint, mailbox
-                )
-                mailbox.pgp_finger_print = fingerprint
+        # sanity check
+        if fingerprint != mailbox.pgp_finger_print:
+            LOG.error("fingerprint %s different for mailbox %s", fingerprint, mailbox)
+            mailbox.pgp_finger_print = fingerprint
 
-        db.session.commit()
+    db.session.commit()
 
     LOG.d("Finish load_pgp_public_keys")
 
@@ -29,4 +26,4 @@ if __name__ == "__main__":
     app = create_app()
 
     with app.app_context():
-        load_pgp_public_keys(app)
+        load_pgp_public_keys()

requirements.in

@@ -41,3 +41,4 @@ python-gnupg
 webauthn
 pyspf
 Flask-Limiter
+memory_profiler

requirements.txt

@@ -8,6 +8,7 @@ aiohttp==3.5.4            # via raven-aiohttp, yacron
 aiosmtpd==1.2             # via -r requirements.in
 aiosmtplib==1.0.6         # via yacron
 alembic==1.0.10           # via flask-migrate
+appnope==0.1.0            # via ipython
 arrow==0.14.2             # via -r requirements.in
 asn1crypto==0.24.0        # via cryptography
 async-timeout==3.0.1      # via aiohttp
@@ -63,6 +64,7 @@ jwcrypto==0.6.0           # via -r requirements.in
 limits==1.5.1             # via flask-limiter
 mako==1.0.12              # via alembic
 markupsafe==1.1.1         # via jinja2, mako
+memory-profiler==0.57.0   # via -r requirements.in
 more-itertools==7.0.0     # via pytest
 multidict==4.5.2          # via aiohttp, yarl
 oauthlib==3.0.2           # via requests-oauthlib
@@ -74,6 +76,7 @@ pickleshare==0.7.5        # via ipython
 pip-tools==3.8.0          # via -r requirements.in
 pluggy==0.12.0            # via pytest
 prompt-toolkit==2.0.9     # via ipython
+psutil==5.7.0             # via memory-profiler
 psycopg2-binary==2.8.2    # via -r requirements.in
 ptyprocess==0.6.0         # via pexpect
 py==1.8.0                 # via pytest