From e30f8f2bc1a3c2cd0ec51e30e4790739ae2188d5 Mon Sep 17 00:00:00 2001
From: Alexandre Beaulieu <alex@segfault.me>
Date: Wed, 15 Apr 2020 14:35:43 -0400
Subject: [PATCH] chore: rebased MP4 conversion to support drawing orders

---
 pyrdp/player/Mp4EventHandler.py       | 33 +++++++++++++++++----------
 pyrdp/player/PlayerEventHandler.py    |  2 ++
 pyrdp/player/RenderingEventHandler.py | 16 ++++++++-----
 3 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/pyrdp/player/Mp4EventHandler.py b/pyrdp/player/Mp4EventHandler.py
index a51b0b59e..67d63dfd7 100644
--- a/pyrdp/player/Mp4EventHandler.py
+++ b/pyrdp/player/Mp4EventHandler.py
@@ -17,19 +17,30 @@
 
 
 class Mp4Sink:
+    """A QRemoteDesktop Mock."""
     def __init__(self):
-        self.screen: QImage = None
+        self._buffer: QImage = None
+
+    @property
+    def screen(self):
+        return self._buffer
 
     def notifyImage(self, x: int, y: int, img: QImage, w: int, h: int):
-        p = QPainter(self.screen)
+        p = QPainter(self._buffer)
         p.drawImage(x, y, img, 0, 0, w, h)
 
+    def resize(self, w: int, h: int):
+        self._buffer = QImage(w, h, QImage.Format_ARGB32_Premultiplied)
+
+    def width(self) -> int:
+        return self._buffer.width()
+
+    def height(self) -> int:
+        return self._buffer.height()
+
     def update(self):
         pass
 
-    def resize(self, w: int, h: int):
-        self.screen = QImage(w, h, QImage.Format_RGB888)
-
 
 class Mp4EventHandler(RenderingEventHandler):
 
@@ -75,14 +86,14 @@ def onPDUReceived(self, pdu: PlayerPDU):
         nframes = (dt // self.delta)
         if nframes > 0:
             for _ in range(nframes):
-                self._writeFrame(self.surface)
+                self._writeFrame(self.sink.screen)
             self.prevTimestamp = ts
             self.log.debug('Rendered %d still frame(s)', nframes)
 
     def cleanup(self):
         # Add one second worth of padding so that the video doesn't end too abruptly.
         for _ in range(self.fps):
-            self._writeFrame(self.surface)
+            self._writeFrame(self.sink.screen)
 
         self.log.info('Flushing to disk: %s', self.filename)
         for pkt in self.stream.encode():
@@ -90,7 +101,6 @@ def cleanup(self):
                 self.progress()
             self.mp4.mux(pkt)
         self.log.info('Export completed.')
-
         self.mp4.close()
 
     def onMousePosition(self, x, y):
@@ -112,18 +122,17 @@ def onCapabilities(self, caps):
         self.stream.width = w
         self.stream.height = h
 
-    def onBeginRender(self):
-        pass
+        super().onCapabilities(caps)
 
     def onFinishRender(self):
         # When the screen is updated, always write a frame.
         self.prevTimestamp = self.timestamp
-        self._writeFrame(self.surface)
+        self._writeFrame(self.sink.screen)
 
     def _writeFrame(self, surface: QImage):
         w = self.stream.width
         h = self.stream.height
-        surface = self.sink.screen.scaled(w, h) if self.scale else self.sink.screen
+        surface = surface.scaled(w, h) if self.scale else surface
         frame = av.VideoFrame.from_image(ImageQt.fromqimage(surface))
 
         # Draw the mouse pointer.
diff --git a/pyrdp/player/PlayerEventHandler.py b/pyrdp/player/PlayerEventHandler.py
index 38223d09a..5b128b0a9 100644
--- a/pyrdp/player/PlayerEventHandler.py
+++ b/pyrdp/player/PlayerEventHandler.py
@@ -46,6 +46,8 @@ def onCapabilities(self, caps):
         (w, h) = (bmp.desktopWidth, bmp.desktopHeight)
         self.viewer.resize(w, h)
 
+        super().onCapabilities(caps)
+
     def onMousePosition(self, x: int, y: int):
         self.viewer.setMousePosition(x, y)
 
diff --git a/pyrdp/player/RenderingEventHandler.py b/pyrdp/player/RenderingEventHandler.py
index 281889913..7f070eec7 100644
--- a/pyrdp/player/RenderingEventHandler.py
+++ b/pyrdp/player/RenderingEventHandler.py
@@ -28,30 +28,34 @@ def __init__(self, sink):
     def onCapabilities(self, caps):
         if CapabilityType.CAPSTYPE_ORDER in caps:
             self.gdi = GdiQtFrontend(self.sink)
-            self.orders = OrdersParser(self.gdi)
-            self.orders.onCapabilities(caps)
+            self._orders = OrdersParser(self.gdi)
+            self._orders.onCapabilities(caps)
 
     # Generic Video Parsing Routines.
     def onFastPathOutput(self, event: FastPathOutputEvent):
+        self.onBeginRender()
         if isinstance(event, FastPathBitmapEvent):
             parsed = self._fastPath.parseBitmapEvent(event)
-            self.onBeginRender()
             for bmp in parsed.bitmapUpdateData:
                 self.onBitmap(bmp)
-            self.onFinishRender()
 
         elif isinstance(event, FastPathOrdersEvent):
-            if self.orders is None:
+            if self._orders is None:
                 self.log.error('Received Unexpected Drawing Orders!')
                 return
+            self.onBeginRender()
             self._orders.parse(event)
 
+        self.onFinishRender()
+
     def onSlowPathUpdate(self, pdu: UpdatePDU):
         if pdu.updateType == SlowPathUpdateType.SLOWPATH_UPDATETYPE_BITMAP:
-            updates = self._bitmap.parseBitmapUpdateData(pdu.updateData)
             self.onBeginRender()
+
+            updates = self._bitmap.parseBitmapUpdateData(pdu.updateData)
             for bmp in updates:
                 self.onBitmap(bmp)
+
             self.onFinishRender()
 
     def onBitmap(self, bitmapData: BitmapUpdateData):