1 files changed, 100 insertions, 0 deletions
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..f4c4f9a
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,100 @@
+import key_loader
+import pgpy
+import email
+from email.message import EmailMessage
+from email.mime.multipart import MIMEMultipart
+from email.mime.application import MIMEApplication
+from email.mime.text import MIMEText
+from copy import deepcopy
+from random import choices as alphabet_random
+from string import ascii_letters, digits
+
+protected_headers = ['to', 'cc', 'from', 'reply-to', 'followup-to', 'subject', 'date',\
+                     'message-id']
+structural_headers = ['content-type']
+overzealous_headers = ['mime-version', 'content-transfer-encoding']
+
+def encrypt(msg: EmailMessage, recipients: list[str]) -> tuple[EmailMessage, bool]:
+  payload = wrap_body(deepcopy(msg))
+
+  rcpt_keys = load_keys(recipients)
+  if len(rcpt_keys) < 1:
+    return msg, False
+
+  enc_msg = pgpy.PGPMessage.new(payload.as_string())
+  for key in rcpt_keys:
+    enc_msg = key.encrypt(enc_msg)
+
+  container = MIMEMultipart(
+    'encrypted',
+    boundary=gen_boundary(),
+    protocol='application/pgp-encrypted'
+  )
+  container.preamble = 'This is an OpenPGP/MIME encrypted message (RFC 4880 and 3156)'
+
+  part1 = MIMEApplication(
+    _data='Version: 1\n',
+    _subtype='pgp-encrypted',
+    _encoder=email.encoders.encode_7or8bit
+  )
+  part1['Content-Description'] = 'PGP/MIME version identification'
+
+  part2 = MIMEApplication(
+    _data=str(enc_msg),
+    _subtype='octet-stream; name="encrypted.asc"',
+    _encoder=email.encoders.encode_7or8bit
+  )
+  part2['Content-Description'] = 'OpenPGP encrypted message'
+  part2['Content-Disposition'] = 'inline; filename="encrypted.asc"'
+
+  strip_extraneous_headers(part1)
+  strip_extraneous_headers(part2)
+  container.attach(part1)
+  container.attach(part2)
+  strip_extraneous_headers(container, ['mime-version'])
+  return container, True
+
+def already_encrypted(msg: EmailMessage) -> bool:
+  if msg.get_content_type() in ['multipart/encrypted', 'application/pgp-encrypted']:
+    return True
+  for part in msg.iter_parts():
+    if already_encrypted(part):
+      return True
+  return False
+
+def wrap_body(msg: EmailMessage) -> EmailMessage:
+  wrapped_msg = MIMEMultipart('mixed', boundary=gen_boundary(), protected_headers='v1')
+  strip_extraneous_headers(wrapped_msg, overzealous_headers + protected_headers)
+  copy_headers(msg, wrapped_msg, protected_headers)
+
+  if msg.is_multipart():
+    strip_extraneous_headers(wrapped_msg, structural_headers)
+    copy_headers(msg, wrapped_msg, structural_headers)
+    wrapped_msg.set_payload(msg.get_payload(decode=False))
+  else:
+    text_holder = MIMEText(msg.get_payload(decode=True), _charset='utf-8')
+    strip_extraneous_headers(text_holder, structural_headers)
+    copy_headers(msg, text_holder, structural_headers)
+    wrapped_msg.attach(text_holder)
+
+  return wrapped_msg
+
+def copy_headers(orig_msg: EmailMessage, container: EmailMessage, headers: list) -> None:
+  for (header, value) in orig_msg.items():
+    if header.lower() in headers:
+      container.add_header(header, value)
+
+def strip_extraneous_headers(msg: EmailMessage, strip_headers=overzealous_headers) -> None:
+  for header in msg.keys():
+    if header.lower() in strip_headers:
+      del msg[header]
+
+def gen_boundary() -> str:
+  return '-' * 12 + ''.join(alphabet_random(ascii_letters + digits, k=24))
+
+def load_keys(recipients: list[str]) -> list[pgpy.PGPKey]:
+  addrs = []
+  for recipient in recipients:
+    _display_name, addr = email.utils.parseaddr(recipient)
+    addrs.append(addr)
+  return key_loader.load_keys(addrs)