Refactor and multi-thread

This commit is contained in:
2023-07-19 17:43:56 +10:00
parent 17622573e4
commit 8531294ca5
+257 -153
View File
@@ -1,190 +1,294 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
import znc import znc
import pymssql
import re import re
import multiprocessing
import warnings
import traceback
from datetime import datetime from datetime import datetime
from contextlib import closing from contextlib import closing
class mssql(znc.Module): class mssql(znc.Module):
module_types = [znc.CModInfo.GlobalModule] module_types = [znc.CModInfo.GlobalModule]
description = "Log all channels to a Microsoft SQL Server database" description = "Log all channels to a Microsoft SQL Server database"
has_args = False has_args = False
sql_server = "db.spd.internal" log_queue = multiprocessing.SimpleQueue()
sql_user = "_srv_znc" internal_log = None
sql_password = "zncznc" hook_debugging = True
sql_database = "znc"
# def OnLoad(self, args, message): def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network):
# self.PutModule("Loading MSSQL...") """
# return True 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): def OnLoad(self, args, message):
self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) """
return True 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): try:
self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) db = MSSQLDatabase({
return True '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): with self.internal_log.error() as target:
self.insert("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) target.write('Could not initialize module caused by: {} {}\n'.format(type(e), str(e)))
return True target.write('Stack trace: ' + traceback.format_exc())
target.write('\n')
def OnVoice(self, user, target_user, channel, noChange): return False
self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True
def OnDevoice(self, user, target_user, channel, noChange): def OnOp(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())) self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True return True
def OnMode(self, user, channel, mode, argument, added, noChange): def OnPrivNotice(self, user, message):
mode = chr(mode) self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
code = "MODE" return True
if mode == "b":
code = "BAN" if added else "UNBAN"
else:
sign = "+" if added else "-"
argument = "{0}{1} {2}".format(sign, mode, argument)
self.insert(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork())) def OnDeop(self, user, target_user, channel, noChange):
return True 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): def OnVoice(self, user, target_user, channel, noChange):
for channel in channels: self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
self.insert("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) return True
return True
def OnNick(self, user, new_nick, channels): def OnDevoice(self, user, target_user, channel, noChange):
for channel in channels: self.insert("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork())) return True
return True
def OnKick(self, user, target_nick, channel, message): def OnMode(self, user, channel, mode, argument, added, noChange):
self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork())) mode = chr(mode)
return True 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(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork()))
self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork())) return True
return True
def OnPart(self, user, channel, message): def OnQuit(self, user, message, channels):
self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) for channel in channels:
return True 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): def OnNick(self, user, new_nick, channels):
(channel, user) = self.resolveTarget(target) for channel in channels:
if channel == None: self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork()))
return True return True
return self.OnChanAction(user, channel, message)
def OnPrivMsg(self, nick, message): def OnKick(self, user, target_nick, channel, message):
self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork()))
return True return True
def OnUserMsg (self, target, message): def OnJoin(self, user, channel):
(channel, user) = self.resolveTarget(target) self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork()))
if channel == None: return True
return True
return self.OnChanMsg(user, channel, message)
def OnChanAction(self, user, channel, message): def OnPart(self, user, channel, message):
self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return True return True
def OnChanMsg(self, user, channel, message): def OnUserAction (self, target, message):
self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork())) (channel, user) = self.resolveTarget(target)
return True if channel == None:
return True
return self.OnChanAction(user, channel, message)
def OnTopic(self, user, channel, message): def OnPrivMsg(self, nick, message):
self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True return True
# def OnLoad(self, args, message): def OnUserMsg (self, target, message):
# match = re.search("(.*?):(.*?)@(.*)", args) (channel, user) = self.resolveTarget(target)
# if match: if channel == None:
# self.username = match.group(1) return True
# self.password = match.group(2) return self.OnChanMsg(user, channel, message)
# self.host = match.group(3)
# return True
# else:
# return False
def findMode(self, channel, user): def OnChanAction(self, user, channel, message):
realUser = channel.FindNick(user.GetNick()) self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
if not realUser: return True
return None
return str(realUser.GetPermChar())
def resolveTarget(self, target): def OnChanMsg(self, user, channel, message):
target = target.s self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork()))
channel = self.GetNetwork().FindChan(target) return True
if channel == None:
return (None, None)
user = self.GetNetwork().GetIRCNick()
return (channel, user)
def colourstrip(self, data): def OnTopic(self, user, channel, message):
find = data.find('\x03') self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
while find > -1: return True
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] == ','): def findMode(self, channel, user):
try: realUser = channel.FindNick(user.GetNick())
assert not done if not realUser:
assert int(data[find+1]) return None
data = data[0:find] + data[find+1:] return str(realUser.GetPermChar())
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') def resolveTarget(self, target):
data = data.replace('\x02','') target = target.s
data = data.replace('\x1d','') channel = self.GetNetwork().FindChan(target)
data = data.replace('\x1f','') if channel == None:
data = data.replace('\x16','') return (None, None)
data = data.replace('\x0f','') user = self.GetNetwork().GetIRCNick()
return data return (channel, user)
def insert(self, code, channel, host, user, ident, user_mode, date, target_user, message, network): def colourstrip(self, data):
# self.PutModule("Inserting row...") find = data.find('\x03')
regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE) 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: if (not done) and (len(data) > find+1) and (data[find] == ','):
with closing(pymssql.connect ("db.spd.internal", "_srv_znc", "znc", "znc")) as conn: try:
with closing(conn.cursor()) as cursor: assert not done
# 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)" assert int(data[find+1])
sql = "EXEC dbo.add_log (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" data = data[0:find] + data[find+1:]
if message is not None: data = data[0:find] + data[find+1:]
# message = message.encode('utf-8').decode('utf-8', 'replace') except:
message = regex.sub("", message) done = True
message = self.colourstrip(message) try:
znc_user = self.GetUser().GetUserName() if self.GetUser() is not None else None assert not done
params = (znc_user, str(date), code, network, channel, host, user, ident, user_mode, target_user, message) while True:
# cursor.execute (sql, params) assert int(data[find])
cursor.callproc("add_log", params) data = data[0:find] + data[find+1:]
conn.commit() except: pass
conn.close()
except Exception as e: find = data.find('\x03')
self.PutModule("Could not save {0} to database caused by: {1} - {2}".format(code, type(e), str(e))) 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()