feat(clipboard): Added timeout to clipboard file carving.

pyrdp/mitm/ClipboardMITM.py

@@ -6,6 +6,7 @@
 
 from logging import LoggerAdapter
 from io import BytesIO
+from functools import partial
 
 from pathlib import Path
 
@@ -18,6 +19,12 @@
 from pyrdp.recording import Recorder
 from pyrdp.mitm.config import MITMConfig
 
+from twisted.internet.interfaces import IDelayedCall
+from twisted.internet import reactor  # Import the current reactor.
+
+
+TRANSFER_TIMEOUT = 5  # delay in seconds after which to kill a stalled transfer.
+
 
 class PassiveClipboardStealer:
     """
@@ -41,6 +48,7 @@ def __init__(self, config: MITMConfig, client: ClipboardLayer, server: Clipboard
         self.forwardNextDataResponse = True
         self.files = []
         self.transfers = {}
+        self.timeouts = {}  # Track active timeout monitoring tasks.
 
         self.fileDir = f"{self.config.fileDir}/{self.log.sessionID}"
 
@@ -88,8 +96,15 @@ def handlePDU(self, pdu: ClipboardPDU, destination: ClipboardLayer):
             destination.sendPDU(pdu)
 
     def onFileContentsRequest(self, pdu: FileContentsRequestPDU):
+        """
+        There are two types of content requests: SIZE and RANGE.
+
+        A new transfer begins with a SIZE request and is followed by multiple
+        RANGE requests to retrieve the file data.
+
+        The file is picked from the advertised clipboard file list with an index.
+        """
         if pdu.flags == FileContentsFlags.SIZE:
-            # This is a new transfer request.
             if pdu.lindex < len(self.files):
                 fd = self.files[pdu.lindex]
                 self.log.info('Starting transfer for file "%s" ClipId=%d', fd.filename, pdu.clipId)
@@ -101,13 +116,18 @@ def onFileContentsRequest(self, pdu: FileContentsRequestPDU):
                 fpath.mkdir(parents=True, exist_ok=True)
 
                 self.transfers[pdu.streamId] = FileTransfer(fpath, fd, pdu.size)
+
+                # Track transfer timeout to prevent hung transfers.
+                cbTimeout = reactor.callLater(TRANSFER_TIMEOUT, partial(self.onTransferTimedOut, pdu.streamId))
+                self.timeouts[pdu.streamId] = cbTimeout
             else:
-                self.log.info('Request for uknown file! (lindex=%d)',  pdu.lindex)
+                self.log.info('Request for uknown file! (list index=%d)', pdu.lindex)
 
         elif pdu.flags == FileContentsFlags.RANGE:
             if pdu.streamId not in self.transfers:
                 self.log.warning('FileContentsRequest for unknown transfer (streamId=%d)', pdu.streamId)
             else:
+                self.refreshTimeout(pdu.streamId)
                 self.transfers[pdu.streamId].onRequest(pdu)
 
         return True
@@ -116,13 +136,35 @@ def onFileContentsResponse(self, pdu: FileContentsResponsePDU):
         if pdu.streamId not in self.transfers:
             self.log.warning('FileContentsResponse for unknown transfer (streamId=%d)', pdu.streamId)
         else:
+            self.refreshTimeout(pdu.streamId)
+
             done = self.transfers[pdu.streamId].onResponse(pdu)
             if done:
                 xfer = self.transfers[pdu.streamId]
                 self.log.info('Transfer completed for file "%s" location: "%s"', xfer.info.filename, xfer.localname)
                 del self.transfers[pdu.streamId]
+
+                # Remove the timeout since the transfer is done.
+                # This cannot throw because if we got this far, the delayed task cannot
+                # have been executed yet.
+                self.timeouts[pdu.streamId].cancel()
+                del self.timeouts[pdu.streamId]
+
         return True
 
+    def onTransferTimedOut(self, streamId: int):
+        if streamId in self.transfers:
+            # If the transfer exists, abort it. Otherwise, most likely the
+            # transfer has been completed. The latter should never happen due to the way
+            # twisted's reactor works.
+            xfer = self.transfers[streamId]
+            self.log.warn('Transfer timed out for "%s" (location: "%s")', xfer.info.filename, xfer.localname)
+            del self.transfers[streamId]
+            del self.timeouts[streamId]
+
+    def refreshTimeout(self, streamId: int):
+        self.timeouts[streamId].delay(TRANSFER_TIMEOUT)
+
     def onFormatDataResponse(self, pdu: FormatDataResponsePDU):
         if pdu.msgFlags == ClipboardMessageFlags.CB_RESPONSE_OK:
             # Keep the file list if there is one.
@@ -238,4 +280,3 @@ def onResponse(self, pdu: FileContentsResponsePDU) -> bool:
             return True
 
         return False
-