2020-06-08 11:53:27 +00:00
|
|
|
import os
|
2020-04-02 19:29:16 +00:00
|
|
|
from io import BytesIO
|
|
|
|
|
2020-03-08 11:51:33 +00:00
|
|
|
import gnupg
|
2020-06-07 07:47:35 +00:00
|
|
|
from memory_profiler import memory_usage
|
2020-03-08 11:51:33 +00:00
|
|
|
|
|
|
|
from app.config import GNUPGHOME
|
2020-06-05 08:10:21 +00:00
|
|
|
from app.log import LOG
|
2020-08-08 08:26:24 +00:00
|
|
|
from app.models import Mailbox, Contact
|
2020-06-05 08:10:21 +00:00
|
|
|
from app.utils import random_string
|
2020-03-08 11:51:33 +00:00
|
|
|
|
|
|
|
gpg = gnupg.GPG(gnupghome=GNUPGHOME)
|
2020-03-13 11:54:52 +00:00
|
|
|
gpg.encoding = "utf-8"
|
2020-03-08 11:51:33 +00:00
|
|
|
|
2020-03-08 22:07:23 +00:00
|
|
|
|
2020-03-08 11:51:33 +00:00
|
|
|
class PGPException(Exception):
|
|
|
|
pass
|
|
|
|
|
2020-03-08 22:07:23 +00:00
|
|
|
|
2020-03-08 11:51:33 +00:00
|
|
|
def load_public_key(public_key: str) -> str:
|
|
|
|
"""Load a public key into keyring and return the fingerprint. If error, raise Exception"""
|
|
|
|
import_result = gpg.import_keys(public_key)
|
|
|
|
try:
|
|
|
|
return import_result.fingerprints[0]
|
|
|
|
except Exception as e:
|
|
|
|
raise PGPException("Cannot load key") from e
|
|
|
|
|
|
|
|
|
2020-06-08 11:53:27 +00:00
|
|
|
def hard_exit():
|
|
|
|
pid = os.getpid()
|
|
|
|
LOG.warning("kill pid %s", pid)
|
|
|
|
os.kill(pid, 9)
|
|
|
|
|
|
|
|
|
2020-04-02 19:29:16 +00:00
|
|
|
def encrypt_file(data: BytesIO, fingerprint: str) -> str:
|
2020-06-07 07:47:35 +00:00
|
|
|
LOG.d("encrypt for %s", fingerprint)
|
2020-06-08 11:53:27 +00:00
|
|
|
mem_usage = memory_usage(-1, interval=1, timeout=1)[0]
|
2020-06-07 07:47:35 +00:00
|
|
|
LOG.d("mem_usage %s", mem_usage)
|
|
|
|
|
2020-04-02 19:29:16 +00:00
|
|
|
r = gpg.encrypt_file(data, fingerprint, always_trust=True)
|
2020-03-08 11:51:33 +00:00
|
|
|
if not r.ok:
|
2020-06-10 20:28:15 +00:00
|
|
|
# maybe the fingerprint is not loaded on this host, try to load it
|
2020-08-08 08:26:24 +00:00
|
|
|
found = False
|
|
|
|
# searching for the key in mailbox
|
2020-06-10 20:28:15 +00:00
|
|
|
mailbox = Mailbox.get_by(pgp_finger_print=fingerprint)
|
|
|
|
if mailbox:
|
|
|
|
LOG.d("(re-)load public key for %s", mailbox)
|
|
|
|
load_public_key(mailbox.pgp_public_key)
|
2020-08-08 08:26:24 +00:00
|
|
|
found = True
|
2020-06-10 20:28:15 +00:00
|
|
|
|
2020-08-08 08:26:24 +00:00
|
|
|
# searching for the key in contact
|
|
|
|
contact = Contact.get_by(pgp_finger_print=fingerprint)
|
|
|
|
if contact:
|
|
|
|
LOG.d("(re-)load public key for %s", contact)
|
|
|
|
load_public_key(contact.pgp_public_key)
|
|
|
|
found = True
|
|
|
|
|
|
|
|
if found:
|
2020-06-10 20:28:15 +00:00
|
|
|
LOG.d("retry to encrypt")
|
2020-06-11 22:01:21 +00:00
|
|
|
data.seek(0)
|
2020-06-10 20:28:15 +00:00
|
|
|
r = gpg.encrypt_file(data, fingerprint, always_trust=True)
|
|
|
|
|
|
|
|
if not r.ok:
|
|
|
|
raise PGPException(f"Cannot encrypt, status: {r.status}")
|
2020-03-08 11:51:33 +00:00
|
|
|
|
|
|
|
return str(r)
|