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 CURRENT_DATE,
	"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(f"INSERT INTO mail (sender, recipient, content) values ('{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()