From f7a6d9376775a5c243d1ec1ba3f0ebb44a7d99e2 Mon Sep 17 00:00:00 2001
From: Andrew Walker <awalker@ixsystems.com>
Date: Tue, 24 Dec 2024 09:18:06 -0600
Subject: [PATCH] Switch to using middleware logger class

---
 scripts/truenas_audit_handler.py | 47 +++-----------------------------
 1 file changed, 4 insertions(+), 43 deletions(-)

diff --git a/scripts/truenas_audit_handler.py b/scripts/truenas_audit_handler.py
index 200976b..49dbee1 100644
--- a/scripts/truenas_audit_handler.py
+++ b/scripts/truenas_audit_handler.py
@@ -13,6 +13,7 @@
 from datetime import datetime
 from collections import defaultdict, deque
 from json import dumps
+from middlewared.logger import TNSyslogHandler
 from queue import Queue
 from random import getrandbits
 from uuid import UUID
@@ -30,53 +31,12 @@
 SYSLOG_IDENT = 'TNAUDIT_SYSTEM: '
 AUDITD_LINE_SEPARATOR = '\x1d'
 JSON_NULL = 'null'
+
 # TODO: generate critical middleware alert if our backlog starts to hit
 # critical levels
 ALERT_QUEUE_DEPTH = 1024
 
 
-class TNAuditHandler(logging.handlers.SysLogHandler):
-    """
-    A logging handler that's never gonna give you up
-
-    Parent class is responsible for maintaining socket via emit()
-    and friends
-    """
-    def __init__(self, address: str, pending_queue: deque):
-        self.pending_queue = pending_queue
-        super().__init__(address, socktype=socket.SOCK_STREAM)
-
-    def drain_pending_queue(self):
-        while self.pending_queue:
-            record = self.pending_queue.popleft()
-            try:
-                super().emit(record)
-            except Exception:
-                # Nope. Still dead. Put it back where we found it
-                self.pending_queue.appendleft(record)
-                return False
-
-        return True
-
-    def emit(self, record):
-        if not self.drain_pending_queue():
-            # Failed to drain our pending queue so add this record to the
-            # ever-growing deque
-            self.pending_queue.append(record)
-            return
-
-        try:
-            super().emit(record)
-        except Exception:
-            # logging framework done broke. Queue up
-            # for drain on next auditd message handled
-            self.pending_queue.append(record)
-
-    def handleError(self, record):
-        # re-raise it back up to the emit call
-        raise
-
-
 class AuditMsgParser(enum.Enum):
     @property
     def idx(self) -> int:
@@ -400,7 +360,8 @@ def __setup_logger(self) -> logging.Logger:
         # Set up logging queue to make sending messages to syslog nonblocking
         logq = Queue()
         queue_handler = logging.handlers.QueueHandler(logq)
-        audit_handler = TNAuditHandler(self.syslog_path, self.pending_queue)
+        queue_handler.setLevel(logging.DEBUG)
+        audit_handler = TNSyslogHandler(self.syslog_path, self.pending_queue)
         audit_handler.setLevel(logging.DEBUG)
         audit_handler.ident = SYSLOG_IDENT