Give some structure to the logging module

3 files changed, 63 insertions, 46 deletions

diff --git a/libpathod/log.py b/libpathod/log.py
index ec21cc7b..69229adf 100644
--- a/libpathod/log.py
+++ b/libpathod/log.py
@@ -7,7 +7,7 @@ import netlib.http
 TIMEFMT = '%d-%m-%y %H:%M:%S'
 
 
-def write(fp, lines):
+def write_raw(fp, lines):
     if fp:
         fp.write(
             "%s: " % datetime.datetime.now().strftime(TIMEFMT)
@@ -18,7 +18,7 @@ def write(fp, lines):
         fp.flush()
 
 
-class Log(object):
+class LogCtx(object):
 
     def __init__(self, fp, hex, rfile, wfile):
         self.lines = []
@@ -46,7 +46,7 @@ class Log(object):
             self("Bytes read:")
             self.dump(rlog, self.hex)
         if self.lines:
-            write(
+            write_raw(
                 self.fp,
                 [
                     "\n".join(self.lines),
@@ -68,3 +68,16 @@ class Log(object):
 
     def __call__(self, line):
         self.lines.append(line)
+
+
+class ConnectionLogger:
+    def __init__(self, fp, hex, rfile, wfile):
+        self.fp = fp
+        self.hex = hex
+        self.rfile, self.wfile = rfile, wfile
+
+    def ctx(self):
+        return LogCtx(self.fp, self.hex, self.rfile, self.wfile)
+
+    def write(self, lines):
+        write_raw(self.fp, lines)
diff --git a/libpathod/pathoc.py b/libpathod/pathoc.py
index e4aadd40..4a2ac5cd 100644
--- a/libpathod/pathoc.py
+++ b/libpathod/pathoc.py
@@ -106,12 +106,10 @@ class WebsocketFrameReader(threading.Thread):
         self.rfile = rfile
         self.terminate = Queue.Queue()
         self.frames_queue = Queue.Queue()
-
-    def log(self, rfile):
-        return log.Log(
+        self.logger = log.ConnectionLogger(
             self.logfp,
             self.hexdump,
-            rfile if self.showresp else None,
+            rfile if showresp else None,
             None
         )
 
@@ -136,8 +134,11 @@ class WebsocketFrameReader(threading.Thread):
                 except Queue.Empty:
                     pass
                 for rfile in r:
-                    with self.log(rfile) as log:
-                        frm = websockets.Frame.from_file(self.rfile)
+                    with self.logger.ctx() as log:
+                        try:
+                            frm = websockets.Frame.from_file(self.rfile)
+                        except tcp.NetLibDisconnect:
+                            return
                         self.frames_queue.put(frm)
                         log("<< %s" % frm.header.human_readable())
                         if self.ws_read_limit is not None:
@@ -218,7 +219,7 @@ class Pathoc(tcp.TCPClient):
 
         if self.use_http2:
             if not OpenSSL._util.lib.Cryptography_HAS_ALPN:  # pragma: nocover
-                log.write(
+                log.write_raw(
                     self.fp,
                     "HTTP/2 requires ALPN support. "
                     "Please use OpenSSL >= 1.0.2. "
@@ -237,14 +238,6 @@ class Pathoc(tcp.TCPClient):
             protocol=self.protocol,
         )
 
-    def log(self):
-        return log.Log(
-            self.fp,
-            self.hexdump,
-            self.rfile if self.showresp else None,
-            self.wfile if self.showreq else None,
-        )
-
     def http_connect(self, connect_to):
         self.wfile.write(
             'CONNECT %s:%s HTTP/1.1\r\n' % tuple(connect_to) +
@@ -348,8 +341,14 @@ class Pathoc(tcp.TCPClient):
         """
             Sends a single websocket frame.
         """
-        with self.log() as log:
-            log(">> %s" % r)
+        logger = log.ConnectionLogger(
+            self.fp,
+            self.hexdump,
+            None,
+            self.wfile if self.showreq else None,
+        )
+        with logger.ctx() as lg:
+            lg(">> %s" % r)
             language.serve(r, self.wfile, self.settings)
             self.wfile.flush()
 
@@ -382,8 +381,14 @@ class Pathoc(tcp.TCPClient):
 
             May raise http.HTTPError, tcp.NetLibError
         """
-        with self.log() as log:
-            log(">> %s" % r)
+        logger = log.ConnectionLogger(
+            self.fp,
+            self.hexdump,
+            self.rfile if self.showresp else None,
+            self.wfile if self.showreq else None,
+        )
+        with logger.ctx() as lg:
+            lg(">> %s" % r)
             resp, req = None, None
             try:
                 req = language.serve(r, self.wfile, self.settings)
@@ -403,19 +408,19 @@ class Pathoc(tcp.TCPClient):
                     resp.append(self.sslinfo)
                     resp = Response(*resp)
             except http.HttpError as v:
-                log("Invalid server response: %s" % v)
+                lg("Invalid server response: %s" % v)
                 raise
             except tcp.NetLibTimeout:
                 if self.ignoretimeout:
-                    log("Timeout (ignored)")
+                    lg("Timeout (ignored)")
                     return None
-                log("Timeout")
+                lg("Timeout")
                 raise
             finally:
                 if resp:
-                    log(self._resp_summary(resp))
+                    lg(self._resp_summary(resp))
                     if resp.status_code in self.ignorecodes:
-                        log.suppress()
+                        lg.suppress()
             return resp
 
     def request(self, r):
diff --git a/libpathod/pathod.py b/libpathod/pathod.py
index e13a3f56..148fa7b2 100644
--- a/libpathod/pathod.py
+++ b/libpathod/pathod.py
@@ -87,7 +87,7 @@ class PathodHandler(tcp.BaseHandler):
     def _handle_sni(self, connection):
         self.sni = connection.get_servername()
 
-    def http_serve_crafted(self, crafted):
+    def http_serve_crafted(self, crafted, logctx):
         """
             This method is HTTP/1 and HTTP/2 capable.
         """
@@ -105,7 +105,7 @@ class PathodHandler(tcp.BaseHandler):
 
         if self.server.explain and not hasattr(crafted, 'is_error_response'):
             crafted = crafted.freeze(self.settings)
-            log.write(self.logfp, ">> Spec: %s" % crafted.spec())
+            logctx(">> Spec: %s" % crafted.spec())
 
         response_log = language.serve(
             crafted,
@@ -116,11 +116,9 @@ class PathodHandler(tcp.BaseHandler):
             return None, response_log
         return self.handle_http_request, response_log
 
-    def handle_websocket(self):
-        lr = self.rfile if self.server.logreq else None
-        lw = self.wfile if self.server.logresp else None
+    def handle_websocket(self, logger):
         while True:
-            with log.Log(self.logfp, self.server.hexdump, lr, lw) as lg:
+            with logger.ctx() as lg:
                 started = time.time()
                 try:
                     frm = websockets.Frame.from_file(self.rfile)
@@ -147,14 +145,13 @@ class PathodHandler(tcp.BaseHandler):
                 try:
                     wf_gen = language.parse_websocket_frame(nest)
                 except language.exceptions.ParseException as v:
-                    log.write(
-                        self.logfp,
+                    logger.write(
                         "Parse error in reflected frame specifcation:"
                         " %s" % v.msg
                     )
-                    break
+                    return None, None
                 for frm in wf_gen:
-                    with log.Log(self.logfp, self.server.hexdump, lr, lw) as lg:
+                    with logger.ctx() as lg:
                         frame_log = language.serve(
                             frm,
                             self.wfile,
@@ -223,7 +220,7 @@ class PathodHandler(tcp.BaseHandler):
         a.serve(flow, self.wfile)
         return self.handle_http_request, None
 
-    def handle_http_request(self):
+    def handle_http_request(self, logger):
         """
             This method is HTTP/1 and HTTP/2 capable.
 
@@ -232,9 +229,7 @@ class PathodHandler(tcp.BaseHandler):
             handler: Handler for the next request, or None to disconnect
             log: A dictionary, or None
         """
-        lr = self.rfile if self.server.logreq else None
-        lw = self.wfile if self.server.logresp else None
-        with log.Log(self.logfp, self.server.hexdump, lr, lw) as lg:
+        with logger.ctx() as lg:
             if self.use_http2:
                 self.protocol.perform_server_connection_preface()
                 stream_id, headers, body = self.protocol.read_request()
@@ -329,7 +324,8 @@ class PathodHandler(tcp.BaseHandler):
 
                 lg("crafting spec: %s" % spec)
                 nexthandler, retlog["response"] = self.http_serve_crafted(
-                    spec
+                    spec,
+                    lg
                 )
                 if nexthandler and websocket_key:
                     return self.handle_websocket, retlog
@@ -420,7 +416,7 @@ class PathodHandler(tcp.BaseHandler):
                         msg=s
                     )
                 )
-                log.write(self.logfp, s)
+                log.write_raw(self.logfp, s)
                 return
 
             alp = self.get_alpn_proto_negotiated()
@@ -433,13 +429,16 @@ class PathodHandler(tcp.BaseHandler):
         # if not self.protocol:
         #     # TODO: create HTTP or Websockets protocol
         #     self.protocol = None
+        lr = self.rfile if self.server.logreq else None
+        lw = self.wfile if self.server.logresp else None
+        logger = log.ConnectionLogger(self.logfp, self.server.hexdump, lr, lw)
 
         self.settings.protocol = self.protocol
 
         handler = self.handle_http_request
 
         while not self.finished:
-            handler, l = handler()
+            handler, l = handler(logger)
             if l:
                 self.addlog(l)
             if not handler:
@@ -552,7 +551,7 @@ class Pathod(tcp.TCPServer):
             h.handle()
             h.finish()
         except tcp.NetLibDisconnect:  # pragma: no cover
-            log.write(self.logfp, "Disconnect")
+            log.write_raw(self.logfp, "Disconnect")
             self.add_log(
                 dict(
                     type="error",
@@ -561,7 +560,7 @@ class Pathod(tcp.TCPServer):
             )
             return
         except tcp.NetLibTimeout:
-            log.write(self.logfp, "Timeout")
+            log.write_raw(self.logfp, "Timeout")
             self.add_log(
                 dict(
                     type="timeout",