diff --git a/mssql.py b/mssql.py index dd163f8..fac960f 100644 --- a/mssql.py +++ b/mssql.py @@ -1,190 +1,294 @@ #!/usr/bin/env python3 import znc -import pymssql import re +import multiprocessing +import warnings +import traceback + from datetime import datetime from contextlib import closing class mssql(znc.Module): - module_types = [znc.CModInfo.GlobalModule] - description = "Log all channels to a Microsoft SQL Server database" - has_args = False + module_types = [znc.CModInfo.GlobalModule] + description = "Log all channels to a Microsoft SQL Server database" + has_args = False - sql_server = "db.spd.internal" - sql_user = "_srv_znc" - sql_password = "zncznc" - sql_database = "znc" + log_queue = multiprocessing.SimpleQueue() + internal_log = None + hook_debugging = True -# def OnLoad(self, args, message): -# self.PutModule("Loading MSSQL...") -# return True + def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network): + """ + Adds the log line to database write queue. + """ + regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE) + message = regex.sub("", message) + message = self.colourstrip(message) + + self.log_queue.put({ + 'znc_user': self.GetUser().GetUserName() if self.GetUser() is not None else None, + 'code': code, + 'channel': channel, + 'host': host, + 'user': user, + 'ident': ident, + 'user_mode': user_mode, + 'date': date, + 'target_user': target_user, + 'message': message, + 'network': network + }) - def OnOp(self, user, target_user, channel, noChange): - self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return True + def OnLoad(self, args, message): + """ + This module hook is called when a module is loaded. + :type args: const CString & + :type args: CString & + :rtype: bool + :param args: The arguments for the modules. + :param message: A message that may be displayed to the user after loading the module. + :return: True if the module loaded successfully, else False. + """ + self.internal_log = InternalLog(self.GetSavePath()) + self.debug_hook() - def OnPrivNotice(self, user, message): - self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return True + try: + db = MSSQLDatabase({ + 'host': "db.spd.internal", + 'user': "_srv_znc", + 'passwd': "zncznc", + 'db': "znc" + }) + + multiprocessing.Process(target=DatabaseThread.worker_safe, + args=(db, self.log_queue, self.internal_log)).start() + return True + except Exception as e: + message.s = str(e) - def OnDeop(self, user, target_user, channel, noChange): - self.insert("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return True + with self.internal_log.error() as target: + target.write('Could not initialize module caused by: {} {}\n'.format(type(e), str(e))) + target.write('Stack trace: ' + traceback.format_exc()) + target.write('\n') - def OnVoice(self, user, target_user, channel, noChange): - self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return True + return False - def OnDevoice(self, user, target_user, channel, noChange): - self.insert("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return True + def OnOp(self, user, target_user, channel, noChange): + self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) + return True - def OnMode(self, user, channel, mode, argument, added, noChange): - mode = chr(mode) - code = "MODE" - if mode == "b": - code = "BAN" if added else "UNBAN" - else: - sign = "+" if added else "-" - argument = "{0}{1} {2}".format(sign, mode, argument) + def OnPrivNotice(self, user, message): + self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) + return True - self.insert(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork())) - return True + def OnDeop(self, user, target_user, channel, noChange): + self.insert("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) + return True - def OnQuit(self, user, message, channels): - for channel in channels: - self.insert("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) - return True + def OnVoice(self, user, target_user, channel, noChange): + self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) + return True - def OnNick(self, user, new_nick, channels): - for channel in channels: - self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork())) - return True + def OnDevoice(self, user, target_user, channel, noChange): + self.insert("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) + return True - def OnKick(self, user, target_nick, channel, message): - self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork())) - return True + def OnMode(self, user, channel, mode, argument, added, noChange): + mode = chr(mode) + code = "MODE" + if mode == "b": + code = "BAN" if added else "UNBAN" + else: + sign = "+" if added else "-" + argument = "{0}{1} {2}".format(sign, mode, argument) - def OnJoin(self, user, channel): - self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork())) - return True + self.insert(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork())) + return True - def OnPart(self, user, channel, message): - self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) - return True + def OnQuit(self, user, message, channels): + for channel in channels: + self.insert("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) + return True - def OnUserAction (self, target, message): - (channel, user) = self.resolveTarget(target) - if channel == None: - return True - return self.OnChanAction(user, channel, message) + def OnNick(self, user, new_nick, channels): + for channel in channels: + self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork())) + return True - def OnPrivMsg(self, nick, message): - self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return True + def OnKick(self, user, target_nick, channel, message): + self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork())) + return True - def OnUserMsg (self, target, message): - (channel, user) = self.resolveTarget(target) - if channel == None: - return True - return self.OnChanMsg(user, channel, message) + def OnJoin(self, user, channel): + self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork())) + return True - def OnChanAction(self, user, channel, message): - self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return True + def OnPart(self, user, channel, message): + self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) + return True - def OnChanMsg(self, user, channel, message): - self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork())) - return True + def OnUserAction (self, target, message): + (channel, user) = self.resolveTarget(target) + if channel == None: + return True + return self.OnChanAction(user, channel, message) - def OnTopic(self, user, channel, message): - self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return True + def OnPrivMsg(self, nick, message): + self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) + return True -# def OnLoad(self, args, message): -# match = re.search("(.*?):(.*?)@(.*)", args) -# if match: -# self.username = match.group(1) -# self.password = match.group(2) -# self.host = match.group(3) -# return True -# else: -# return False + def OnUserMsg (self, target, message): + (channel, user) = self.resolveTarget(target) + if channel == None: + return True + return self.OnChanMsg(user, channel, message) - def findMode(self, channel, user): - realUser = channel.FindNick(user.GetNick()) - if not realUser: - return None - return str(realUser.GetPermChar()) + def OnChanAction(self, user, channel, message): + self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) + return True - def resolveTarget(self, target): - target = target.s - channel = self.GetNetwork().FindChan(target) - if channel == None: - return (None, None) - user = self.GetNetwork().GetIRCNick() - return (channel, user) + def OnChanMsg(self, user, channel, message): + self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork())) + return True - def colourstrip(self, data): - find = data.find('\x03') - while find > -1: - done = False - data = data[0:find] + data[find+1:] - if len(data) <= find+1: - done = True - try: - assert not done - assert int(data[find]) - while True: - assert int(data[find]) - data = data[0:find] + data[find+1:] - except: - if not done: - if data[find] != ',': done = True + def OnTopic(self, user, channel, message): + self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) + return True - if (not done) and (len(data) > find+1) and (data[find] == ','): - try: - assert not done - assert int(data[find+1]) - data = data[0:find] + data[find+1:] - data = data[0:find] + data[find+1:] - except: - done = True - try: - assert not done - while True: - assert int(data[find]) - data = data[0:find] + data[find+1:] - except: pass + def findMode(self, channel, user): + realUser = channel.FindNick(user.GetNick()) + if not realUser: + return None + return str(realUser.GetPermChar()) - find = data.find('\x03') - data = data.replace('\x02','') - data = data.replace('\x1d','') - data = data.replace('\x1f','') - data = data.replace('\x16','') - data = data.replace('\x0f','') - return data + def resolveTarget(self, target): + target = target.s + channel = self.GetNetwork().FindChan(target) + if channel == None: + return (None, None) + user = self.GetNetwork().GetIRCNick() + return (channel, user) - def insert(self, code, channel, host, user, ident, user_mode, date, target_user, message, network): -# self.PutModule("Inserting row...") - regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE) + def colourstrip(self, data): + find = data.find('\x03') + while find > -1: + done = False + data = data[0:find] + data[find+1:] + if len(data) <= find+1: + done = True + try: + assert not done + assert int(data[find]) + while True: + assert int(data[find]) + data = data[0:find] + data[find+1:] + except: + if not done: + if data[find] != ',': done = True - try: - with closing(pymssql.connect ("db.spd.internal", "_srv_znc", "znc", "znc")) as conn: - with closing(conn.cursor()) as cursor: -# sql = "INSERT INTO log (date, code, network, channel, host, [user], ident, mode, target, message) VALUES(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)" - sql = "EXEC dbo.add_log (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - if message is not None: -# message = message.encode('utf-8').decode('utf-8', 'replace') - message = regex.sub("", message) - message = self.colourstrip(message) - znc_user = self.GetUser().GetUserName() if self.GetUser() is not None else None - params = (znc_user, str(date), code, network, channel, host, user, ident, user_mode, target_user, message) -# cursor.execute (sql, params) - cursor.callproc("add_log", params) - conn.commit() - conn.close() + if (not done) and (len(data) > find+1) and (data[find] == ','): + try: + assert not done + assert int(data[find+1]) + data = data[0:find] + data[find+1:] + data = data[0:find] + data[find+1:] + except: + done = True + try: + assert not done + while True: + assert int(data[find]) + data = data[0:find] + data[find+1:] + except: pass - except Exception as e: - self.PutModule("Could not save {0} to database caused by: {1} - {2}".format(code, type(e), str(e))) \ No newline at end of file + find = data.find('\x03') + data = data.replace('\x02','') + data = data.replace('\x1d','') + data = data.replace('\x1f','') + data = data.replace('\x16','') + data = data.replace('\x0f','') + return data + +class DatabaseThread: + @staticmethod + def worker_safe(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None: + try: + DatabaseThread.worker(db, log_queue, internal_log) + except Exception as e: + with internal_log.error() as target: + target.write('Unrecoverable exception in worker thread: {0} {1}\n'.format(type(e), str(e))) + target.write('Stack trace: ' + traceback.format_exc()) + target.write('\n') + raise + + @staticmethod + def worker(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None: + db.connect() + + while True: + item = log_queue.get() + if item is None: + break + + try: + db.ensure_connected() + db.insert_into('logs', item) + except Exception as e: + sleep_for = 10 + + with internal_log.error() as target: + target.write('Could not save to database caused by: {0} {1}\n'.format(type(e), str(e))) + if 'open' in dir(db.conn): + target.write('Database handle state: {}\n'.format(db.conn.open)) + target.write('Stack trace: ' + traceback.format_exc()) + target.write('Current log: ') + json.dump(item, target) + target.write('\n\n') + target.write('Retry in {} s\n'.format(sleep_for)) + + sleep(sleep_for) + + with internal_log.error() as target: + target.write('Retrying now.\n'.format(sleep_for)) + log_queue.put(item) + +class InternalLog: + def __init__(self, save_path: str): + self.save_path = save_path + + def debug(self): + return self.open('debug') + + def error(self): + return self.open('error') + + def open(self, level: str): + target = open(os.path.join(self.save_path, level + '.log'), 'a') + line = 'Log opened at: {} UTC\n'.format(datetime.utcnow()) + target.write(line) + target.write('=' * len(line) + '\n\n') + return target + + + +class Database: + def __init__(self, dsn: dict): + self.dsn = dsn + self.conn = None + +class MSSQLDatabase(Database): + def connect(self) -> None: + import pymssql + self.conn = pymssql.connect (**self.dsn) + + def ensure_connected(self): + if self.conn.status == 0: + self.connect() + + def insert_into(self, table, row): +# self.PutModule("Inserting row...") + cols = ', '.join('"{}"'.format(col) for col in row.keys()) + vals = ', '.join('%({})s'.format(col) for col in row.keys()) + self.conn.cursor().callproc("add_log", vals) + self.conn.commit() \ No newline at end of file