Made replay filenames easier to link with logs and vice versa

Fixes #180
GoSecure · Feb 5, 2020 · 490886f · 490886f
1 parent 53c5fdd
commit 490886f
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 5 deletions.
diff --git a/pyrdp/logging/StatCounter.py b/pyrdp/logging/StatCounter.py
@@ -4,8 +4,10 @@
 # Licensed under the GPLv3 or later.
 #
 
+import collections
 import time
 from logging import LoggerAdapter
+from typing import Optional
 
 
 class STAT:
@@ -154,10 +156,17 @@ def stop(self):
         if self.stats[STAT.TOTAL_OUTPUT] > 0:
             self.stats[STAT.CLIENT_SERVER_RATIO] = self.stats[STAT.TOTAL_INPUT] / self.stats[STAT.TOTAL_OUTPUT]
 
-    def logReport(self, log: LoggerAdapter):
+    def logReport(self, log: LoggerAdapter, more_info: Optional[collections.Mapping] = None):
         """
         Create an INFO log message to log the Connection report using the keys in self.stats.
         :param log: Logger to use to log the report
         """
-        keys = ", ".join([f"{key}: %({key})s" for key in self.stats.keys()])
-        log.info(f"Connection report: {keys}", self.stats)
+        # merge in the additional data if required
+        if isinstance(more_info, collections.Mapping):
+            report_data = {**self.stats, **more_info}
+        else:
+            report_data = self.stats
+
+        keys = ", ".join([f"{key}: %({key})s" for key in report_data.keys()])
+
+        log.info(f"Connection report: {keys}", report_data)
diff --git a/pyrdp/mitm/RDPMITM.py b/pyrdp/mitm/RDPMITM.py
@@ -131,7 +131,11 @@ def __init__(self, mainLogger: SessionLogger, crawlerLogger: SessionLogger, conf
 
         if config.recordReplays:
             date = datetime.datetime.now()
-            replayFileName = "rdp_replay_{}_{}.pyrdp".format(date.strftime('%Y%m%d_%H-%M-%S'), date.microsecond // 1000)
+            replayFileName = "rdp_replay_{}_{}_{}.pyrdp"\
+                    .format(date.strftime('%Y%m%d_%H-%M-%S'),
+                            date.microsecond // 1000,
+                            self.log.sessionID)
+            self.recorder.setRecordFilename(replayFileName)
             self.recorder.addTransport(FileLayer(self.config.replayDir / replayFileName))
 
         if config.enableCrawler:

diff --git a/pyrdp/mitm/TCPMITM.py b/pyrdp/mitm/TCPMITM.py
@@ -89,7 +89,11 @@ def onClientDisconnection(self, reason):
         self.statCounter.stop()
         self.recordConnectionClose()
         self.log.info("Client connection closed. %(reason)s", {"reason": reason.value})
-        self.statCounter.logReport(self.log)
+        if self.recorder.recordFilename:
+            self.statCounter.logReport(self.log, {"replayFilename":
+                                                  self.recorder.recordFilename})
+        else:
+            self.statCounter.logReport(self.log)
         self.serverConnector.close()
         self.server.disconnect(True)
 

diff --git a/pyrdp/recording/recorder.py b/pyrdp/recording/recorder.py
@@ -33,10 +33,18 @@ def __init__(self, transports: List[LayerChainItem]):
         }
 
         self.topLayers = []
+        self.recordFilename = None
 
         for transport in transports:
             self.addTransport(transport)
 
+    def setRecordFilename(self, filename: str):
+        """
+        Sets the filename used for the session recording.
+        :param filename: the filename
+        """
+        self.recordFilename = filename
+
     def addTransport(self, transportLayer: LayerChainItem):
         player = PlayerLayer()
         player.setPrevious(transportLayer)