paulbsd-salt/states/maildb/templates/maildb.py

#!/usr/bin/env python3

import argparse
import sys
import configparser
from email.utils import parseaddr
import email
import psycopg2


def DBConfig(configfile):
    config = configparser.ConfigParser()
    config.read(configfile)

    d_dbconfig = {
        "host"   : config.get("maildb", "hostname"),
        "user"   : config.get("maildb", "username"),
        "passwd" : config.get("maildb", "password"),
        "db"     : config.get("maildb", "database")
    }

    return d_dbconfig


def ConnectDB(dbconfig):
    connstr = f"host={dbconfig['host']} dbname={dbconfig['db']} user={dbconfig['user']} password={dbconfig['passwd']}"
    conn = psycopg2.connect(connstr)

    return conn, conn.cursor()


def CloseDB(conn):
    conn.close()


def CreateDB(conn, cursor):
    try:
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS mail (
            id serial NOT NULL,
            sender text NULL,
            recipient text NULL,
            "date" timestamp NOT NULL DEFAULT NOW(),
            "content" text NULL,
            CONSTRAINT mail_pkey PRIMARY KEY (id)
        )
        ''')
        conn.commit()
    except Exception as e:
        print(e)


def InsertMail(conn, cursor, sender, recipient, content):
    try:
        cursor.execute("PREPARE InsertMail as INSERT INTO mail (sender, recipient, content) values ($1, $2, $3);")
        cursor.execute("EXECUTE InsertMail (%s, %s, %s);", (sender, recipient, content))
        conn.commit()
    except Exception as e:
        print(e)


def ParseMail(input_data):
    content = ""
    for line in input_data:
        if line != "":
            content += line

    msgobj = email.message_from_string(content)
    sender = parseaddr(msgobj["From"])[1]
    recipient = parseaddr(msgobj["To"])[1]

    return sender, recipient, content


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", help="config file path", default="maildb.cfg")
    parser.add_argument("sender", help="sender")
    parser.add_argument("recipient", help="recipient")
    args = parser.parse_args()

    d_dbconfig = DBConfig(args.config)

    maildata = ParseMail(sys.stdin)
    sender = maildata[0]
    recipient = maildata[1]
    mailcontent = maildata[2]

    conn, cursor = ConnectDB(d_dbconfig)
    CreateDB(conn, cursor)

    InsertMail(conn, cursor, sender, recipient, mailcontent)

    CloseDB(conn)


if __name__ == "__main__":
    main()