Merge pull request #333 from GoSecure/improved-file-handling

Improved file handling on disconnection

pyrdp/mitm/DeviceRedirectionMITM.py

@@ -204,12 +204,13 @@ def handleDeviceListAnnounceRequest(self, pdu: DeviceListAnnounceRequest):
 
     def handleCreateResponse(self, request: DeviceCreateRequestPDU, response: DeviceCreateResponsePDU):
         """
-        Prepare to intercept a file: create a FileProxy object, which will only create the file when we actually write
+        Prepare to intercept a file: create a FileMapping object, which will only create the file when we actually write
         to it. When listing a directory, Windows sends a lot of create requests without actually reading the files. We
-        use a FileProxy object to avoid creating a lot of empty files whenever a directory is listed.
+        verify the request to avoid creating a lot of empty files whenever a directory is listed.
         :param request: the device create request
         :param response: the device IO response to the request
         """
+        self.log.debug("Handling a DeviceCreateRequest. Path: '%(path)s'", {"path": request.path})
 
         isFileRead = request.desiredAccess & (FileAccessMask.GENERIC_READ | FileAccessMask.FILE_READ_DATA) != 0
         isDirectory = request.createOptions & CreateOption.FILE_NON_DIRECTORY_FILE == 0

pyrdp/mitm/FileMapping.py

@@ -88,8 +88,12 @@ def finalize(self):
 
     def onDisconnection(self, reason):
         if not self.file.closed:
+            self.log.info("Got disconnected with an active file mapping. Path: '%(path)s'. "
+                          "Non-zero partial file transfer is kept here: '%(dataPath)s'. Closing.",
+                          {"path": str(self.filesystemPath.relative_to(self.filesystemDir)), "dataPath": str(self.dataPath)})
             self.file.close()
-            Path(self.file.name).unlink(missing_ok=True)
+            if not self.written:
+                self.dataPath.unlink(missing_ok=True)
 
     @staticmethod
     def generate(remotePath: str, outDir: Path, filesystemDir: Path, log: LoggerAdapter):