rewrite to use scapy.all.sniff callback parameter, better, still sucks

decode.py

@@ -25,38 +25,13 @@ args = parser.parse_args()
 
 
 # TODO - probably a better way to do this
-def live_capture():
-    while True:
-        # TODO: This is wrong:
-        # during scapy.all.sniff construction we lose about 27 packets.
-        # we work around it by capturing 2* packets(frame), so we just lose
-        # about half the frames.
-        cap = scapy.all.sniff(iface="enp1s0f0", count=32 * 2)
-        for pkt in cap:
-            l = len(bytes(pkt))
-            if l != 6972:
-                continue
-            data = bytes(pkt)
-            yield data
-
-
-if args.live:
-    print('live stream, import scapy')
-    import scapy.all
-    print('open stream')
-    stream = live_capture()
-    start = datetime.now()
-    timestamp = start.strftime('%Y%m%d_%H%M%S')
-    basename = f'live_{timestamp}'
-    blocks = tqdm(stream)
-else:
-    # Now use args.input_file as the file to process
-    input_file = args.input_file
-    basename = os.path.splitext(os.path.basename(input_file))[0]
-    stream = open(input_file, 'rb')
-    # Read packets from a pcap file
-    scanner = pcapng.scanner.FileScanner(stream)
-    blocks = tqdm(scanner)
+def live_capture_cb(cb):
+    def outer(pkt):
+        data = bytes(pkt)
+        l = len(data)
+        if l == 6972:
+            cb(data)
+    scapy.all.sniff(iface="enp1s0f0", filter='udp', prn=outer)
 
 
 def rightsize(it):
@@ -77,7 +52,11 @@ def rightsize(it):
 def removestart(it):
     "Remove the UDP header from the packets"
     for x in it:
-        yield x[0x2A:]
+        yield removestart_inner(x)
+
+
+def removestart_inner(x):
+    return x[0x2A:]
 
 
 # Function to parse packet data
@@ -97,28 +76,51 @@ def parsed(it):
         yield parse(x)
 
 
+class FrameCollector:
+    def __init__(self):
+        self.current = []
+
+    def handle(self, obj):
+        ret = None
+        if obj['part'] == 0:
+            if len(self.current) > 0:
+                ret = b"".join(self.current)
+            self.current = []
+            #otherdata = []
+        self.current.append(obj["data"])
+        return ret
+        #otherdata.append(obj)
+
+    def last(self):
+        if len(self.current) > 0:
+            return b"".join(current)
+        return None
+
+
 # Function to group data into frames
 def frames(it):
-    current = []
+    handler = FrameCollector()
     #otherdata = []
     for obj in it:
-        if obj['part'] == 0:
-            if len(current) > 0:
-                yield b"".join(current)
-            current = []
-            #otherdata = []
-        current.append(obj["data"])
-        #otherdata.append(obj)
-    if len(current) > 0:
-        yield b"".join(current)
+        ret = handler.handle(obj)
+        if ret:
+            yield ret
+    last = handler.last()
+    if last:
+        yield last
+
 
 WIDTH = 384
 HEIGHT = 288
 
 
+def bad_frame(frame, width=WIDTH, height=HEIGHT):
+    return len(frame) != width * height * 2  # 16 bpp
+
+
 def skip_bad_frames(it, width=WIDTH, height=HEIGHT):
     for frame in it:
-        if len(frame) != width * height * 2:  # 16 bpp
+        if bad_frame(frame):  # 16 bpp
             # Will be fixed when we stopped doing restarts
             #print(f'{len(frame)} != {width} * {height} * 2')
             continue
@@ -130,16 +132,24 @@ def iterimages(it, width=WIDTH, height=HEIGHT, pixelformat=">H"):
         yield Image.fromarray(np.frombuffer(frame, dtype=pixelformat).reshape(width, height))
 
 
-# Get frames and convert them to images
-frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks)))))
+def process_video():
+    # Now use args.input_file as the file to process
+    input_file = args.input_file
+    basename = os.path.splitext(os.path.basename(input_file))[0]
+    stream = open(input_file, 'rb')
+    # Read packets from a pcap file
+    scanner = pcapng.scanner.FileScanner(stream)
+    blocks = tqdm(scanner)
 
-# Create the directory for frames if not exists
-frame_dir = f"frames/{basename}"
-if not os.path.exists(frame_dir):
-    os.makedirs(frame_dir)
+    # Get frames and convert them to images
+    frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks)))))
 
-# Save each image as a PNG file
-if not args.live:
+    # Create the directory for frames if not exists
+    frame_dir = f"frames/{basename}"
+    if not os.path.exists(frame_dir):
+        os.makedirs(frame_dir)
+
+    # Save each image as a PNG file
     images = iterimages(it=frames)
     for i, img in enumerate(images):
         img.save(f'frames/{basename}/{basename}_{i:04}.png')
@@ -162,7 +172,8 @@ if not args.live:
     subprocess.run(command)
     print("to play: ffplay thermal.mp4")
 
-else:
+
+if args.live:
     # TODO: to video via ffmpeg; right now just a single png
     # of the last frame
     def todo_live_ffmpeg():
@@ -175,7 +186,33 @@ else:
         for frame in frames:
             fd.write(frame)
 
-    for frame in frames:
-        print('.')
+    print('live stream, import scapy')
+    import scapy.all
+    print('open stream')
+
+    class PacketHandler:
+        def __init__(self, cb):
+            self.frame_collector = FrameCollector()
+            self.cb = cb
+
+        def handle(self, pkt):
+            pkt = removestart_inner(pkt)
+            parsed = parse(pkt)
+            frame_maybe = self.frame_collector.handle(parsed)
+            if not frame_maybe or bad_frame(frame_maybe):
+                return
+            self.cb(frame_maybe)
+
+    progress = tqdm()
+
+    def on_frame(frame):
+        progress.update(1)
         Image.fromarray(np.frombuffer(frame, dtype='>H').reshape(WIDTH, HEIGHT)).save(f'live.new.png')
         os.rename('live.new.png', 'live.png')
+
+    handler = PacketHandler(on_frame)
+    live_capture_cb(handler.handle)
+
+else:
+    process_video()
+