paulbsd-salt/states/maildb/maildb.py

98 lines
2.3 KiB
Python

#!/usr/bin/env python3
import argparse
import sys
import configparser
from email.utils import parseaddr
import email
#import MySQLdb
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 ConnectMysqlDB(dbconfig):
connstr = f"host={dbconfig['host']} dbname={dbconfig['db']} user={dbconfig['user']} password={dbconfig['passwd']}"
conn = psycopg2.connect(connstr)
#conn = MySQLdb.connect(**dbconfig)
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 = ConnectMysqlDB(d_dbconfig)
CreateDB(conn, cursor)
InsertMail(conn, cursor, sender, recipient, mailcontent)
CloseDB(conn)
if __name__ == '__main__':
main()