diff --git a/pyrdp/mitm/ClipboardMITM.py b/pyrdp/mitm/ClipboardMITM.py index 10105c9d4..8b6d8779a 100644 --- a/pyrdp/mitm/ClipboardMITM.py +++ b/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 -