diff --git a/app/pgp_utils.py b/app/pgp_utils.py index 638a0bfc..45497988 100644 --- a/app/pgp_utils.py +++ b/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) diff --git a/email_handler.py b/email_handler.py index 17bd1056..f9868635 100644 --- a/email_handler.py +++ b/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) diff --git a/init_app.py b/init_app.py index 54bb1239..22389fc2 100644 --- a/init_app.py +++ b/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() diff --git a/requirements.in b/requirements.in index cefa88e5..76df3e9b 100644 --- a/requirements.in +++ b/requirements.in @@ -41,3 +41,4 @@ python-gnupg webauthn pyspf Flask-Limiter +memory_profiler \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index b1e2c52c..41832294 100644 --- a/requirements.txt +++ b/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