Initial version with debug prints

1 files changed, 85 insertions, 0 deletions

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()