import Milter
import sys
import traceback
import email
import json

class DsnProcessorMilter(Milter.Base):
  def __init__(self):
    self.mail_data: str = ""

  def process_dsn(self, mail_data: str):
    msg = email.message_from_string(mail_data)
    recipients = []
    for part in msg.walk():
      if part.get_content_type() == "message/delivery-status":
        for subpart in part.walk():
          rcpt = {}
          if 'Action' in subpart and 'Original-Recipient' in subpart:
            rcpt["action"] = subpart['Action']
            rcpt["orig_rcpt"] = subpart['Original-Recipient'].replace('rfc822;', '')
            rcpt["diag_code"] = None
            if 'Diagnostic-Code' in subpart:
              rcpt["diag_code"] = subpart['Diagnostic-Code']
            rcpt["status"] = None
            if 'Status' in subpart:
              rcpt["status"] = subpart['Status']
            recipients.append(rcpt)
    return recipients
  
  # Not registered/used callbacks
  @Milter.nocallback
  def hello(self, heloname):
    return Milter.CONTINUE
  @Milter.nocallback
  def data(self):
    return Milter.CONTINUE
  
  def connect(self, IPname, family, hostaddr):
    print("CONN: Client: {}".format(hostaddr[0])) 
    return Milter.CONTINUE
  
  # Mandatory callback
  def envfrom(self, mailfrom, *str):
    print("MAIL FROM: {}".format(mailfrom))
    self.mail_data = ""
    return Milter.CONTINUE

  # Mandatory callback
  def envrcpt(self, to, *str):
    print("RCPT TO: {}".format(to))
    return Milter.CONTINUE
  
  def header(self, name, hval):
    print("HDR: {0}: {1}".format(name, hval))
    self.mail_data += "{0}: {1}\n".format(name, hval)
    return Milter.CONTINUE
  
  def eoh(self):
    self.mail_data += "\n"
    return Milter.CONTINUE

  def body(self, chunk):
    self.mail_data += chunk.decode()
    return Milter.CONTINUE

  def eom(self):
    dsn_details = self.process_dsn(self.mail_data)
    for d in dsn_details:
      print(json.dumps(d))
    return Milter.CONTINUE
  
  def close(self):
    print("CLOSE")
    self.mail_data = ""
    return Milter.CONTINUE

if __name__ == "__main__":
  try:
    timeout = 600
    Milter.factory = DsnProcessorMilter
    print("Starting dsn-processor-milter...")
    Milter.runmilter("dsn-processor-milter", "inet:4321@127.0.0.1" , timeout, True)
  except:
    print("MAIN-EXCEPTION: {}".format(traceback.format_exc()))
    sys.exit(1)