[WIP] Seek using thumbnails

bin/pyrdp-player.py

@@ -54,6 +54,10 @@ def main():
     parser.add_argument("-b", "--bind", help="Bind address (default: 127.0.0.1)", default="127.0.0.1")
     parser.add_argument("-p", "--port", help="Bind port (default: 3000)", default=3000)
     parser.add_argument("-o", "--output", help="Output folder", default="pyrdp_output")
+    parser.add_argument("-t", "--thumbnails-directory", help="Directory where the thumbnails for "
+                                                            "the current replay files are located."
+                                                            " they are obtained using the -f png "
+                                                            "argument of pyrdp-convert.")
     parser.add_argument("-L", "--log-level", help="Log level", default=None,
                         choices=["INFO", "DEBUG", "WARNING", "ERROR", "CRITICAL"], nargs="?")
     parser.add_argument("-F", "--log-filter",
@@ -87,7 +91,8 @@ def main():
 
     if not args.headless:
         app = QApplication(sys.argv)
-        mainWindow = MainWindow(args.bind, int(args.port), args.replay)
+        mainWindow = MainWindow(args.bind, int(args.port), args.replay,
+                                thumbnails_directory=args.thumbnails_directory)
         mainWindow.showMaximized()
         mainWindow.show()
 

pyrdp/core/ssl.py

@@ -11,7 +11,6 @@
 
 import OpenSSL
 from OpenSSL import SSL
-
 from twisted.internet import ssl
 
 
@@ -87,7 +86,10 @@ def clone(self, cert: OpenSSL.crypto.X509) -> (OpenSSL.crypto.PKey, OpenSSL.cryp
     def lookup(self, cert: OpenSSL.crypto.X509) -> (str, str):
         subject = cert.get_subject()
         parts = dict(subject.get_components())
-        commonName = parts[b'CN'].decode()
+        try:
+            commonName = parts[b'CN'].decode('utf-8')
+        except UnicodeDecodeError:
+            commonName = parts[b'CN'].decode('utf-16be')
         base = str(self._root / commonName)
 
         if path.exists(base + '.pem'):

pyrdp/player/MainWindow.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2018-2020 GoSecure Inc.
 # Licensed under the GPLv3 or later.
 #
-from typing import List
+from typing import List, Optional
 
 from PySide2.QtCore import Qt, Signal
 from PySide2.QtWidgets import QAction, QFileDialog, QInputDialog, QMainWindow, QTabWidget
@@ -20,7 +20,7 @@ class MainWindow(QMainWindow):
 
     updateCountSignal = Signal()
 
-    def __init__(self, bind_address: str, port: int, filesToRead: [str]):
+    def __init__(self, bind_address: str, port: int, filesToRead: [str], thumbnails_directory: Optional[str]):
         """
         :param bind_address: address to bind to when listening for live connections.
         :param port: port to bind to when listening for live connections.
@@ -35,7 +35,8 @@ def __init__(self, bind_address: str, port: int, filesToRead: [str]):
         }
 
         self.liveWindow = LiveWindow(bind_address, port, self.updateCountSignal, self.options, parent=self)
-        self.replayWindow = ReplayWindow(self.options, parent=self)
+        self.replayWindow = ReplayWindow(self.options, parent=self,
+                                         thumbnails_directory=thumbnails_directory)
         self.tabManager = QTabWidget()
         self.tabManager.addTab(self.liveWindow, "Live connections")
         self.tabManager.addTab(self.replayWindow, "Replays")

pyrdp/player/ReplayTab.py

@@ -3,11 +3,17 @@
 # Copyright (C) 2019-2021 GoSecure Inc.
 # Licensed under the GPLv3 or later.
 #
-from PySide2.QtGui import QResizeEvent
+import os
+import pickle
+from typing import Dict, Optional
+
+from PySide2.QtCore import QByteArray
+from PySide2.QtGui import QImage, QResizeEvent
 from PySide2.QtWidgets import QApplication, QWidget
 
-from pyrdp.layer import PlayerLayer
 from pyrdp.player.BaseTab import BaseTab
+from pyrdp.player.gdi.cache import BitmapCache
+from pyrdp.player.gdi.SerializableBitmapCache import SerializableBitmapCache
 from pyrdp.player.PlayerEventHandler import PlayerEventHandler
 from pyrdp.player.Replay import Replay, ReplayReader
 from pyrdp.player.ReplayBar import ReplayBar
@@ -20,25 +26,29 @@ class ReplayTab(BaseTab):
     Tab that displays a RDP Connection that is being replayed from a file.
     """
 
-    def __init__(self, fileName: str, parent: QWidget):
+    def __init__(self, fileName: str, parent: QWidget, thumbnail_directory: Optional[str] = None):
         """
         :param fileName: name of the file to read.
         :param parent: parent widget.
         """
         self.viewer = QRemoteDesktop(800, 600, parent)
+
         super().__init__(self.viewer, parent)
         QApplication.instance().aboutToQuit.connect(self.onClose)
 
         self.fileName = fileName
         self.file = open(self.fileName, "rb")
         self.eventHandler = PlayerEventHandler(self.widget, self.text)
 
+        thumbnails = self.makeThumbnailsMap(thumbnail_directory)
+
         replay = Replay(self.file)
         self.reader = ReplayReader(replay)
-        self.thread = ReplayThread(replay)
+        self.thread = ReplayThread(replay, thumbnails=thumbnails if thumbnail_directory else None)
         self.thread.eventReached.connect(self.readEvent)
         self.thread.timeUpdated.connect(self.onTimeUpdated)
         self.thread.clearNeeded.connect(self.clear)
+        self.thread.paintThumbnail.connect(self.paintThumbnail)
         self.thread.start()
 
         self.controlBar = ReplayBar(replay.duration)
@@ -51,6 +61,20 @@ def __init__(self, fileName: str, parent: QWidget):
 
         self.tabLayout.insertWidget(0, self.controlBar)
 
+    def makeThumbnailsMap(self, thumbnail_directory):
+        thumbnails = {}
+        first_thumbnail_timestamp = 0
+        if thumbnail_directory:
+            for i, file in enumerate(os.listdir(thumbnail_directory)):
+                if file.endswith(".png"):
+                    timestamp = int(file.replace('.png', ''))
+                    if i == 0:
+                        first_thumbnail_timestamp = timestamp
+                    thumbnails[timestamp - first_thumbnail_timestamp] = (
+                    f"{thumbnail_directory}/{file}",
+                    f"{thumbnail_directory}/gdi_cache/{file}".replace('.png', '.bitmapcache'))
+        return thumbnails
+
     def play(self):
         self.controlBar.button.setPlaying(True)
         self.controlBar.play.emit()
@@ -79,6 +103,25 @@ def clear(self):
         self.viewer.clear()
         self.text.setText("")
 
+    def paintThumbnail(self, thumbnailFilePath: str, bitmapCachePath: str):
+        with open(thumbnailFilePath, 'rb') as file:
+            png_content = file.read()
+
+        qimage = QImage()
+        qimage.loadFromData(QByteArray(png_content))
+        self.viewer.notifyImage(0, 0, qimage, qimage.width(), qimage.height())
+
+        # We also need to deserialize the bitmap cache.
+        print(f"load {bitmapCachePath}")
+        with open(bitmapCachePath, 'rb') as file:
+            serializableBitmapCaches: Dict[int, SerializableBitmapCache] = pickle.load(file)
+
+        bitmapCache = BitmapCache()
+        for key, serializableCache in serializableBitmapCaches.items():
+            bitmapCache.caches[key] = serializableCache.cache
+
+        self.eventHandler.gdi.bitmaps = bitmapCache
+
     def onClose(self):
         self.thread.close()
         self.thread.wait()

pyrdp/player/ReplayThread.py

@@ -8,6 +8,7 @@
 from enum import IntEnum
 from multiprocessing import Queue
 from time import sleep
+from typing import Dict, Optional, Tuple
 
 from PySide2.QtCore import QThread, Signal
 
@@ -36,15 +37,18 @@ class ReplayThread(QThread):
     # We use the object type instead of int for this signal to prevent Python integers from being converted to 32-bit integers
     eventReached = Signal(object)
     clearNeeded = Signal()
+    paintThumbnail = Signal(str, str)
 
-    def __init__(self, replay: Replay):
+    def __init__(self, replay: Replay, thumbnails: Optional[Dict[int, Tuple[str, str]]] = None):
         super().__init__()
 
+        self.thumbnails = thumbnails
         self.queue = Queue()
         self.lastSeekTime = 0
         self.requestedSpeed = 1
         self.replay = replay
         self.timer = Timer()
+        self.thumbnail_timestamp = 0
 
     def run(self):
         step = 16 / 1000
@@ -64,7 +68,12 @@ def run(self):
                     elif event == ReplayThreadEvent.PAUSE:
                         self.timer.stop()
                     elif event == ReplayThreadEvent.SEEK:
-                        if self.lastSeekTime < self.timer.getElapsedTime():
+                        if self.thumbnails:
+                            self.thumbnail_timestamp = self.get_thumbnail_timestamp(self.lastSeekTime)
+                            self.paintThumbnail.emit(self.thumbnails[self.thumbnail_timestamp][0],
+                                                     self.thumbnails[self.thumbnail_timestamp][1])
+                            currentIndex = 0
+                        elif self.lastSeekTime < self.timer.getElapsedTime():
                             currentIndex = 0
                             self.clearNeeded.emit()
 
@@ -83,15 +92,25 @@ def run(self):
 
                 while currentIndex < len(timestamps) and timestamps[currentIndex] / 1000.0 <= currentTime:
                     nextTimestamp = timestamps[currentIndex]
-                    positions = self.replay.events[nextTimestamp]
 
-                    for position in positions:
-                        self.eventReached.emit(position)
+                    # Only replay an event if it's after the latest thumbnail printed.
+                    if not self.thumbnails or nextTimestamp >= self.thumbnail_timestamp:
+                        positions = self.replay.events[nextTimestamp]
+
+                        for position in positions:
+                            self.eventReached.emit(position)
 
                     currentIndex += 1
 
             sleep(step)
 
+    def get_thumbnail_timestamp(self, seek_time: int):
+        best_thumbnail = -1
+        for thumbnail_timestamp in self.thumbnails.keys():
+            if best_thumbnail < thumbnail_timestamp < seek_time * 1000:
+                best_thumbnail = thumbnail_timestamp
+        return best_thumbnail
+
     def play(self):
         self.queue.put(ReplayThreadEvent.PLAY)
 

pyrdp/player/ReplayWindow.py

@@ -4,7 +4,7 @@
 # Licensed under the GPLv3 or later.
 #
 
-from typing import Dict
+from typing import Dict, Optional
 
 from PySide2.QtGui import QResizeEvent
 from PySide2.QtWidgets import QWidget
@@ -18,15 +18,17 @@ class ReplayWindow(BaseWindow):
     Class for managing replay tabs.
     """
 
-    def __init__(self, options: Dict[str, object], parent: QWidget):
+    def __init__(self, options: Dict[str, object], parent: QWidget,
+                 thumbnails_directory: Optional[str]):
         super().__init__(options, parent=parent)
+        self.thumbnails_directory = thumbnails_directory
 
     def openFile(self, fileName: str, autoplay: bool = False):
         """
         Open a replay file and open a new tab.
         :param fileName: replay path.
         """
-        tab = ReplayTab(fileName, parent=self)
+        tab = ReplayTab(fileName, parent=self, thumbnail_directory=self.thumbnails_directory)
         self.addTab(tab, fileName)
         self.log.debug("Loading replay file %(arg1)s", {"arg1": fileName})
         if autoplay:

pyrdp/player/ThumbnailEventHandler.py

@@ -3,22 +3,19 @@
 # Copyright (C) 2020-2021 GoSecure Inc.
 # Licensed under the GPLv3 or later.
 #
-
-from pyrdp.enum import BitmapFlags, CapabilityType
-from pyrdp.pdu import BitmapUpdateData, PlayerPDU
-from pyrdp.player.RenderingEventHandler import RenderingEventHandler
-from pyrdp.ui import RDPBitmapToQtImage
-
 import logging
-
-import av
-from PIL import ImageQt
-from PySide2.QtGui import QImage, QPainter, QColor
+import os
+import pickle
 from os import path
 from pathlib import Path
 
-from pyrdp.player.Mp4EventHandler import Mp4Sink
+from PySide2.QtGui import QColor, QImage, QPainter
 
+from pyrdp.enum import CapabilityType
+from pyrdp.pdu import PlayerPDU
+from pyrdp.player.gdi.SerializableBitmapCache import SerializableBitmapCache
+from pyrdp.player.Mp4EventHandler import Mp4Sink
+from pyrdp.player.RenderingEventHandler import RenderingEventHandler
 
 DEFAULT_DELTA = 300  # seconds
 
@@ -104,5 +101,16 @@ def _writeFrame(self, surface: QImage):
         tmp.save(path.join(self.dst, f'{self.timestamp}.png'))
         del tmp
 
-
-
+        self.save_bitmap_cache()
+
+    def save_bitmap_cache(self):
+        gdi_cache_directory = path.join(self.dst, 'gdi_cache')
+        serializable_bitmaps = {}
+        for key, bitmap in self.gdi.bitmaps.caches.items():
+            serializable_bitmaps[key] = SerializableBitmapCache(bitmap)
+        if not os.path.exists(gdi_cache_directory):
+            os.makedirs(gdi_cache_directory, exist_ok=True)
+        bitmapCachePath = path.join(gdi_cache_directory, f'{self.timestamp}.bitmapcache')
+        with open(bitmapCachePath, "wb") as file:
+            pickle.dump(serializable_bitmaps, file)
+            del serializable_bitmaps

pyrdp/player/gdi/SerializableBitmapCache.py

@@ -0,0 +1,32 @@
+from typing import Dict
+
+from PySide2.QtCore import QByteArray
+from PySide2.QtGui import QImage
+
+
+class SerializableBitmapCache:
+
+    def __init__(self, cache: Dict[int, QImage]):
+        super().__init__()
+        self.cache = cache
+
+    def __getstate__(self):
+        state = []
+        for key, qImage in self.cache.items():
+            data = {
+                'width': qImage.width(),
+                'height': qImage.height(),
+                'format': qImage.format(),
+                'data': QByteArray(bytes(qImage.bits()))
+            }
+            state.append((key, data))
+        return state
+
+    def __setstate__(self, state):
+
+        self.cache = {}
+        for key, data in state:
+            self.cache[key] = QImage(width=data['width'],
+                                     height=data['height'],
+                                     format=data['format'],
+                                     data=data['data'])