#!/usr/bin/env python3 import znc import re import queue import threading import warnings import traceback import inspect import os import pprint import json 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 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. """ 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 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())) 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. """ if self.hook_debugging is not True: return 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 """ 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: {}'.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 ') 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('Failed to parse args: {} (args_str: "{}", type: {})\n'.format(e, args_str, type(args))) 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') 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') self.debug_hook() 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' } self.worker_threads = [] self.next_worker_id = 0 self._start_workers(self.num_workers) return True except Exception as e: message.s = str(e) 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') return False 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 if self.worker_threads: for _ in self.worker_threads: self.log_queue.put(None) self.log_queue.join() 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 __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_nowait(None) except Exception: 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 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 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(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) 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 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) 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') 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, 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.""" 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 def insert_into(self, table, row): self.insert_batch([row])