Optimize MP4 forensic replay export (~12x faster, idle skip)

pyrdp/bin/convert.py

@@ -57,6 +57,14 @@ def main():
         "otherwise the result is output next to the source file with the proper extension. "
         "However if the source of the conversion is a .pcap then this option will create a directory where all files will be stored.",
     )
+    parser.add_argument(
+        "--idle-skip",
+        type=int,
+        default=0,
+        metavar="N",
+        help="Skip idle periods longer than N seconds in MP4 output (0 = disabled, default: 0)",
+    )
+
 
     args = parser.parse_args()
 
@@ -83,15 +91,19 @@ def main():
     else:
         outputPrefix = ""
 
+    handler_kwargs = {}
+    if args.idle_skip > 0:
+        handler_kwargs['idle_skip'] = args.idle_skip
+
     if inputFile.suffix in [".pcap"]:
         secrets = loadSecrets(args.secrets) if args.secrets else None
-        converter = PCAPConverter(inputFile, outputPrefix, args.format, secrets=secrets, srcFilter=args.src, dstFilter=args.dst, listOnly=args.list_only)
+        converter = PCAPConverter(inputFile, outputPrefix, args.format, secrets=secrets, srcFilter=args.src, dstFilter=args.dst, listOnly=args.list_only, handler_kwargs=handler_kwargs)
     elif inputFile.suffix in [".pyrdp"]:
         if args.format == "replay":
             sys.stderr.write("Refusing to convert a replay file to a replay file. Choose another format.")
             sys.exit(1)
 
-        converter = ReplayConverter(inputFile, outputPrefix, args.format)
+        converter = ReplayConverter(inputFile, outputPrefix, args.format, handler_kwargs=handler_kwargs)
     else:
         sys.stderr.write("Unknown file extension. (Supported: .pcap, .pyrdp)")
         sys.exit(1)

pyrdp/convert/Converter.py

@@ -7,10 +7,11 @@
 
 
 class Converter:
-    def __init__(self, inputFile: Path, outputPrefix: str, format: str):
+    def __init__(self, inputFile: Path, outputPrefix: str, format: str, handler_kwargs: dict = None):
         self.inputFile = inputFile
         self.outputPrefix = outputPrefix
         self.format = format
+        self.handler_kwargs = handler_kwargs or {}
 
     def process(self):
         raise NotImplementedError("Converter.process is not implemented")

pyrdp/convert/MP4EventHandler.py

@@ -4,7 +4,7 @@
 # Licensed under the GPLv3 or later.
 #
 
-from pyrdp.enum import CapabilityType
+from pyrdp.enum import CapabilityType, PlayerPDUType
 from pyrdp.pdu import PlayerPDU
 from pyrdp.player.ImageHandler import ImageHandler
 from pyrdp.player.RenderingEventHandler import RenderingEventHandler
@@ -38,35 +38,42 @@ def screen(self) -> QImage:
 
 class MP4EventHandler(RenderingEventHandler):
 
-    def __init__(self, filename: str, fps=30, progress=None):
+    def __init__(self, filename: str, fps=10, progress=None, idle_skip=0):
         """
         Construct an event handler that outputs to an Mp4 file.
 
         :param filename: The output file to write to.
-        :param fps: The frame rate (30 recommended).
+        :param fps: The frame rate (10 recommended for forensic captures).
         :param progress: An optional callback (sig: `() -> ()`) whenever a frame is muxed.
+        :param idle_skip: Seconds of inactivity before compressing idle gaps (0 = disabled).
         """
         self.filename = filename
-        # The movflags puts the encoder in an MP4 Streaming Format. This has two benefits:
-        # - recover partial videos in case of a pyrdp-convert crash
-        # - reduce memory consumption (especially for long captures)
-        # See: https://ffmpeg.org/ffmpeg-formats.html#mov_002c-mp4_002c-ismv
-        self.mp4 = f = av.open(filename, 'w', options={'movflags': 'frag_keyframe+empty_moov'})
+        # faststart moves the moov atom to the front for seekable playback.
+        self.mp4 = f = av.open(filename, 'w', options={'movflags': 'faststart'})
         self.stream = f.add_stream('h264', rate=fps)
-        # TODO: this undocumented PyAV stream feature needs to be properly investigated
-        #       we could probably batch the encoding of several frames and benefit from threads
-        #       but trying this as-is lead to no gains
-        #       (actually a degradation but that could be statistically irrelevant)
-        #self.stream.thread_count = 4
         self.stream.pix_fmt = 'yuv420p'
+        self.stream.options = {'preset': 'ultrafast'}
+        self.stream.gop_size = fps * 5  # Keyframe every 5s for seeking
         self.progress = progress
-        self.scale = False
+        self.padW = 0
+        self.padH = 0
         self.mouse = (0, 0)
         self.fps = fps
         self.delta = 1000 // fps  # ms per frame
         self.log = logging.getLogger(__name__)
         self.log.info('Begin MP4 export to %s: %d FPS', filename, fps)
         self.timestamp = self.prevTimestamp = None
+        # PTS counter in stream time_base units for correct playback timing
+        self.pts = 0
+        # Track whether the surface has changed since the last encoded frame
+        self.dirty = False
+        # Idle skip: based on user input (keyboard/mouse), not screen changes
+        self.idle_skip_ms = idle_skip * 1000 if idle_skip > 0 else 0
+        self.total_skipped_ms = 0
+        self.lastInputTimestamp = None
+        self._in_idle = False
+        self._idle_enter_ts = None
+        self._last_idle_frame_ts = None
 
         super().__init__(MP4Image())
 
@@ -79,19 +86,87 @@ def onPDUReceived(self, pdu: PlayerPDU):
 
         ts = pdu.timestamp
         self.timestamp = ts
+        is_input = pdu.header == PlayerPDUType.FAST_PATH_INPUT
+
+        # Track user input for idle detection
+        if is_input:
+            self.lastInputTimestamp = ts
 
         if self.prevTimestamp is None:
-            dt = self.delta
-        else:
-            dt = self.timestamp - self.prevTimestamp  # ms
+            # First PDU: assume active at start
+            if self.lastInputTimestamp is None:
+                self.lastInputTimestamp = ts
+            if self.dirty:
+                self.writeFrame()
+                self.dirty = False
+            self.prevTimestamp = ts
+            return
+
+        # Check input-idle state
+        input_idle_ms = ts - self.lastInputTimestamp
+        now_idle = self.idle_skip_ms > 0 and input_idle_ms > self.idle_skip_ms
+
+        if now_idle and not self._in_idle:
+            # Entering idle: encode last dirty frame, then freeze
+            self._in_idle = True
+            self._idle_enter_ts = ts
+            self._last_idle_frame_ts = ts
+            if self.dirty:
+                self.writeFrame()
+                self.dirty = False
+
+        if self._in_idle and not now_idle:
+            # Exiting idle (input resumed)
+            self._in_idle = False
+            self.total_skipped_ms += ts - self._idle_enter_ts
+            self.pts += self.fps  # 1s pause in output
+            self.prevTimestamp = ts
+            return
+
+        if self._in_idle:
+            # During idle: encode 1 frame every 10s to capture screen state
+            if self.dirty and (ts - self._last_idle_frame_ts) >= 10000:
+                self.writeFrame()
+                self.dirty = False
+                self._last_idle_frame_ts = ts
+            else:
+                self.dirty = False
+            self.prevTimestamp = ts
+            return
+
+        # Normal (non-idle) processing
+        dt = ts - self.prevTimestamp  # ms
         nframes = (dt // self.delta)
+
         if nframes > 0:
-            for _ in range(nframes):
+            # Frame boundary crossed. Encode one frame if surface changed,
+            # then advance PTS to cover any remaining idle gap.
+            if self.dirty:
                 self.writeFrame()
+                self.dirty = False
+                nframes -= 1  # One frame was just encoded
+
+            # True gap (no PDUs at all): compress
+            gap_threshold = int(self.idle_skip_ms / self.delta) if self.idle_skip_ms > 0 else 0
+            if gap_threshold > 0 and nframes > gap_threshold:
+                self.total_skipped_ms += nframes * self.delta
+                nframes = self.fps  # Replace gap with 1s pause
+
+            # Skip remaining frames (player holds last frame)
+            self.pts += nframes
             self.prevTimestamp = ts
-            self.log.debug('Rendered %d still frame(s)', nframes)
 
     def cleanup(self):
+        # Close out idle state if capture ends while idle
+        if self._in_idle and self._idle_enter_ts and self.timestamp:
+            self.total_skipped_ms += self.timestamp - self._idle_enter_ts
+            self._in_idle = False
+
+        # Flush any pending dirty frame
+        if self.dirty:
+            self.writeFrame()
+            self.dirty = False
+
         # Add one second worth of padding so that the video doesn't end too abruptly.
         for _ in range(self.fps):
             self.writeFrame()
@@ -101,6 +176,13 @@ def cleanup(self):
             if self.progress:
                 self.progress()
             self.mp4.mux(pkt)
+
+        if self.total_skipped_ms > 0:
+            skipped_s = self.total_skipped_ms / 1000
+            m, s = divmod(int(skipped_s), 60)
+            h, m = divmod(m, 60)
+            self.log.info('Total idle time skipped: %dh %dm %ds (%.1fs)', h, m, s, skipped_s)
+
         self.log.info('Export completed.')
         self.mp4.close()
 
@@ -111,39 +193,45 @@ def onMousePosition(self, x, y):
     def onCapabilities(self, caps):
         bmp = caps[CapabilityType.CAPSTYPE_BITMAP]
         (w, h) = (bmp.desktopWidth, bmp.desktopHeight)
-        self.imageHandler.resize(w, h)
-
-        if w % 2 != 0:
-            self.scale = True
-            w += 1
-        if h % 2 != 0:
-            self.scale = True
-            h += 1
 
-        self.stream.width = w
-        self.stream.height = h
+        # H264 requires even dimensions. Pad by 1px instead of scaling
+        # every frame (scaling 2556x929 is ~6ms per frame).
+        self.padW = w % 2
+        self.padH = h % 2
+        self.stream.width = w + self.padW
+        self.stream.height = h + self.padH
 
+        self.imageHandler.resize(w, h)
         super().onCapabilities(caps)
 
     def onFinishRender(self):
-        # When the screen is updated, always write a frame.
-        self.prevTimestamp = self.timestamp
-        self.writeFrame()
+        # Mark surface as changed. The frame will be encoded at the next
+        # frame boundary in onPDUReceived, batching multiple renders.
+        self.dirty = True
 
     def writeFrame(self):
         w = self.stream.width
         h = self.stream.height
-        surface = self.imageHandler.screen.scaled(w, h) if self.scale else self.imageHandler.screen.copy()
 
-        # Draw the mouse pointer. Render mouse clicks?
-        p = QPainter(surface)
+        if self.padW or self.padH:
+            # Create even-sized surface and draw screen into it (avoids full scale)
+            surface = QImage(w, h, QImage.Format_ARGB32_Premultiplied)
+            p = QPainter(surface)
+            p.drawImage(0, 0, self.imageHandler.screen)
+        else:
+            surface = self.imageHandler.screen.copy()
+            p = QPainter(surface)
+
+        # Draw the mouse pointer.
         p.setBrush(QColor.fromRgb(255, 255, 0, 180))
         (x, y) = self.mouse
         p.drawEllipse(x, y, 5, 5)
         p.end()
 
-        # Output frame.
+        # Output frame with explicit PTS for correct playback timing.
         frame = av.VideoFrame.from_ndarray(qimage2ndarray.rgb_view(surface))
+        frame.pts = self.pts
+        self.pts += 1
         for packet in self.stream.encode(frame):
             if self.progress:
                 self.progress()

pyrdp/convert/PCAPConverter.py

@@ -24,8 +24,8 @@
 class PCAPConverter(Converter):
     SESSIONID_FORMAT = "{timestamp}_{src}-{dst}"
 
-    def __init__(self, inputFile: Path, outputPrefix: str, format: str, secrets: Dict = None, srcFilter = None, dstFilter = None, listOnly = False):
-        super().__init__(inputFile, outputPrefix, format)
+    def __init__(self, inputFile: Path, outputPrefix: str, format: str, secrets: Dict = None, srcFilter = None, dstFilter = None, listOnly = False, handler_kwargs: dict = None):
+        super().__init__(inputFile, outputPrefix, format, handler_kwargs=handler_kwargs)
         self.secrets = secrets if secrets is not None else {}
         self.srcFilter = srcFilter if srcFilter is not None else srcFilter
         self.dstFilter = dstFilter if dstFilter is not None else dstFilter
@@ -106,7 +106,7 @@ def processStream(self, startTimeStamp: int, stream: PCAPStream):
         })
         sessionID = sessionID.replace(":", "_")
 
-        handler, _ = createHandler(self.format, self.outputPrefix + sessionID)
+        handler, _ = createHandler(self.format, self.outputPrefix + sessionID, **self.handler_kwargs)
         replayer = RDPReplayer(handler, self.outputPrefix, sessionID)
 
         print(f"[*] Processing {stream.client} -> {stream.server}")

pyrdp/convert/ReplayConverter.py

@@ -20,7 +20,7 @@ def process(self):
             print(f"[*] Converting '{self.inputFile}' to {self.format.upper()}")
 
             outputFileBase = self.outputPrefix + self.inputFile.stem
-            handler, outputPath = createHandler(self.format, outputFileBase)
+            handler, outputPath = createHandler(self.format, outputFileBase, **self.handler_kwargs)
 
             if not handler:
                 print("The input file is already a replay file. Nothing to do.")

pyrdp/convert/utils.py

@@ -74,7 +74,7 @@ def extractInetAddressesFromPDUPacket(packet) -> Tuple[InetAddress, InetAddress]
     return (InetAddress(x.src, x.sport), InetAddress(x.dst, x.dport))
 
 
-def createHandler(format: str, outputFileBase: str, progress=None) -> Tuple[str, str]:
+def createHandler(format: str, outputFileBase: str, progress=None, **kwargs) -> Tuple[str, str]:
     """
     Gets the appropriate handler and returns the filename with extension.
     Returns None if the format is replay.
@@ -87,7 +87,9 @@ def createHandler(format: str, outputFileBase: str, progress=None) -> Tuple[str,
 
     HandlerClass, ext = HANDLERS[format]
     outputFileBase += f".{ext}"
-    return HandlerClass(outputFileBase, progress=progress) if HandlerClass else None, outputFileBase
+    if HandlerClass:
+        return HandlerClass(outputFileBase, progress=progress, **kwargs), outputFileBase
+    return None, outputFileBase
 
 
 class ExportedPDU(Packet):