diff --git a/mssql.py b/mssql.py index 84a7fb7..e943381 100644 --- a/mssql.py +++ b/mssql.py @@ -4,7 +4,6 @@ import znc import re import queue import threading -import warnings import traceback import inspect import os @@ -14,586 +13,883 @@ import uuid from datetime import datetime from time import sleep -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 = True + module_types = [znc.CModInfo.GlobalModule] + description = "Log all channels to a Microsoft SQL Server database" + has_args = True - log_queue = None - worker_threads = None - internal_log = None - hook_debugging = False - debug_log_saves = False - default_workers = 4 - num_workers = 4 - max_workers = 32 - batch_size = 50 - db_config = None - next_worker_id = 0 + log_queue = None + worker_threads = None + internal_log = None + hook_debugging = False + debug_log_saves = False + default_workers = 4 + num_workers = 4 + max_workers = 32 + batch_size = 50 + db_config = None + next_worker_id = 0 - def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network): - """ - Adds the log line to database write queue. - """ + def put_log( + self, + code, + channel, + host, + user, + ident, + user_mode, + date, + target_user, + message, + network): + """ + Adds the log line to database write queue. + """ - decolouriser = re.compile(r"\x1f|\x02|\x1d|\x16|\x11|\x1e|\x0f|\x03(?:\d{1,2}(?:,\d{1,2})?)?|\x04(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE) + decolouriser = re.compile( + r"\x1f|\x02|\x1d|\x16|\x11|\x1e|\x0f|\x03(?:\d{1,2}(?:,\d{1,2})?)?|\x04(?:\d{1,2}(?:,\d{1,2})?)?", + re.UNICODE) - if message is not None: - message.encode('utf8', 'replace').decode('utf8') - message = decolouriser.sub("", message) + if message is not None: + message = message.encode('utf8', 'replace').decode('utf8') + message = decolouriser.sub("", message) - if self.log_queue is None: - return + if self.log_queue is None: + return - queue_len = self.log_queue.qsize() if self.log_queue is not None else 0 - if queue_len > 10: - with self.internal_log.error() as target: - target.write('Queue length high: {} entries at {}\n'.format(queue_len, datetime.utcnow())) + queue_len = self.log_queue.qsize() if self.log_queue is not None else 0 + if queue_len > 10: + with self.internal_log.error() as target: + target.write( + 'Queue length high: {} entries at {}\n'.format( + queue_len, datetime.utcnow())) - try: - 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': str(date), - 'target_user': target_user, - 'message': message, - 'network': network, - '_id': uuid.uuid4() - }) - except Exception as e: - with self.internal_log.error() as target: - target.write('Failed to enqueue log event: {0} {1}\n'.format(type(e), str(e))) - target.write('Stack trace: ' + traceback.format_exc()) - target.write('\n') + try: + 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': str(date), + 'target_user': target_user, + 'message': message, + 'network': network, + '_id': uuid.uuid4() + }) + except Exception as e: + with self.internal_log.error() as target: + target.write( + 'Failed to enqueue log event: {0} {1}\n'.format( + type(e), str(e))) + target.write('Stack trace: ' + traceback.format_exc()) + target.write('\n') - def debug_hook(self): - """ - Dumps parent calling method name and its arguments to debug logfile. - """ + def debug_hook(self): + """ + Dumps parent calling method name and its arguments to debug logfile. + """ - if self.hook_debugging is not True: - return + if (self.hook_debugging is not True + or not self.internal_log.debug_enabled): + return - frameinfo = inspect.stack()[1] - argvals = frameinfo.frame.f_locals + frameinfo = inspect.stack()[1] + argvals = frameinfo.frame.f_locals + with self.internal_log.debug() as target: + target.write('Called method: ' + frameinfo.function + '()\n') + for argname in argvals: + if argname == 'self': + continue + target.write( + ' ' + + argname + + ' -> ' + + pprint.pformat( + argvals[argname]) + + '\n') + target.write('\n') + + def _start_workers(self, count): + """Start `count` more worker threads on the shared queue.""" + for _ in range(count): + idx = self.next_worker_id + self.next_worker_id += 1 + db = MSSQLDatabase(self.db_config) + worker_thread = threading.Thread( + target=DatabaseThread.worker_safe, + args=( + db, + self.log_queue, + self.internal_log, + self.debug_log_saves, + self.batch_size), + daemon=True, + name=f'DatabaseWorker-{idx}') + worker_thread.start() + self.worker_threads.append(worker_thread) + + def OnModCommand(self, command): + """ + Commands: Status, AddWorkers , Debug [on|off|status] + """ + parts = str(command).split() + if not parts: + return + cmd = parts[0].lower() + + if cmd == 'status': + alive = len([t for t in self.worker_threads if t.is_alive()] + ) if self.worker_threads else 0 + qsize = self.log_queue.qsize() if self.log_queue is not None else 0 + self.PutModule( + 'Workers alive: {}, queue length: {}, batch size: {}, ' + 'debug: {}'.format( + alive, + qsize, + self.batch_size, + 'on' if self.internal_log.debug_enabled else 'off')) + elif cmd == 'addworkers': + try: + n = int(parts[1]) + except (IndexError, ValueError): + self.PutModule('Usage: AddWorkers ') + return + current = len(self.worker_threads) + n = max(0, min(n, self.max_workers - current)) + if n == 0: + self.PutModule( + 'No workers added (limit is {}, currently {})'.format( + self.max_workers, current)) + return + self._start_workers(n) + self.PutModule( + 'Started {} workers, now {}'.format( + n, len( + self.worker_threads))) + elif cmd == 'debug': + setting = parts[1].lower() if len(parts) > 1 else 'status' + if setting == 'status': + state = 'on' if self.internal_log.debug_enabled else 'off' + self.PutModule('Debug logging is {}'.format(state)) + elif setting in ('on', 'off'): + enabled = setting == 'on' + self.internal_log.set_debug_enabled(enabled) + self.PutModule('Debug logging {}'.format( + 'enabled' if enabled else 'disabled')) + else: + self.PutModule('Usage: Debug [on|off|status]') + else: + self.PutModule( + 'Commands: Status, AddWorkers , Debug [on|off|status]') + + def OnLoad(self, args, message): + """ + This module hook is called when a module is loaded. + """ + # Initialize internal log first for debugging + self.internal_log = InternalLog(self.GetSavePath()) + + # Parse num_workers from args (default 4) + self.num_workers = self.default_workers + args_str = "" + if args: + try: + args_str = str(args) if not hasattr(args, 's') else args.s + if args_str: + self.num_workers = int(args_str.strip()) + if self.num_workers < 1: + self.num_workers = 1 + if self.num_workers > self.max_workers: + self.num_workers = self.max_workers + except (ValueError, AttributeError, TypeError) as e: with self.internal_log.debug() as target: - target.write('Called method: ' + frameinfo.function + '()\n') - for argname in argvals: - if argname == 'self': - continue - target.write(' ' + argname + ' -> ' + pprint.pformat(argvals[argname]) + '\n') - target.write('\n') + target.write( + 'Failed to parse args: {} (args_str: "{}", type: {})\n'.format( + e, args_str, type(args))) - def _start_workers(self, count): - """Start `count` more worker threads on the shared queue.""" - for _ in range(count): - idx = self.next_worker_id - self.next_worker_id += 1 - db = MSSQLDatabase(self.db_config) - worker_thread = threading.Thread( - target=DatabaseThread.worker_safe, - args=(db, self.log_queue, self.internal_log, self.debug_log_saves, self.batch_size), - daemon=True, - name=f'DatabaseWorker-{idx}' - ) - worker_thread.start() - self.worker_threads.append(worker_thread) + with self.internal_log.debug() as target: + target.write( + 'mssql module loaded at: {}\n'.format( + datetime.utcnow())) + target.write('save_path: {}\n'.format(self.GetSavePath())) + target.write( + 'raw args: {} (type: {})\n'.format( + repr(args), type(args).__name__)) + target.write('parsed args_str: "{}"\n'.format(args_str)) + target.write('num_workers: {}\n'.format(self.num_workers)) + target.write('\n') - def OnModCommand(self, command): - """ - Commands: Status, AddWorkers - """ - parts = str(command).split() - if not parts: - return - cmd = parts[0].lower() + with self.internal_log.error() as target: + target.write( + 'mssql module loaded at: {}\n'.format( + datetime.utcnow())) + target.write('hook_debugging: {}\n'.format(self.hook_debugging)) + target.write('num_workers: {}\n'.format(self.num_workers)) + target.write('\n') - if cmd == 'status': - alive = len([t for t in self.worker_threads if t.is_alive()]) if self.worker_threads else 0 - qsize = self.log_queue.qsize() if self.log_queue is not None else 0 - self.PutModule('Workers alive: {}, queue length: {}, batch size: {}'.format(alive, qsize, self.batch_size)) - elif cmd == 'addworkers': - try: - n = int(parts[1]) - except (IndexError, ValueError): - self.PutModule('Usage: AddWorkers ') - return - current = len(self.worker_threads) - n = max(0, min(n, self.max_workers - current)) - if n == 0: - self.PutModule('No workers added (limit is {}, currently {})'.format(self.max_workers, current)) - return - self._start_workers(n) - self.PutModule('Started {} workers, now {}'.format(n, len(self.worker_threads))) - else: - self.PutModule('Commands: Status, AddWorkers ') + self.debug_hook() - def OnLoad(self, args, message): - """ - This module hook is called when a module is loaded. - """ - # Initialize internal log first for debugging - self.internal_log = InternalLog(self.GetSavePath()) + try: + self.log_queue = queue.Queue() + self.db_config = { + 'host': "dbx.spd.internal", + 'user': "_srv_znc", + 'password': "znc", + 'database': "znc", + 'port': '1433', + 'appname': 'ZNC Logger' + } - # Parse num_workers from args (default 4) - self.num_workers = self.default_workers - args_str = "" - if args: - try: - args_str = str(args) if not hasattr(args, 's') else args.s - if args_str: - self.num_workers = int(args_str.strip()) - if self.num_workers < 1: - self.num_workers = 1 - if self.num_workers > self.max_workers: - self.num_workers = self.max_workers - except (ValueError, AttributeError, TypeError) as e: - with self.internal_log.debug() as target: - target.write('Failed to parse args: {} (args_str: "{}", type: {})\n'.format(e, args_str, type(args))) + self.worker_threads = [] + self.next_worker_id = 0 + self._start_workers(self.num_workers) - with self.internal_log.debug() as target: - target.write('mssql module loaded at: {}\n'.format(datetime.utcnow())) - target.write('save_path: {}\n'.format(self.GetSavePath())) - target.write('raw args: {} (type: {})\n'.format(repr(args), type(args).__name__)) - target.write('parsed args_str: "{}"\n'.format(args_str)) - target.write('num_workers: {}\n'.format(self.num_workers)) - target.write('\n') + return True + except Exception as e: + message.s = str(e) - with self.internal_log.error() as target: - target.write('mssql module loaded at: {}\n'.format(datetime.utcnow())) - target.write('hook_debugging: {}\n'.format(self.hook_debugging)) - target.write('num_workers: {}\n'.format(self.num_workers)) - target.write('\n') + 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') - self.debug_hook() + return False - try: - self.log_queue = queue.Queue() - self.db_config = { - 'host': "dbx.spd.internal", - 'user': "_srv_znc", - 'password': "znc", - 'database': "znc", - 'port': '1433', - 'appname': 'ZNC Logger' - } + def OnUnload(self): + if self.log_queue is None: + return - self.worker_threads = [] - self.next_worker_id = 0 - self._start_workers(self.num_workers) + with self.internal_log.debug() as target: + target.write( + 'Unloading module, draining queue length {}\n'.format( + self.log_queue.qsize())) - return True - except Exception as e: - message.s = str(e) + # Send one sentinel value per worker thread + if self.worker_threads: + for _ in self.worker_threads: + self.log_queue.put(None) - 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') + self.log_queue.join() - return False + if self.worker_threads is not None: + for worker_thread in self.worker_threads: + worker_thread.join(timeout=30) + with self.internal_log.debug() as target: + target.write('Unload complete, queue drained\n') - def OnUnload(self): - if self.log_queue is None: - return - - with self.internal_log.debug() as target: - target.write('Unloading module, draining queue length {}\n'.format(self.log_queue.qsize())) - - # Send one sentinel value per worker thread + def __del__(self): + # Terminate worker threads cleanly if possible. + try: + if self.log_queue is not None: + # Send sentinel values for all workers if self.worker_threads: - for _ in self.worker_threads: - self.log_queue.put(None) - - self.log_queue.join() + for _ in self.worker_threads: + self.log_queue.put_nowait(None) + except Exception: + pass - if self.worker_threads is not None: - for worker_thread in self.worker_threads: - worker_thread.join(timeout=30) + def OnIRCConnected(self): + """ + This module hook is called after a successful login to IRC. + :rtype: None + """ + self.debug_hook() + self.put_log("CONNECT", None, None, None, None, None, + datetime.now(), None, None, str(self.GetNetwork())) + return znc.CONTINUE - with self.internal_log.debug() as target: - target.write('Unload complete, queue drained\n') + def OnIRCDisconnected(self): + """ + This module hook is called when a user gets disconnected from IRC. + :rtype: None + """ + self.debug_hook() + self.put_log("DISCONNECT", None, None, None, None, None, + datetime.now(), None, None, str(self.GetNetwork())) + return znc.CONTINUE - def __del__(self): - # Terminate worker threads cleanly if possible. + def OnOp(self, user, target_user, channel, noChange): + self.debug_hook() + self.put_log("OP", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + target_user.GetNick(), + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnChanNotice(self, nick, channel, message): + """ + Called when we receive a channel NOTICE message from IRC. + :type nick: CNick & + :type channel: CChan & + :type message: CString & + :rtype: EModRet + """ + self.debug_hook() + self.put_log("NOTICE", + channel.GetName(), + nick.GetHost(), + nick.GetNick(), + nick.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnPrivNotice(self, user, message): + self.debug_hook() + self.put_log("NOTICE", + None, + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnUserNotice(self, target, message): + """ + This module hook is called when a user sends a NOTICE message. + :type target: CString & + :type message: CString & + :rtype: EModRet + """ + self.debug_hook() + network = self.GetNetwork() + if network: + self.put_log("NOTICE", + None, + target.GetHost(), + target.GetNick(), + target.GetIdent(), + None, + datetime.now(), + target.GetNick(), + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnDeop(self, user, target_user, channel, noChange): + self.debug_hook() + self.put_log("DEOP", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + target_user.GetNick(), + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnVoice(self, user, target_user, channel, noChange): + self.debug_hook() + self.put_log("VOICE", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + target_user.GetNick(), + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnDevoice(self, user, target_user, channel, noChange): + self.debug_hook() + self.put_log("DEVOICE", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + target_user.GetNick(), + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnMode(self, user, channel, mode, argument, added, noChange): + self.debug_hook() + 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) + + self.put_log(code, + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + argument, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnQuit(self, user, message, channels): + self.debug_hook() + for channel in channels: + self.put_log("QUIT", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnNick(self, user, new_nick, channels): + self.debug_hook() + for channel in channels: + self.put_log("NICK", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + new_nick, + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnKick(self, user, target_nick, channel, message): + self.debug_hook() + self.put_log("KICK", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + target_nick, + message, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnJoin(self, user, channel): + self.debug_hook() + self.put_log("JOIN", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + None, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnPart(self, user, channel, message): + self.debug_hook() + self.put_log("PART", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnUserAction(self, target, message): + self.debug_hook() + (channel, user) = self.resolveTarget(target) + if channel is None: + local_nick = self.GetNetwork().GetIRCNick() + self.put_log( + "ME", + None, + local_nick.GetHost() if local_nick is not None else None, + local_nick.GetNick() if local_nick is not None else None, + local_nick.GetIdent() if local_nick is not None else None, + None, + datetime.now(), + target.s, + message.s, + str(self.GetNetwork()), + ) + return znc.CONTINUE + return self.OnChanAction(user, channel, message) + + def OnPrivAction(self, nick, message): + """ + Called when we receive a private CTCP ACTION ("/me" in query) from IRC. + :type nick: CNick & + :type message: CString & + :rtype: EModRet + """ + self.debug_hook() + self.put_log("ME", + None, + nick.GetHost(), + nick.GetNick(), + nick.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnPrivMsg(self, nick, message): + self.debug_hook() + self.put_log("MSG", + None, + nick.GetHost(), + nick.GetNick(), + nick.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnPrivCTCP(self, nick, message): + self.debug_hook() + self.put_log("CTCP", + None, + nick.GetHost(), + nick.GetNick(), + nick.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnChanCTCP(self, user, channel, message): + self.debug_hook() + self.put_log("CTCP", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnUserCTCP(self, target, message): + self.debug_hook() + local_nick = self.GetNetwork().GetIRCNick() + if local_nick is not None: + self.put_log( + "CTCP", + None, + local_nick.GetHost(), + local_nick.GetNick(), + local_nick.GetIdent(), + None, + datetime.now(), + target.s, + message.s, + str(self.GetNetwork()), + ) + else: + self.put_log( + "CTCP", + None, + None, + None, + None, + None, + datetime.now(), + target.s, + message.s, + str(self.GetNetwork()), + ) + return znc.CONTINUE + + def OnUserMsg(self, target, message): + self.debug_hook() + (channel, user) = self.resolveTarget(target) + if channel is None: + local_nick = self.GetNetwork().GetIRCNick() + if local_nick is not None: + self.put_log( + "MSG", + None, + local_nick.GetHost(), + local_nick.GetNick(), + local_nick.GetIdent(), + None, + datetime.now(), + target.s, + message.s, + str(self.GetNetwork()), + ) + else: + self.put_log( + "MSG", + None, + None, + None, + None, + None, + datetime.now(), + target.s, + message.s, + str(self.GetNetwork()), + ) + return znc.CONTINUE + return self.OnChanMsg(user, channel, message) + + def OnChanAction(self, user, channel, message): + self.debug_hook() + self.put_log("ME", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return True + + def OnChanMsg(self, user, channel, message): + self.debug_hook() + self.put_log("SAY", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + self.findMode(channel, + user), + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def OnTopic(self, user, channel, message): + self.debug_hook() + self.put_log("TOPIC", + channel.GetName(), + user.GetHost(), + user.GetNick(), + user.GetIdent(), + None, + datetime.now(), + None, + message.s, + str(self.GetNetwork())) + return znc.CONTINUE + + def findMode(self, channel, user): + realUser = channel.FindNick(user.GetNick()) + if not realUser: + return None + return str(realUser.GetPermChar()) + + def resolveTarget(self, target): + target = target.s + channel = self.GetNetwork().FindChan(target) + if channel is None: + return (None, None) + user = self.GetNetwork().GetIRCNick() + return (channel, user) + + 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 BaseException: + if not done: + if data[find] != ',': + done = True + + if (not done) and (len(data) > find + 1) and (data[find] == ','): try: - if self.log_queue is not None: - # Send sentinel values for all workers - if self.worker_threads: - for _ in self.worker_threads: - self.log_queue.put_nowait(None) - except Exception: - pass + assert not done + assert int(data[find + 1]) + data = data[0:find] + data[find + 1:] + data = data[0:find] + data[find + 1:] + except BaseException: + done = True + try: + assert not done + while True: + assert int(data[find]) + data = data[0:find] + data[find + 1:] + except BaseException: + pass - def OnIRCConnected(self): - """ - This module hook is called after a successful login to IRC. - :rtype: None - """ - self.debug_hook() - self.put_log("CONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork())) - return znc.CONTINUE + find = data.find('\x03') - def OnIRCDisconnected(self): - """ - This module hook is called when a user gets disconnected from IRC. - :rtype: None - """ - self.debug_hook() - self.put_log("DISCONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork())) - return znc.CONTINUE + data = data.replace('\x02', '') + data = data.replace('\x1d', '') + data = data.replace('\x1f', '') + data = data.replace('\x16', '') + data = data.replace('\x0f', '') + return data - def OnOp(self, user, target_user, channel, noChange): - self.debug_hook() - self.put_log("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnChanNotice(self, nick, channel, message): - """ - Called when we receive a channel NOTICE message from IRC. - :type nick: CNick & - :type channel: CChan & - :type message: CString & - :rtype: EModRet - """ - self.debug_hook() - self.put_log("NOTICE", channel.GetName(), nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnPrivNotice(self, user, message): - self.debug_hook() - self.put_log("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnUserNotice(self, target, message): - """ - This module hook is called when a user sends a NOTICE message. - :type target: CString & - :type message: CString & - :rtype: EModRet - """ - self.debug_hook() - network = self.GetNetwork() - if network: - self.put_log("NOTICE", None, target.GetHost(), target.GetNick(), target.GetIdent(), None, datetime.now(), target.GetNick(), message, str(self.GetNetwork())) - - def OnDeop(self, user, target_user, channel, noChange): - self.debug_hook() - self.put_log("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnVoice(self, user, target_user, channel, noChange): - self.debug_hook() - self.put_log("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnDevoice(self, user, target_user, channel, noChange): - self.debug_hook() - self.put_log("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnMode(self, user, channel, mode, argument, added, noChange): - self.debug_hook() - 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) - - self.put_log(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork())) - return znc.CONTINUE - - def OnQuit(self, user, message, channels): - self.debug_hook() - for channel in channels: - self.put_log("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) - return znc.CONTINUE - - def OnNick(self, user, new_nick, channels): - self.debug_hook() - for channel in channels: - self.put_log("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnKick(self, user, target_nick, channel, message): - self.debug_hook() - self.put_log("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork())) - return znc.CONTINUE - - def OnJoin(self, user, channel): - self.debug_hook() - self.put_log("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork())) - return znc.CONTINUE - - def OnPart(self, user, channel, message): - self.debug_hook() - self.put_log("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) - return znc.CONTINUE - - def OnUserAction (self, target, message): - self.debug_hook() - (channel, user) = self.resolveTarget(target) - if channel == None: - self.put_log("ME", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) - return znc.CONTINUE - return self.OnChanAction(user, channel, message) - - def OnPrivAction(self, nick, message): - """ - Called when we receive a private CTCP ACTION ("/me" in query) from IRC. - :type nick: CNick & - :type message: CString & - :rtype: EModRet - """ - self.debug_hook() - self.put_log("ME", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnPrivMsg(self, nick, message): - self.debug_hook() - self.put_log("MSG", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnPrivCTCP(self, nick, message): - self.debug_hook() - self.put_log("CTCP", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnChanCTCP(self, user, channel, message): - self.debug_hook() - self.put_log("CTCP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnUserCTCP(self, target, message): - self.debug_hook() - local_nick = self.GetNetwork().GetIRCNick() - if local_nick is not None: - self.put_log( - "CTCP", - None, - local_nick.GetHost(), - local_nick.GetNick(), - local_nick.GetIdent(), - None, - datetime.now(), - target.s, - message.s, - str(self.GetNetwork()), - ) - else: - self.put_log( - "CTCP", - None, - None, - None, - None, - None, - datetime.now(), - target.s, - message.s, - str(self.GetNetwork()), - ) - return znc.CONTINUE - - def OnUserMsg (self, target, message): - self.debug_hook() - (channel, user) = self.resolveTarget(target) - if channel == None: - local_nick = self.GetNetwork().GetIRCNick() - if local_nick is not None: - self.put_log( - "MSG", - None, - local_nick.GetHost(), - local_nick.GetNick(), - local_nick.GetIdent(), - None, - datetime.now(), - target.s, - message.s, - str(self.GetNetwork()), - ) - else: - self.put_log( - "MSG", - None, - None, - None, - None, - None, - datetime.now(), - target.s, - message.s, - str(self.GetNetwork()), - ) - return znc.CONTINUE - return self.OnChanMsg(user, channel, message) - - def OnChanAction(self, user, channel, message): - self.debug_hook() - self.put_log("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return True - - def OnChanMsg(self, user, channel, message): - self.debug_hook() - self.put_log("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def OnTopic(self, user, channel, message): - self.debug_hook() - self.put_log("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) - return znc.CONTINUE - - def findMode(self, channel, user): - realUser = channel.FindNick(user.GetNick()) - if not realUser: - return None - return str(realUser.GetPermChar()) - - 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 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 - - 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 - - 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, internal_log, debug_log_saves=False, batch_size=50) -> None: - while True: - try: - DatabaseThread.worker(db, log_queue, internal_log, debug_log_saves, batch_size) - break - except Exception as e: - with internal_log.error() as target: - target.write('Worker thread crashed: {0} {1}\n'.format(type(e), str(e))) - target.write('Stack trace: ' + traceback.format_exc()) - target.write('Restarting worker thread in 5 seconds...\n\n') - sleep(5) - continue + @staticmethod + def worker_safe( + db, + log_queue, + internal_log, + debug_log_saves=False, + batch_size=50) -> None: + while True: + try: + DatabaseThread.worker( + db, log_queue, internal_log, debug_log_saves, batch_size) + break + except Exception as e: + with internal_log.error() as target: + target.write( + 'Worker thread crashed: {0} {1}\n'.format( + type(e), str(e))) + target.write('Stack trace: ' + traceback.format_exc()) + target.write( + 'Restarting worker thread in 5 seconds...\n\n') + sleep(5) + continue - @staticmethod - def worker(db, log_queue, internal_log, debug_log_saves=False, batch_size=50) -> None: - stop = False - while not stop: - first = log_queue.get() - if first is None: - log_queue.task_done() # count the sentinel too, or join() hangs - return + @staticmethod + def worker( + db, + log_queue, + internal_log, + debug_log_saves=False, + batch_size=50) -> None: + stop = False + while not stop: + first = log_queue.get() + if first is None: + log_queue.task_done() # count the sentinel too, or join() hangs + return - # Collect more rows without blocking, up to batch_size - taken = 1 - items = [first] - while len(items) < batch_size: - try: - nxt = log_queue.get_nowait() - except queue.Empty: - break - taken += 1 - if nxt is None: - stop = True - break - items.append(nxt) + # Collect more rows without blocking, up to batch_size + taken = 1 + items = [first] + while len(items) < batch_size: + try: + nxt = log_queue.get_nowait() + except queue.Empty: + break + taken += 1 + if nxt is None: + stop = True + break + items.append(nxt) - try: - DatabaseThread.save_batch(db, items, log_queue, internal_log, debug_log_saves) - finally: - # exactly one task_done() per get(), after the batch is finished - for _ in range(taken): - log_queue.task_done() + try: + DatabaseThread.save_batch( + db, items, log_queue, internal_log, debug_log_saves) + finally: + # exactly one task_done() per get(), after the batch is + # finished + for _ in range(taken): + log_queue.task_done() - @staticmethod - def save_batch(db, items, log_queue, internal_log, debug_log_saves) -> None: - while True: - try: - db.ensure_connected(internal_log, str(items[0]['_id'])) - db.insert_batch(items) - if debug_log_saves: - with internal_log.debug() as target: - target.write('Saved {} log items queue_size={}\n'.format(len(items), log_queue.qsize())) - return - except Exception as e: - sleep_for = 10 + @staticmethod + def save_batch( + db, + items, + log_queue, + internal_log, + debug_log_saves) -> None: + while True: + try: + db.ensure_connected(internal_log, str(items[0]['_id'])) + db.insert_batch(items) + if debug_log_saves: + with internal_log.debug() as target: + target.write( + 'Saved {} log items queue_size={}\n'.format( + len(items), log_queue.qsize())) + return + 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))) - target.write('Stack trace: ' + traceback.format_exc()) - target.write('Batch of {} rows, first log: '.format(len(items))) - item_copy = items[0].copy() - item_copy['_id'] = str(item_copy['_id']) - json.dump(item_copy, target) - target.write('\n\n') - target.write('Retry in {} s\n'.format(sleep_for)) + with internal_log.error() as target: + target.write( + 'Could not save to database caused by: {0} {1}\n'.format( + type(e), str(e))) + target.write('Stack trace: ' + traceback.format_exc()) + target.write( + 'Batch of {} rows, first log: '.format( + len(items))) + item_copy = items[0].copy() + item_copy['_id'] = str(item_copy['_id']) + json.dump(item_copy, target) + target.write('\n\n') + target.write('Retry in {} s\n'.format(sleep_for)) + + sleep(sleep_for) - sleep(sleep_for) class InternalLog: - def __init__(self, save_path: str): - self.save_path = save_path + def __init__(self, save_path: str): + self.save_path = save_path + self.debug_enabled = False - def debug(self): - return self.open('debug') + def debug(self): + if not self.debug_enabled: + from io import StringIO + return StringIO() + return self.open('debug') - def error(self): - return self.open('error') + def error(self): + return self.open('error') - def open(self, level: str): - target = open(os.path.join(self.save_path, level + '.log'), 'a') - return target + def set_debug_enabled(self, enabled: bool): + self.debug_enabled = enabled + def open(self, level: str): + target = open(os.path.join(self.save_path, level + '.log'), 'a') + return target class Database: @@ -601,43 +897,55 @@ class Database: self.dsn = dsn self.conn = None + class MSSQLDatabase(Database): - def connect(self) -> None: - import pymssql - self.conn = pymssql.connect(**self.dsn) + def connect(self) -> None: + import pymssql + self.conn = pymssql.connect(**self.dsn) - def ensure_connected(self, internal_log, log_id): - # No per-row SELECT 1. A dead connection is detected when an - # insert fails; insert_batch() then drops it and the retry reconnects. - if self.conn is None: - self.connect() + def ensure_connected(self, internal_log, log_id): + # No per-row SELECT 1. A dead connection is detected when an + # insert fails; insert_batch() then drops it and the retry reconnects. + if self.conn is None: + self.connect() - def insert_batch(self, rows): - """Insert all rows, then commit once. All or nothing: on any error - the transaction is rolled back and the connection is dropped.""" + def insert_batch(self, rows): + """Insert all rows, then commit once. All or nothing: on any error + the transaction is rolled back and the connection is dropped.""" + try: + for row in rows: + params = ( + row['znc_user'], + row['date'], + row['code'], + row['network'], + row['channel'], + row['host'], + row['user'], + row['ident'], + row['user_mode'], + row['target_user'], + row['message']) + cur = self.conn.cursor() try: - for row in rows: - params = (row['znc_user'], row['date'], row['code'], row['network'], row['channel'], row['host'], row['user'], row['ident'], row['user_mode'], row['target_user'], row['message']) - cur = self.conn.cursor() - try: - cur.callproc("add_log", params) - finally: - try: - cur.close() - except Exception: - pass - self.conn.commit() - except Exception: - try: - self.conn.rollback() - except Exception: - pass - try: - self.conn.close() - except Exception: - pass - self.conn = None - raise + cur.callproc("add_log", params) + finally: + try: + cur.close() + except Exception: + pass + self.conn.commit() + except Exception: + try: + self.conn.rollback() + except Exception: + pass + try: + self.conn.close() + except Exception: + pass + self.conn = None + raise - def insert_into(self, table, row): - self.insert_batch([row]) \ No newline at end of file + def insert_into(self, table, row): + self.insert_batch([row])