From e442cb677a73a1f1bb70b7110b068b710a0f1f41 Mon Sep 17 00:00:00 2001
From: ed <s@ocv.me>
Date: Sun, 6 Nov 2022 13:30:16 +0000
Subject: [PATCH] improve ftp/smb logging

---
 copyparty/ftpd.py   |  5 ++---
 copyparty/smbd.py   | 40 +++++-----------------------------------
 copyparty/svchub.py |  7 +++++++
 copyparty/util.py   | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 53 insertions(+), 38 deletions(-)

diff --git a/copyparty/ftpd.py b/copyparty/ftpd.py
index 6a9d64fd..133a91e6 100644
--- a/copyparty/ftpd.py
+++ b/copyparty/ftpd.py
@@ -11,7 +11,6 @@ import time
 from pyftpdlib.authorizers import AuthenticationFailed, DummyAuthorizer
 from pyftpdlib.filesystems import AbstractedFS, FilesystemError
 from pyftpdlib.handlers import FTPHandler
-from pyftpdlib.log import config_logging
 from pyftpdlib.servers import FTPServer
 
 from .__init__ import PY2, TYPE_CHECKING, E
@@ -401,8 +400,8 @@ class Ftpd(object):
             if self.args.ftp_nat:
                 h2.masquerade_address = self.args.ftp_nat
 
-        if self.args.ftp_dbg:
-            config_logging(level=logging.DEBUG)
+        lgr = logging.getLogger("pyftpdlib")
+        lgr.setLevel(logging.DEBUG if self.args.ftp_dbg else logging.INFO)
 
         ioloop = IOLoop()
         for ip in self.args.i:
diff --git a/copyparty/smbd.py b/copyparty/smbd.py
index 7e0ec42c..4e6f8746 100644
--- a/copyparty/smbd.py
+++ b/copyparty/smbd.py
@@ -17,51 +17,21 @@ from .util import Daemon, min_ex
 if True:  # pylint: disable=using-constant-test
     from typing import Any
 
-    from .util import RootLogger
-
 if TYPE_CHECKING:
     from .svchub import SvcHub
 
 
-class HLog(logging.Handler):
-    def __init__(self, log_func: "RootLogger") -> None:
-        logging.Handler.__init__(self)
-        self.log_func = log_func
-
-    def __repr__(self) -> str:
-        level = logging.getLevelName(self.level)
-        return "<%s cpp(%s)>" % (self.__class__.__name__, level)
-
-    def flush(self) -> None:
-        pass
-
-    def emit(self, record: logging.LogRecord) -> None:
-        msg = self.format(record)
-        lv = record.levelno
-        if lv < logging.INFO:
-            c = 6
-        elif lv < logging.WARNING:
-            c = 0
-        elif lv < logging.ERROR:
-            c = 3
-        else:
-            c = 1
-
-        self.log_func("smb", msg, c)
-
-
 class SMB(object):
     def __init__(self, hub: "SvcHub") -> None:
         self.hub = hub
         self.args = hub.args
         self.asrv = hub.asrv
-        self.log_func = hub.log
+        self.log = hub.log
         self.files: dict[int, tuple[float, str]] = {}
 
-        handler = HLog(hub.log)
-        lvl = logging.DEBUG if self.args.smb_dbg else logging.INFO
-        logging.getLogger().addHandler(handler)
-        logging.getLogger().setLevel(lvl)
+        for x in ["impacket", "impacket.smbserver"]:
+            lgr = logging.getLogger(x)
+            lgr.setLevel(logging.DEBUG if self.args.smb_dbg else logging.INFO)
 
         try:
             from impacket import smbserver
@@ -125,7 +95,7 @@ class SMB(object):
 
         self.srv = srv
         self.stop = srv.stop
-        logging.info("listening @ %s:%s", ip, port)
+        self.log("smb", "listening @ {}:{}".format(ip, port))
 
     def start(self) -> None:
         Daemon(self.srv.start)
diff --git a/copyparty/svchub.py b/copyparty/svchub.py
index 23cf7c4e..c68c5b1d 100644
--- a/copyparty/svchub.py
+++ b/copyparty/svchub.py
@@ -8,6 +8,7 @@ import argparse
 import base64
 import calendar
 import gzip
+import logging
 import os
 import re
 import shlex
@@ -34,6 +35,7 @@ from .up2k import Up2k
 from .util import (
     VERSIONS,
     Daemon,
+    HLog,
     HMaccas,
     alltrace,
     ansi_re,
@@ -106,6 +108,11 @@ class SvcHub(object):
         if args.lo:
             self._setup_logfile(printed)
 
+        lg = logging.getLogger()
+        lh = HLog(self.log)
+        lg.handlers = [lh]
+        lg.setLevel(logging.INFO)
+
         if args.stackmon:
             start_stackmon(args.stackmon, 0)
 
diff --git a/copyparty/util.py b/copyparty/util.py
index 76c1bdda..64309165 100644
--- a/copyparty/util.py
+++ b/copyparty/util.py
@@ -6,6 +6,7 @@ import contextlib
 import errno
 import hashlib
 import hmac
+import logging
 import math
 import mimetypes
 import os
@@ -381,6 +382,44 @@ class Cooldown(object):
             return ret
 
 
+class HLog(logging.Handler):
+    def __init__(self, log_func: "RootLogger") -> None:
+        logging.Handler.__init__(self)
+        self.log_func = log_func
+        self.ptn_ftp = re.compile(r"^([0-9a-f:\.]+:[0-9]{1,5})-\[")
+        self.ptn_smb_ign = re.compile(r"^(Callback added|Config file parsed)")
+
+    def __repr__(self) -> str:
+        level = logging.getLevelName(self.level)
+        return "<%s cpp(%s)>" % (self.__class__.__name__, level)
+
+    def flush(self) -> None:
+        pass
+
+    def emit(self, record: logging.LogRecord) -> None:
+        msg = self.format(record)
+        lv = record.levelno
+        if lv < logging.INFO:
+            c = 6
+        elif lv < logging.WARNING:
+            c = 0
+        elif lv < logging.ERROR:
+            c = 3
+        else:
+            c = 1
+
+        if record.name == "pyftpdlib":
+            m = self.ptn_ftp.match(msg)
+            if m:
+                record.name = ip = m.group(1)
+                msg = msg[len(ip) + 1 :]
+        elif record.name.startswith("impacket"):
+            if self.ptn_smb_ign.match(msg):
+                return
+
+        self.log_func(record.name[-21:], msg, c)
+
+
 class UnrecvEOF(OSError):
     pass