Clipboard files deduplication

README.md

@@ -178,13 +178,16 @@ pyrdp_output/
 │   └── WinDev2108Eval.pem
 ├── files
 │   ├── e91c6a5eb3ca15df5a5cb4cf4ebb6f33b2d379a3a12d7d6de8c412d4323feb4c
+│   ├── b14b26b7d02c85e74ab4f0d847553b2fdfaf8bc616f7c3efcc4771aeddd55700
 ├── filesystems
-│   ├── Kimberly835337
+│   ├── romantic_kalam_8214773
 │   │   └── device1
-│   └── Stephen215343
+│   │   └── clipboard
+|           └── priv-esc.exe -> ../../../files/b14b26b7d02c85e74ab4f0d847553b2fdfaf8bc616f7c3efcc4771aeddd55700
+│   └── happy_stonebraker_1992243
 │       ├── device1
 │       └── device2
-|           └── Users/User/3D Objects/desktop.ini
+|           └── Users/User/3D Objects/desktop.ini -> ../../../../../../e91c6a5eb3ca15df5a5cb4cf4ebb6f33b2d379a3a12d7d6de8c412d4323feb4c
 ├── logs
 │   ├── crawl.json
 │   ├── crawl.log
@@ -195,8 +198,8 @@ pyrdp_output/
 │   ├── player.log
 │   └── ssl.log
 └── replays
-    ├── rdp_replay_20210826_12-15-33_512_Stephen215343.pyrdp
-    └── rdp_replay_20211125_12-55-42_352_Kimberly835337.pyrdp
+    ├── rdp_replay_20231214_01-20-28_965_happy_stonebraker_1992243.pyrdp
+    └── rdp_replay_20231214_00-42-24_295_romantic_kalam_8214773.pyrdp
 ```
 
 * `certs/` contains the certificates generated stored using the `CN` of the certificate as the file name

pyrdp/mitm/ClipboardMITM.py

@@ -19,12 +19,14 @@
 from pyrdp.parser.rdp.virtual_channel.clipboard import FileDescriptor
 from pyrdp.recording import Recorder
 from pyrdp.mitm.config import MITMConfig
+from pyrdp.mitm.FileMapping import FileMapping
 
-from twisted.internet.interfaces import IDelayedCall
 from twisted.internet import reactor  # Import the current reactor.
+from twisted.python.failure import Failure
 
 
 TRANSFER_TIMEOUT = 5  # delay in seconds after which to kill a stalled transfer.
+CLIPBOARD_FILEDIR = "clipboard" # special directory name under filesystems/<sessionId> for collected clipboard files
 
 
 class PassiveClipboardStealer:
@@ -52,7 +54,7 @@ def __init__(self, config: MITMConfig, client: ClipboardLayer, server: Clipboard
         self.transfers = {}
         self.timeouts = {}  # Track active timeout monitoring tasks.
 
-        self.fileDir = f"{self.config.fileDir}/{self.state.sessionID}"
+        self.filesystemRoot = self.config.filesystemDir / self.state.sessionID
 
         self.client.createObserver(
             onPDUReceived = self.onClientPDUReceived,
@@ -113,12 +115,10 @@ def onFileContentsRequest(self, pdu: FileContentsRequestPDU):
                               {"filename": fd.filename, "clipId": pdu.clipId})
 
                 if pdu.streamId in self.transfers:
-                    self.log.warning('File transfer already started')
+                    self.log.warning("Clipboard file transfer already started file '%(filename)s', clipId=%(clipId)d",
+                                     {"filename": fd.filename, "clipId": pdu.clipId})
 
-                fpath = Path(self.fileDir)
-                fpath.mkdir(parents=True, exist_ok=True)
-
-                self.transfers[pdu.streamId] = FileTransfer(fpath, fd, pdu.size)
+                self.transfers[pdu.streamId] = FileTransferMappingProxy(fd, self.config.fileDir, self.filesystemRoot, self.log, pdu.size)
 
                 # Track transfer timeout to prevent hung transfers.
                 cbTimeout = reactor.callLater(TRANSFER_TIMEOUT, partial(self.onTransferTimedOut, pdu.streamId))
@@ -147,8 +147,10 @@ def onFileContentsResponse(self, pdu: FileContentsResponsePDU):
             done = self.transfers[pdu.streamId].onResponse(pdu)
             if done:
                 xfer = self.transfers[pdu.streamId]
-                self.log.info("Transfer completed for file '%(filename)s', saved to: '%(localPath)s'",
-                              {"filename": xfer.info.filename, "localPath": xfer.localname})
+                self.log.info("Clipboard transfer completed for file '%(filename)s', saved as: '%(localPath)s', "
+                              "linked from: '%(linkPath)s'",
+                              {"filename": xfer.info.filename, "localPath": str(self.config.fileDir / xfer.getFileHash()),
+                               "linkPath": str(xfer.getFilesystemPath())})
                 del self.transfers[pdu.streamId]
 
                 # Remove the timeout since the transfer is done.
@@ -165,8 +167,9 @@ def onTransferTimedOut(self, streamId: int):
             # transfer has been completed. The latter should never happen due to the way
             # twisted's reactor works.
             xfer = self.transfers[streamId]
+            xfer.onDisconnection(Failure(Exception("Clipboard transfer timeout")))
             self.log.warn("Transfer timed out for '%(filename)s' saved to: '%(localPath)s'",
-                          {"filename": xfer.info.filename, "localPath": xfer.localname})
+                          {"filename": xfer.info.filename, "localPath": str(xfer.getDataPath())})
             del self.transfers[streamId]
             del self.timeouts[streamId]
 
@@ -236,27 +239,19 @@ def sendPasteRequest(self, destination: ClipboardLayer):
         destination.sendPDU(formatDataRequestPDU)
         self.forwardNextDataResponse = False
 
-
-class FileTransfer:
-    """Encapsulate the state of a clipboard file transfer."""
-    def __init__(self, dst: Path, info: FileDescriptor, size: int):
-        self.info = info
+class FileTransferMappingProxy():
+    """Encapsulate the state of a clipboard file transfer but proxies to FileMapping for storage and logging consistency"""
+    def __init__(self, fd: FileDescriptor, outDir: Path, filesystemSessionIdRoot: Path, log: LoggerAdapter, size: int):
+        self.info = fd
         self.size = size
         self.transferred: int = 0
-        self.data = b''
         self.prev = None  # Pending file content request.
 
-        self.localname = dst / Path(info.filename).name  # Avoid path traversal.
-
-        # Handle duplicates.
-        c = 1
-        localname = self.localname
-        while localname.exists():
-            localname = self.localname.parent / f'{self.localname.stem}_{c}{self.localname.suffix}'
-            c += 1
-        self.localname = localname
+        # We store files under filesystems/<sessionID>/ under a special clipboard directory
+        symlinkDst = filesystemSessionIdRoot / CLIPBOARD_FILEDIR
+        symlinkDst.mkdir(parents=True, exist_ok=True)
 
-        self.handle = open(str(self.localname), 'wb')
+        self.fileMapping = FileMapping.generate("/" + fd.filename, outDir, symlinkDst, log)
 
     def onRequest(self, pdu: FileContentsRequestPDU):
         # TODO: Handle out of order ranges. Are they even possible?
@@ -276,11 +271,20 @@ def onResponse(self, pdu: FileContentsResponsePDU) -> bool:
 
         received = len(pdu.data)
 
-        self.handle.write(pdu.data)
+        self.fileMapping.write(pdu.data)
         self.transferred += received
 
         if self.transferred == self.size:
-            self.handle.close()
+            self.fileMapping.finalize()
             return True
 
         return False
+
+    def getFileHash(self) -> str:
+        return self.fileMapping.fileHash
+
+    def getFilesystemPath(self) -> Path:
+        return self.fileMapping.filesystemPath
+
+    def onDisconnection(self, reason):
+        return self.fileMapping.onDisconnection(reason)

pyrdp/mitm/FileMapping.py

@@ -31,6 +31,8 @@ def __init__(self, file: io.BinaryIO, dataPath: Path, filesystemPath: Path, file
         self.filesystemDir = filesystemDir
         self.log = log
         self.written = False
+        # only available once finalized (since we hash to find the name we can't know ahead of time)
+        self.fileHash: str = None
 
     def seek(self, offset: int):
         if not self.file.closed:
@@ -40,7 +42,7 @@ def write(self, data: bytes):
         self.file.write(data)
         self.written = True
 
-    def getShaHash(self):
+    def _getShaHash(self):
         with open(self.dataPath, "rb") as f:
             # Note: In early 2022 we switched to sha256 for file hashes. If you
             #       want to use sha1, uncomment the next line and comment the
@@ -65,10 +67,10 @@ def finalize(self):
         self.log.debug("Closing file %(path)s", {"path": self.dataPath})
         self.file.close()
 
-        fileHash = self.getShaHash()
+        self.fileHash = self._getShaHash()
 
         # Go up one directory because files are saved to outDir / tmp while we're downloading them
-        hashPath = (self.dataPath.parents[1] / fileHash)
+        hashPath = (self.dataPath.parents[1] / self.fileHash)
 
         # Don't keep the file if we haven't written anything to it or it's a duplicate, otherwise rename and move to files dir
         if not self.written or hashPath.exists():
@@ -87,7 +89,7 @@ def finalize(self):
             self.filesystemPath.symlink_to(Path(os.path.relpath(hashPath, self.filesystemPath.parent)))
 
             self.log.info("SHA-256 '%(path)s' = '%(shasum)s'", {
-                "path": str(self.filesystemPath.relative_to(self.filesystemDir)), "shasum": fileHash
+                "path": str(self.filesystemPath.relative_to(self.filesystemDir)), "shasum": self.fileHash
             })
 
     def onDisconnection(self, reason):

test/test_FileMapping.py

@@ -22,7 +22,7 @@ def setUp(self):
     def createMapping(self, mkdir: MagicMock, mkstemp: MagicMock, mock_open_object):
         mkstemp.return_value = (1, str(self.outDir / "tmp" / "tmp_test"))
         mapping = FileMapping.generate("/test", self.outDir, Path("filesystems"), self.log)
-        mapping.getShaHash = Mock(return_value = self.hash)
+        mapping._getShaHash = Mock(return_value = self.hash)
         mapping.file.closed = False
         return mapping, mkdir, mkstemp, mock_open_object