Merge pull request #244 from GoSecure/fix-222

fix(#222): Added device ID to I/O Requests.
GoSecure · Sep 14, 2020 · ba41c66 · ba41c66
2 parents 5b4223c + 38b2e64
commit ba41c66
Showing 1 changed file with 18 additions and 15 deletions.
diff --git a/pyrdp/mitm/DeviceRedirectionMITM.py b/pyrdp/mitm/DeviceRedirectionMITM.py
@@ -54,7 +54,6 @@ class DeviceRedirectionMITM(Subject):
 
     FORGED_COMPLETION_ID = 1000000
 
-
     def __init__(self, client: DeviceRedirectionLayer, server: DeviceRedirectionLayer, log: LoggerAdapter,
                  statCounter: StatCounter, state: RDPMITMState):
         """
@@ -70,12 +69,13 @@ def __init__(self, client: DeviceRedirectionLayer, server: DeviceRedirectionLaye
         self.state = state
         self.log = log
         self.statCounter = statCounter
-        self.currentIORequests: Dict[int, DeviceIORequestPDU] = {}
         self.openedFiles: Dict[int, FileProxy] = {}
         self.openedMappings: Dict[int, FileMapping] = {}
         self.fileMap: Dict[str, FileMapping] = {}
         self.fileMapPath = self.config.outDir / "mapping.json"
-        self.forgedRequests: Dict[int, DeviceRedirectionMITM.ForgedRequest] = {}
+
+        self.currentIORequests: Dict[(int, int), DeviceIORequestPDU] = {}
+        self.forgedRequests: Dict[(int, int), DeviceRedirectionMITM.ForgedRequest] = {}
 
         if self.config.extractFiles:
             self.responseHandlers: Dict[MajorFunction, callable] = {
@@ -155,7 +155,8 @@ def handleIORequest(self, pdu: DeviceIORequestPDU):
         """
 
         self.statCounter.increment(STAT.DEVICE_REDIRECTION_IOREQUEST)
-        self.currentIORequests[pdu.completionID] = pdu
+        key = (pdu.deviceID, pdu.completionID)
+        self.currentIORequests[key] = pdu
 
     def handleIOResponse(self, pdu: DeviceIOResponsePDU):
         """
@@ -164,16 +165,16 @@ def handleIOResponse(self, pdu: DeviceIOResponsePDU):
         """
 
         self.statCounter.increment(STAT.DEVICE_REDIRECTION_IORESPONSE)
-
-        if pdu.completionID in self.forgedRequests:
-            request = self.forgedRequests[pdu.completionID]
+        key = (pdu.deviceID, pdu.completionID)
+        if key in self.forgedRequests:
+            request = self.forgedRequests[key]
             request.handleResponse(pdu)
 
             if request.isComplete:
-                self.forgedRequests.pop(pdu.completionID)
+                self.forgedRequests.pop(key)
 
-        elif pdu.completionID in self.currentIORequests:
-            requestPDU = self.currentIORequests.pop(pdu.completionID)
+        elif key in self.currentIORequests:
+            requestPDU = self.currentIORequests.pop(key)
 
             if pdu.ioStatus >> 30 == IOOperationSeverity.STATUS_SEVERITY_ERROR:
                 self.statCounter.increment(STAT.DEVICE_REDIRECTION_IOERROR)
@@ -216,8 +217,9 @@ def handleCreateResponse(self, request: DeviceCreateRequestPDU, response: Device
             mapping = FileMapping.generate(remotePath, self.config.fileDir)
             proxy = FileProxy(mapping.localPath, "wb")
 
-            self.openedFiles[response.fileID] = proxy
-            self.openedMappings[response.fileID] = mapping
+            key = (response.deviceID, response.completionID, response.fileID)
+            self.openedFiles[key] = proxy
+            self.openedMappings[key] = mapping
 
             proxy.createObserver(
                 onFileCreated = lambda _: self.log.info("Saving file '%(remotePath)s' to '%(localPath)s'", {
@@ -335,7 +337,8 @@ def sendForgedFileRead(self, deviceID: int, path: str) -> int:
 
         completionID = self.findNextRequestID()
         request = DeviceRedirectionMITM.ForgedFileReadRequest(deviceID, completionID, self, path)
-        self.forgedRequests[completionID] = request
+        key = (deviceID, completionID)
+        self.forgedRequests[key] = request
 
         request.send()
         return completionID
@@ -358,13 +361,13 @@ def sendForgedDirectoryListing(self, deviceID: int, path: str) -> int:
 
         completionID = self.findNextRequestID()
         request = DeviceRedirectionMITM.ForgedDirectoryListingRequest(deviceID, completionID, self, path)
-        self.forgedRequests[completionID] = request
+        key = (deviceID, completionID)
+        self.forgedRequests[key] = request
 
         request.send()
         return completionID
 
 
-
     class ForgedRequest:
         """
         Base class for forged requests that simulate the server asking for information.