Skip to content
Snippets Groups Projects
Commit 9db8544a authored by Bo-Chun Chen's avatar Bo-Chun Chen
Browse files

Add database logic

parent 3a0474f4
No related branches found
No related tags found
6 merge requests!147Merge previous default branch feat-cod-rmq into main,!85kill nginx process running under user from login node,!51Fix acct create wait,!39WIP:Feat cod rmq,!38WIP: Feat cod rmq,!24Add notify user agent
...@@ -3,8 +3,10 @@ import sys ...@@ -3,8 +3,10 @@ import sys
import json import json
import rc_util import rc_util
import smtplib import smtplib
import dataset
from rc_rmq import RCRMQ from rc_rmq import RCRMQ
from jinja2 import Template from jinja2 import Template
from datetime import datetime
import mail_config as mail_cfg import mail_config as mail_cfg
task = 'notify_user' task = 'notify_user'
...@@ -12,6 +14,9 @@ task = 'notify_user' ...@@ -12,6 +14,9 @@ task = 'notify_user'
args = rc_util.get_args() args = rc_util.get_args()
logger = rc_util.get_logger(args) logger = rc_util.get_logger(args)
db = dataset.connect(f'sqlite:///.agent_db/{task}.db')
table = db['notified_user']
# Instantiate rabbitmq object # Instantiate rabbitmq object
rc_rmq = RCRMQ({'exchange': 'RegUsr', 'exchange_type': 'topic'}) rc_rmq = RCRMQ({'exchange': 'RegUsr', 'exchange_type': 'topic'})
...@@ -24,16 +29,40 @@ def notify_user(ch, method, properties, body): ...@@ -24,16 +29,40 @@ def notify_user(ch, method, properties, body):
msg['success'] = False msg['success'] = False
try: try:
#Send email to user
receiver = [user_mail, mail_cfg.My_mail]
message = Template(mail_cfg.Whole_mail).render(username=username, to=user_email)
if args.dry_run: # Search username in database
logger.info("smtp.sendmail(sender, receiver, message)") record = table.find_one(username=username)
if record:
# Update counter
count = record['count']
table.update({'username': username, 'count': count + 1}, ['username'])
logger.debug(f'User {username} counter updated to {count + 1}')
else: else:
smtp = smtplib.SMTP(mail_cfg.Server) # Send email to user
smtp.sendmail(sender, receiver, message) receiver = [user_mail, mail_cfg.My_mail]
message = Template(mail_cfg.Whole_mail).render(username=username, to=user_email)
if args.dry_run:
logger.info(f'smtp = smtplib.SMTP({mail_cfg.Server})')
logger.info(f'smtp.sendmail({sender}, {receiver}, message)')
logger.info(f"table.insert({'username': {username}, 'count': 1, 'sent_at': {datetime.now()}})")
else:
smtp = smtplib.SMTP(mail_cfg.Server)
smtp.sendmail(sender, receiver, message)
logger.debug(f'Email sent to: {user_email}')
table.insert({
'username': username,
'count': 1,
'sent_at': datetime.now()
})
logger.debug(f'User {username} inserted into database')
msg['success'] = True msg['success'] = True
except Exception as exception: except Exception as exception:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment