46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
from io import BytesIO
|
|
|
|
import gnupg
|
|
from memory_profiler import memory_usage
|
|
|
|
from app.config import GNUPGHOME
|
|
from app.log import LOG
|
|
from app.utils import random_string
|
|
|
|
gpg = gnupg.GPG(gnupghome=GNUPGHOME)
|
|
gpg.encoding = "utf-8"
|
|
|
|
|
|
class PGPException(Exception):
|
|
pass
|
|
|
|
|
|
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
|
|
|
|
|
|
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)
|
|
r = gpg.encrypt_file(data, fingerprint, always_trust=True)
|
|
if not r.ok:
|
|
# save the content for debugging
|
|
random_file_name = random_string(20) + ".eml"
|
|
full_path = f"/tmp/{random_file_name}"
|
|
with open(full_path, "wb") as f:
|
|
f.write(data.getbuffer())
|
|
LOG.error("PGP fail - log to %s", full_path)
|
|
raise PGPException("Cannot encrypt")
|
|
|
|
return str(r)
|