From 5db2214e2e05d94ae03ad9db5faa227e3560bf73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89milio=20Gonzalez?= Date: Sat, 3 Jul 2021 14:55:22 -0400 Subject: [PATCH 1/2] Add seeking using thumbnails generated with pyrdp-convert. It works, but there is a major bug (possibly) because of the GDI cache. It's hackish as well. --- bin/pyrdp-player.py | 7 ++- pyrdp/player/MainWindow.py | 7 +-- pyrdp/player/ReplayTab.py | 51 +++++++++++++++++++-- pyrdp/player/ReplayThread.py | 29 ++++++++++-- pyrdp/player/ReplayWindow.py | 8 ++-- pyrdp/player/ThumbnailEventHandler.py | 34 ++++++++------ pyrdp/player/gdi/SerializableBitmapCache.py | 32 +++++++++++++ 7 files changed, 139 insertions(+), 29 deletions(-) create mode 100644 pyrdp/player/gdi/SerializableBitmapCache.py diff --git a/bin/pyrdp-player.py b/bin/pyrdp-player.py index 25aa44ef0..f0a1acf13 100755 --- a/bin/pyrdp-player.py +++ b/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() diff --git a/pyrdp/player/MainWindow.py b/pyrdp/player/MainWindow.py index 3f587e75a..b7fe2800f 100644 --- a/pyrdp/player/MainWindow.py +++ b/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") diff --git a/pyrdp/player/ReplayTab.py b/pyrdp/player/ReplayTab.py index a4ccc2d6a..fa62f3fdb 100644 --- a/pyrdp/player/ReplayTab.py +++ b/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,12 +26,13 @@ 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) @@ -33,12 +40,15 @@ def __init__(self, fileName: str, parent: QWidget): 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() diff --git a/pyrdp/player/ReplayThread.py b/pyrdp/player/ReplayThread.py index be1a15cf5..c6b3ee928 100644 --- a/pyrdp/player/ReplayThread.py +++ b/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) diff --git a/pyrdp/player/ReplayWindow.py b/pyrdp/player/ReplayWindow.py index 8b4636d23..4cc40c52e 100644 --- a/pyrdp/player/ReplayWindow.py +++ b/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: diff --git a/pyrdp/player/ThumbnailEventHandler.py b/pyrdp/player/ThumbnailEventHandler.py index 01e77defc..679fb7445 100644 --- a/pyrdp/player/ThumbnailEventHandler.py +++ b/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 diff --git a/pyrdp/player/gdi/SerializableBitmapCache.py b/pyrdp/player/gdi/SerializableBitmapCache.py new file mode 100644 index 000000000..6e7f0b968 --- /dev/null +++ b/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']) From fe5286a1fe8d7fb25dd80c65911d5503485004c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89milio=20Gonzalez?= Date: Sat, 3 Jul 2021 14:56:19 -0400 Subject: [PATCH 2/2] Fix a bug when the certificate of the server contains emojis --- pyrdp/core/ssl.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pyrdp/core/ssl.py b/pyrdp/core/ssl.py index 817211525..145e10bfd 100644 --- a/pyrdp/core/ssl.py +++ b/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'):