Add a statcounter that prints statistics about the connection

pyrdp/logging/StatCounter.py

@@ -0,0 +1,157 @@
+#
+# This file is part of the PyRDP project.
+# Copyright (C) 2019 GoSecure Inc.
+# Licensed under the GPLv3 or later.
+#
+
+import time
+from logging import LoggerAdapter
+
+
+class STAT:
+    """
+    Type of statistics that a StatCounter object can hold.
+    """
+
+    CONNECTION_TIME = "connectionTime"
+    # Duration (in secs) for the TCP connection
+
+    CLIENT_SERVER_RATIO = "clientServerRatio"
+    # Ratio of the # of messages coming from the client vs from the server. High value (>1) means high client interaction.
+
+    TOTAL_INPUT = "totalInput"
+    # # of messages coming from the client to the server.
+
+    TOTAL_OUTPUT = "totalOutput"
+    # # of messages coming from the server to the client.
+
+    IO_INPUT = "input"
+    # Packet coming from the client to the server for the io channel
+
+    IO_INPUT_FASTPATH = "fastPathInput"
+    # Packet coming from the client to the server for the io channel as a fastpath packet
+
+    IO_INPUT_SLOWPATH = "slowPathInput"
+    # Packet coming from the client to the server for the io channel as a slowpath packet
+
+    IO_OUTPUT = "output"
+    # Packet coming from the server to the client for the io channel
+
+    IO_OUTPUT_FASTPATH = "fastPathOutput"
+    # Packet coming from the server to the client for the io channel as a fastpath packet
+
+    IO_OUTPUT_SLOWPATH = "slowPathOutput"
+    # Packet coming from the server to the client for the io channel as a slowpath packet
+
+    MCS = "mcs"
+    # Packet Coming from either end for any channel
+
+    MCS_OUTPUT = "mcsOutput"
+    # Packet Coming from the server to the client for any channel
+
+    MCS_OUTPUT_ = "mcsOutput_"
+    # Packet Coming from the server to the client for a given channel (must append channel # after it)
+
+    MCS_INPUT = "mcsInput"
+    # Packet Coming from the client to the server for any channel
+
+    MCS_INPUT_ = "mcsInput_"
+    # Packet Coming from the client to the server for a given channel (must append channel # after it)
+
+    VIRTUAL_CHANNEL = "virtualChannel"
+    # Packet Coming from either end for any virtual channel that doesnt have a specific implementation (ex clipboard)
+
+    VIRTUAL_CHANNEL_INPUT = "virtualChannelInput"
+    # Packet Coming from the client to the server for any virtual channel that doesnt have a specific implementation (ex clipboard)
+
+    VIRTUAL_CHANNEL_OUTPUT = "virtualChannelOutput"
+    # Packet Coming from the server to the client for any virtual channel that doesnt have a specific implementation (ex clipboard)
+
+    DEVICE_REDIRECTION = "deviceRedirection"
+    # Packet coming from either end for the rdpdr channel
+
+    DEVICE_REDIRECTION_CLIENT = "deviceRedirectionClient"
+    # Packet coming from the client to the server for the rdpdr channel
+
+    DEVICE_REDIRECTION_SERVER = "deviceRedirectionServer"
+    # Packet coming from the server to the client for the rdpdr channel
+
+    DEVICE_REDIRECTION_IOREQUEST = "deviceRedirectionIORequest"
+    # IORequest packets for the rdpdr channel
+
+    DEVICE_REDIRECTION_IORESPONSE = "deviceRedirectionIOResponse"
+    # IOResponse packets for the rdpdr channel
+
+    DEVICE_REDIRECTION_IOERROR = "deviceRedirectionIOError"
+    # IO error packets for the rdpdr channel
+
+    DEVICE_REDIRECTION_FILE_CLOSE = "deviceRedirectionFileClose"
+    # File Close packets for the rdpdr channel
+
+    DEVICE_REDIRECTION_FORGED_FILE_READ = "deviceRedirectionForgedFileRead"
+    # File read packets forged by pyrdp for the rdpdr channel
+
+    DEVICE_REDIRECTION_FORGED_DIRECTORY_LISTING = "deviceRedirectionForgedDirectoryListing"
+    # Directory listing packets forged by pyrdp for the rdpdr channel
+
+    CLIPBOARD = "clipboard"
+    # Number of clipboard PDUs coming from either end
+
+    CLIPBOARD_CLIENT = "clipboardClient"
+    # Number of clipboard PDUs coming from the client
+
+    CLIPBOARD_SERVER = "clipboardServer"
+    # Number of clipboard PDUs coming from the server
+
+    CLIPBOARD_COPY = "clipboardCopies"
+    # Number of times data has been copied by either end
+
+    CLIPBOARD_PASTE = "clipboardPastes"
+    # Number of times data has been pasted by either end
+
+
+class StatCounter:
+    """
+    Class that keeps track of various statistics during an RDP connection (See STAT)
+    """
+
+    def __init__(self):
+        self.stats = {"report": 1.0}
+
+    def increment(self, *args: str):
+        """
+        Increments all statistics passed in arguments
+        :param args: list of statistics to increment by one. See STAT for list of allowed values.
+        """
+        for stat in args:
+            if stat not in self.stats:
+                self.stats[stat] = 0
+            self.stats[stat] += 1
+
+    def incrementWith(self, statDestination: str, *statsSource: str):
+        """
+        Increments statDestination by all provided statSources
+        """
+        if statDestination not in self.stats:
+            self.stats[statDestination] = 0
+        for statSource in statsSource:
+            if statSource in self.stats:
+                self.stats[statDestination] += self.stats[statSource]
+
+    def start(self):
+        """
+        Initialize some statistics such as connectionTime
+        """
+        self.stats[STAT.CONNECTION_TIME] = time.time()
+
+    def stop(self):
+        """
+        Calculates the last statistics such as interaction ratio and connectionTime
+        """
+        self.stats[STAT.CONNECTION_TIME] = time.time() - self.stats[STAT.CONNECTION_TIME]
+        self.incrementWith(STAT.TOTAL_INPUT, STAT.MCS_INPUT, STAT.IO_INPUT_FASTPATH, STAT.VIRTUAL_CHANNEL_INPUT, STAT.CLIPBOARD_CLIENT, STAT.DEVICE_REDIRECTION_CLIENT)
+        self.incrementWith(STAT.TOTAL_OUTPUT, STAT.MCS_OUTPUT, STAT.IO_OUTPUT_FASTPATH, STAT.VIRTUAL_CHANNEL_OUTPUT, STAT.CLIPBOARD_SERVER, STAT.DEVICE_REDIRECTION_SERVER)
+        self.stats[STAT.CLIENT_SERVER_RATIO] = self.stats[STAT.TOTAL_INPUT] / self.stats[STAT.TOTAL_OUTPUT]
+
+    def logReport(self, log: LoggerAdapter):
+        log.info(self.stats)

pyrdp/mitm/ClipboardMITM.py

@@ -9,6 +9,7 @@
 from pyrdp.core import decodeUTF16LE
 from pyrdp.enum import ClipboardFormatNumber, ClipboardMessageFlags, ClipboardMessageType, PlayerPDUType
 from pyrdp.layer import ClipboardLayer
+from pyrdp.logging.StatCounter import StatCounter, STAT
 from pyrdp.pdu import ClipboardPDU, FormatDataRequestPDU, FormatDataResponsePDU
 from pyrdp.recording import Recorder
 
@@ -18,13 +19,15 @@ class PassiveClipboardStealer:
     MITM component for the clipboard layer. Logs clipboard data when it is pasted.
     """
 
-    def __init__(self, client: ClipboardLayer, server: ClipboardLayer, log: LoggerAdapter, recorder: Recorder):
+    def __init__(self, client: ClipboardLayer, server: ClipboardLayer, log: LoggerAdapter, recorder: Recorder,
+                 statCounter: StatCounter):
         """
         :param client: clipboard layer for the client side
         :param server: clipboard layer for the server side
         :param log: logger for this component
         :param recorder: recorder for clipboard data
         """
+        self.statCounter = statCounter
         self.client = client
         self.server = server
         self.log = log
@@ -40,9 +43,11 @@ def __init__(self, client: ClipboardLayer, server: ClipboardLayer, log: LoggerAd
         )
 
     def onClientPDUReceived(self, pdu: ClipboardPDU):
+        self.statCounter.increment(STAT.CLIPBOARD, STAT.CLIPBOARD_CLIENT)
         self.handlePDU(pdu, self.server)
 
     def onServerPDUReceived(self, pdu: ClipboardPDU):
+        self.statCounter.increment(STAT.CLIPBOARD, STAT.CLIPBOARD_SERVER)
         self.handlePDU(pdu, self.client)
 
     def handlePDU(self, pdu: ClipboardPDU, destination: ClipboardLayer):
@@ -63,6 +68,10 @@ def handlePDU(self, pdu: ClipboardPDU, destination: ClipboardLayer):
                 self.log.info("Clipboard data: %(clipboardData)r", {"clipboardData": clipboardData})
                 self.recorder.record(pdu, PlayerPDUType.CLIPBOARD_DATA)
 
+                if self.forwardNextDataResponse:
+                    # Means it's NOT a crafted response
+                    self.statCounter.increment(STAT.CLIPBOARD_PASTE)
+
             self.forwardNextDataResponse = True
 
     def decodeClipboardData(self, data: bytes) -> str:
@@ -79,8 +88,9 @@ class ActiveClipboardStealer(PassiveClipboardStealer):
     clipboard is updated.
     """
 
-    def __init__(self, client: ClipboardLayer, server: ClipboardLayer, log: LoggerAdapter, recorder: Recorder):
-        super().__init__(client, server, log, recorder)
+    def __init__(self, client: ClipboardLayer, server: ClipboardLayer, log: LoggerAdapter, recorder: Recorder,
+                 statCounter: StatCounter):
+        super().__init__(client, server, log, recorder, statCounter)
 
     def handlePDU(self, pdu: ClipboardPDU, destination: ClipboardLayer):
         """
@@ -99,6 +109,8 @@ def sendPasteRequest(self, destination: ClipboardLayer):
         Sets forwardNextDataResponse to False to make sure that this request is not actually transferred to the other end.
         """
 
+        self.statCounter.increment(STAT.CLIPBOARD_COPY)
+
         formatDataRequestPDU = FormatDataRequestPDU(ClipboardFormatNumber.GENERIC)
         destination.sendPDU(formatDataRequestPDU)
-        self.forwardNextDataResponse = False
+        self.forwardNextDataResponse = False

pyrdp/mitm/DeviceRedirectionMITM.py

@@ -15,6 +15,7 @@
     FileCreateDisposition, FileCreateOptions, FileShareAccess, FileSystemInformationClass, IOOperationSeverity, \
     MajorFunction, MinorFunction
 from pyrdp.layer import DeviceRedirectionLayer
+from pyrdp.logging.StatCounter import StatCounter, STAT
 from pyrdp.mitm.config import MITMConfig
 from pyrdp.mitm.FileMapping import FileMapping, FileMappingDecoder, FileMappingEncoder
 from pyrdp.pdu import DeviceAnnounce, DeviceCloseRequestPDU, DeviceCloseResponsePDU, DeviceCreateRequestPDU, \
@@ -53,7 +54,8 @@ class DeviceRedirectionMITM(Subject):
     FORGED_COMPLETION_ID = 1000000
 
 
-    def __init__(self, client: DeviceRedirectionLayer, server: DeviceRedirectionLayer, log: LoggerAdapter, config: MITMConfig):
+    def __init__(self, client: DeviceRedirectionLayer, server: DeviceRedirectionLayer, log: LoggerAdapter,
+                 config: MITMConfig, statCounter: StatCounter):
         """
         :param client: device redirection layer for the client side
         :param server: device redirection layer for the server side
@@ -65,6 +67,7 @@ def __init__(self, client: DeviceRedirectionLayer, server: DeviceRedirectionLaye
         self.client = client
         self.server = server
         self.log = log
+        self.statCounter = statCounter
         self.config = config
         self.currentIORequests: Dict[int, DeviceIORequestPDU] = {}
         self.openedFiles: Dict[int, FileProxy] = {}
@@ -106,9 +109,11 @@ def saveMapping(self):
             f.write(json.dumps(self.fileMap, cls=FileMappingEncoder, indent=4, sort_keys=True))
 
     def onClientPDUReceived(self, pdu: DeviceRedirectionPDU):
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION, STAT.DEVICE_REDIRECTION_CLIENT)
         self.handlePDU(pdu, self.server)
 
     def onServerPDUReceived(self, pdu: DeviceRedirectionPDU):
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION, STAT.DEVICE_REDIRECTION_SERVER)
         self.handlePDU(pdu, self.client)
 
     def handlePDU(self, pdu: DeviceRedirectionPDU, destination: DeviceRedirectionLayer):
@@ -137,6 +142,7 @@ def handleIORequest(self, pdu: DeviceIORequestPDU):
         :param pdu: the device IO request
         """
 
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION_IOREQUEST)
         self.currentIORequests[pdu.completionID] = pdu
 
     def handleIOResponse(self, pdu: DeviceIOResponsePDU):
@@ -145,6 +151,8 @@ def handleIOResponse(self, pdu: DeviceIOResponsePDU):
         :param pdu: the device IO response.
         """
 
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION_IORESPONSE)
+
         if pdu.completionID in self.forgedRequests:
             request = self.forgedRequests[pdu.completionID]
             request.handleResponse(pdu)
@@ -156,6 +164,7 @@ def handleIOResponse(self, pdu: DeviceIOResponsePDU):
             requestPDU = self.currentIORequests.pop(pdu.completionID)
 
             if pdu.ioStatus >> 30 == IOOperationSeverity.STATUS_SEVERITY_ERROR:
+                self.statCounter.increment(STAT.DEVICE_REDIRECTION_IOERROR)
                 self.log.warning("Received an IO Response with an error IO status: %(responsePDU)s for request %(requestPDU)s", {"responsePDU": repr(pdu), "requestPDU": repr(requestPDU)})
 
             if pdu.majorFunction in self.responseHandlers:
@@ -234,6 +243,8 @@ def handleCloseResponse(self, request: DeviceCloseRequestPDU, _: DeviceCloseResp
         :param _: the device IO response to the request
         """
 
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION_FILE_CLOSE)
+
         if request.fileID in self.openedFiles:
             file = self.openedFiles.pop(request.fileID)
             file.close()
@@ -292,6 +303,8 @@ def sendForgedFileRead(self, deviceID: int, path: str) -> int:
         :param path: path of the file to download. The path should use '\' instead of '/' to separate directories.
         """
 
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION_FORGED_FILE_READ)
+
         completionID = self.findNextRequestID()
         request = DeviceRedirectionMITM.ForgedFileReadRequest(deviceID, completionID, self, path)
         self.forgedRequests[completionID] = request
@@ -310,6 +323,8 @@ def sendForgedDirectoryListing(self, deviceID: int, path: str) -> int:
         \Documents\*
         """
 
+        self.statCounter.increment(STAT.DEVICE_REDIRECTION_FORGED_DIRECTORY_LISTING)
+
         completionID = self.findNextRequestID()
         request = DeviceRedirectionMITM.ForgedDirectoryListingRequest(deviceID, completionID, self, path)
         self.forgedRequests[completionID] = request

pyrdp/mitm/FastPathMITM.py

@@ -5,6 +5,7 @@
 #
 
 from pyrdp.layer import FastPathLayer
+from pyrdp.logging.StatCounter import StatCounter, STAT
 from pyrdp.mitm.state import RDPMITMState
 from pyrdp.pdu import FastPathPDU
 
@@ -14,29 +15,32 @@ class FastPathMITM:
     MITM component for the fast-path layer.
     """
 
-    def __init__(self, client: FastPathLayer, server: FastPathLayer, state: RDPMITMState):
+    def __init__(self, client: FastPathLayer, server: FastPathLayer, state: RDPMITMState, statCounter: StatCounter):
         """
         :param client: fast-path layer for the client side
         :param server: fast-path layer for the server side
         :param state: the MITM state.
         """
 
+        self.statCounter = statCounter
         self.client = client
         self.server = server
         self.state = state
 
         self.client.createObserver(
-            onPDUReceived = self.onClientPDUReceived,
+            onPDUReceived=self.onClientPDUReceived,
         )
 
         self.server.createObserver(
-            onPDUReceived = self.onServerPDUReceived,
+            onPDUReceived=self.onServerPDUReceived,
         )
 
     def onClientPDUReceived(self, pdu: FastPathPDU):
+        self.statCounter.increment(STAT.IO_INPUT_FASTPATH)
         if self.state.forwardInput:
             self.server.sendPDU(pdu)
 
     def onServerPDUReceived(self, pdu: FastPathPDU):
+        self.statCounter.increment(STAT.IO_OUTPUT_FASTPATH)
         if self.state.forwardOutput:
-            self.client.sendPDU(pdu)
+            self.client.sendPDU(pdu)