Initial version with debug prints

6 files changed, 257 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..382ddaa
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+config.py
diff --git a/config.py.dist b/config.py.dist
new file mode 100644
index 0000000..e35e8ae
--- /dev/null
+++ b/config.py.dist
@@ -0,0 +1,16 @@
+socket = '/var/spool/postfix/pgp-milter/pgp-milter.sock'
+db_config = {
+  'host': 'localhost',
+  'user': 'user',
+  'password': 'password',
+  'database': 'database'
+}
+
+# Things to bear in mind when writing this query:
+# - `fingerprint` can be any deduplicating attribute of any type
+# - `raw_data` can either be armored or unarmored key data
+# - column names don't matter
+# - the input argument (%s) is a series of addresses, like so:
+#   - `'a@example.com', 'b@example.com', 'c@example.com'`
+# - a prepared statement is used so input is sanitized for you
+key_loader_query = 'SELECT fingerprint, raw_data FROM pgp_keys WHERE addr IN (%s)'
diff --git a/key_loader.py b/key_loader.py
new file mode 100644
index 0000000..4867d07
--- /dev/null
+++ b/key_loader.py
@@ -0,0 +1,38 @@
+import config
+import mysql.connector as mysql
+import pgpy
+
+def load_keys(addrs: list[str]) -> list[pgpy.PGPKey]:
+  if len(addrs) < 1:
+    return []
+
+  db = mysql.connect(**config.db_config)
+  try:
+    stmt = db.cursor()
+  except Exception as e:
+    db.close()
+    raise e
+  
+  try:
+    stmt.execute(\
+      format_multiquery(config.key_loader_query, len(addrs)), \
+      tuple(addrs) \
+    )
+    
+    keys = {}
+    for (keyfpr, keydata) in stmt:
+      if keyfpr in keys:
+        continue
+      key = pgpy.PGPKey()
+      try:
+        key.parse(keydata)
+      except:
+        continue
+
+    return list(keys.values())
+  finally:
+    stmt.close()
+    db.close()
+
+def format_multiquery(query: str, n_items: int) -> str:
+  return query % (', '.join(['%s'] * n_items),)
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..4d2e7c2
--- /dev/null
+++ b/main.py
@@ -0,0 +1,85 @@
+import config
+import utils
+import Milter
+import email
+
+class PGPMilter(Milter.Base):
+  def __init__(self):
+    self.recipients = []
+    self.headers = []
+    self.content = bytes()
+
+  @Milter.noreply
+  def connect(_self, _ip_name, _family, _hostaddr):
+    return Milter.CONTINUE
+
+  @Milter.noreply
+  def envfrom(self, name, *esmtp_params):
+    self.__init__()
+    return Milter.CONTINUE
+
+  @Milter.noreply
+  def envrcpt(self, name, *strings):
+    self.recipients.append(name)
+    return Milter.CONTINUE
+
+  @Milter.noreply
+  def header(self, k: str, v: str):
+    self.headers.append((k.encode(), v.encode()))
+    return Milter.CONTINUE
+
+  def eoh(self):
+    return Milter.CONTINUE
+
+  def body(self, chunk):
+    self.content += chunk
+    return Milter.CONTINUE
+
+  def eom(self):
+    raw_headers = b'\n'.join(map(lambda header : b'%s: %s' % header, self.headers))
+    msg = email.message_from_bytes(raw_headers + b'\n\n' + self.content,\
+                                   policy=email.policy.default)
+
+    if b'-----BEGIN PGP MESSAGE-----' in self.content or utils.already_encrypted(msg):
+      return Milter.ACCEPT
+
+    # TODO: remove debug print
+    print('Encrypting message to recipients: [%s]' % ', '.join(self.recipients))
+    enc_msg, encrypted = utils.encrypt(msg, self.recipients)
+    if not encrypted:
+      # TODO: remove debug print
+      print('No keys found, passing through.')
+      return Milter.ACCEPT
+
+    # `Content-Transfer-Encoding: quoted-printable`
+    # can prevent the message from being decrypted by clients 
+    self.set_header(msg, 'Content-Transfer-Encoding', '')
+    for (k, v) in enc_msg.items():
+      self.set_header(msg, k, v)
+
+    enc_bytes = enc_msg.as_bytes()
+    enc_body = enc_bytes[enc_bytes.find(b'\n\n')+2:]
+    self.replacebody(enc_body)
+
+    return Milter.ACCEPT
+
+  def close(self):
+    self.__init__()
+    return Milter.CONTINUE
+
+  def set_header(self, old_msg, k, v):
+    old_headers = old_msg.get_all(k)
+    if old_headers != None:
+      for i in range(len(old_headers)-1, -1, -1):
+        self.chgheader(k, i, '')
+    if v != None and len(v) > 0:
+      self.addheader(k, v)
+
+def main():
+  Milter.factory = PGPMilter
+  Milter.set_flags(Milter.ADDHDRS + Milter.CHGHDRS + Milter.CHGBODY)
+  Milter.runmilter('cmail-pgp-milter', config.socket)
+
+if __name__ == '__main__':
+  print('Starting CMail PGP Milter')
+  main()
diff --git a/pgp-milter.service b/pgp-milter.service
new file mode 100644
index 0000000..2b93f40
--- /dev/null
+++ b/pgp-milter.service
@@ -0,0 +1,16 @@
+[Unit]
+Description=CMail PGP Milter Service
+Documentation=https://github.com/cflems/pgp-milter
+After=network.target
+
+[Service]
+Type=exec
+User=pgp-milter
+Group=pgp-milter
+StandardOutput=append:/var/log/pgp-milter.log
+StandardError=append:/var/log/pgp-milter.err
+ExecStart=/usr/bin/python3 /opt/pgp-milter/main.py
+Restart=on-failure
+
+[Install]
+WantedBy=multi-user.target
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)