diff --git a/pyrdp/bin/convert.py b/pyrdp/bin/convert.py index bbe2685bc..5123e67f8 100755 --- a/pyrdp/bin/convert.py +++ b/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) diff --git a/pyrdp/convert/Converter.py b/pyrdp/convert/Converter.py index 604059475..d81f4a8ac 100644 --- a/pyrdp/convert/Converter.py +++ b/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") diff --git a/pyrdp/convert/MP4EventHandler.py b/pyrdp/convert/MP4EventHandler.py index 9a8509d1f..6b98ec571 100644 --- a/pyrdp/convert/MP4EventHandler.py +++ b/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() diff --git a/pyrdp/convert/PCAPConverter.py b/pyrdp/convert/PCAPConverter.py index 7f87ad130..90a8999fc 100644 --- a/pyrdp/convert/PCAPConverter.py +++ b/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}") diff --git a/pyrdp/convert/ReplayConverter.py b/pyrdp/convert/ReplayConverter.py index 0e66ccd6c..193c6f7c5 100644 --- a/pyrdp/convert/ReplayConverter.py +++ b/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.") diff --git a/pyrdp/convert/utils.py b/pyrdp/convert/utils.py index 0a75e1a86..6fbc70203 100644 --- a/pyrdp/convert/utils.py +++ b/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):