fix(331): fix json conversion

pyrdp/convert/ExportedPDUStream.py

@@ -15,31 +15,27 @@ class ExportedPDUStream(PCAPStream):
     def __init__(self, client: str, server: str, packets: PacketList):
         super().__init__(client, server)
         self.packets = packets
+        self.n = 0
 
     def __len__(self):
         return len(self.packets)
 
     def __iter__(self):
-        return self.parsePDUs()
+        return self
 
-    def parsePDUs(self):
-        """
-        Generator function that parses Exported PDUs from Wireshark and outputs them.
-        """
-
-        n = 0
+    def __next__(self):
 
         while True:
-            if n >= len(self):
+            if self.n >= len(self):
                 raise StopIteration
 
-            packet = self.packets[n]
+            packet = self.packets[self.n]
             src = ".".join(str(b) for b in packet.load[12:16])
             dst = ".".join(str(b) for b in packet.load[20:24])
             data = packet.load[60:]
-            n += 1
+            self.n += 1
 
             if any(ip not in self.ips for ip in [src, dst]):
-                continue
+                continue  # Skip packets not meant for this stream.
 
-            yield PCAPStream.output(data, packet.time, src, dst)
+            return PCAPStream.output(data, packet.time, src, dst)

pyrdp/convert/PCAPConverter.py

@@ -10,13 +10,13 @@
 from progressbar import progressbar
 from scapy.layers.inet import TCP
 from scapy.layers.tls.record import TLS
+from pyrdp.convert.pyrdp_scapy import *
 
 from pyrdp.convert.Converter import Converter
 from pyrdp.convert.ExportedPDUStream import ExportedPDUStream
+from pyrdp.convert.TLSPDUStream import TLSPDUStream
 from pyrdp.convert.PCAPStream import PCAPStream
 from pyrdp.convert.RDPReplayer import RDPReplayer
-from pyrdp.convert.TLSPDUStream import TLSPDUStream
-from pyrdp.convert.pyrdp_scapy import *
 from pyrdp.convert.utils import tcp_both, getSessionInfo, findClientRandom, createHandler, canExtractSessionInfo
 
 
@@ -106,12 +106,17 @@ def processStream(self, startTimeStamp: int, stream: PCAPStream):
 
         print(f"[*] Processing {stream.client} -> {stream.server}")
 
-        for data, timeStamp, src, _ in progressbar(stream):
-            replayer.setTimeStamp(timeStamp)
-            replayer.recv(data, src == stream.client)
+        try:
+            for data, timeStamp, src, _dst in progressbar(stream):
+                replayer.setTimeStamp(timeStamp)
+                replayer.recv(data, src == stream.client)
+        except StopIteration:
+            # Done processing the stream.
+            pass
 
         try:
             replayer.tcp.recordConnectionClose()
+            handler.cleanup()
         except struct.error:
             sys.stderr.write("[!] Couldn't close the session cleanly. Make sure that --src and --dst are correct.")